xAI SDK

A Rust SDK for xAI's API, providing type-safe gRPC clients for all xAI services including Grok language models, embeddings, image generation, and more.

Note: This is an unofficial SDK and is not affiliated with or endorsed by xAI. [!WARNING] This project is in early development and is not production-ready. The API is subject to change without notice.

Features

Complete API Coverage: Full gRPC client implementation for all xAI services including chat completions, embeddings, image generation, video generation, batch processing, and billing management
Type Safety: Auto-generated Rust types from Protocol Buffers
Multiple Models: Support for all xAI language models (Grok-3, Grok-4, etc.)
Streaming Support: Real-time streaming for chat completions and text generation
Response Assembly: Convert streaming chunks into complete responses
Batch Processing: Asynchronous batch processing for large-scale requests
Secure: TLS encryption with automatic certificate validation

Quick Start

Prerequisites

Rust 1.88+ installed
xAI API key

Installation

Add to your Cargo.toml:

[dependencies]
xai-sdk = "0.10"
tokio = { version = "1.0", features = ["full"] }
anyhow = "1.0"

Running the Examples

Set your API key as an environment variable:
```
export XAI_API_KEY="your-api-key-here"
```
Run the authentication info example:
```
cargo run --example auth_info
```
Run the raw text sampling example:
```
cargo run --example raw_text_sample
```

Run the chat completion example (supports multiple modes):

# Blocking completion
cargo run --example chat -- --complete

# Streaming completion
cargo run --example chat -- --stream

# Streaming with assembly
cargo run --example chat -- --assemble

Run the multi-client example (demonstrates using multiple services with shared channel):
```
cargo run --example multi_client
```
Run the interceptor composition example:
```
cargo run --example interceptor_compose
```
Run the tool calls example:
```
cargo run --example tool_calls
```

API Services

The SDK provides clients for all xAI services:

Chat Service

get_completion - Get chat completion
get_completion_chunk - Stream chat completion in chunks
start_deferred_completion - Start deferred chat completion
get_deferred_completion - Retrieve deferred completion
get_stored_completion - Get stored chat completion
delete_stored_completion - Delete stored chat completion

Sample Service

sample_text - Raw text generation
sample_text_streaming - Streaming text generation

Models Service

list_language_models - List available language models
list_embedding_models - List embedding models
list_image_generation_models - List image generation models

Embed Service

embed - Generate embeddings from text or images

Image Service

generate_image - Create images from text prompts

Video Service

generate_video - Create videos with deferred processing
get_deferred_video - Retrieve generated videos

Auth Service

get_api_key_info - Get API key information

Billing Service

set_billing_info - Set billing information for a team
get_billing_info - Get billing information for a team
list_payment_methods - List payment methods on file
set_default_payment_method - Set default payment method
get_amount_to_pay - Preview current billing period amount
analyze_billing_items - Analyze historical billing usage
list_invoices - List invoices for a team
list_prepaid_balance_changes - List prepaid credit balance changes
top_up_or_get_existing_pending_change - Top up prepaid credits
get_spending_limits - Get spending limits
set_soft_spending_limit - Set soft spending limit

Files Service

upload_file - Upload files via streaming chunks
list_files - List uploaded files with pagination/sorting
retrieve_file - Retrieve file metadata by ID
delete_file - Delete a file by ID
retrieve_file_content - Stream file content (original bytes or extracted text)

Batch Service

create_batch - Create a new batch for asynchronous processing
get_batch - Retrieve information about a specific batch
list_batches - List all batches owned by the team
cancel_batch - Cancel processing of all requests in a batch
add_batch_requests - Add completion requests to an existing batch
list_batch_request_metadata - List metadata for individual requests in a batch
list_batch_results - List processing results for a batch
get_batch_request_result - Retrieve the result of a specific batch request

Client Modules

The SDK is organized into focused modules, each providing easy client creation:

Available Modules

auth - Authentication services
billing - Billing and payment management
chat - Chat completions and streaming
documents - Document processing
embed - Text and image embeddings
files - File upload and management
image - Image generation
models - Model listing and information
sample - Text sampling and generation
tokenize - Text tokenization
video - Video generation with deferred processing
batch - Batch processing and management

Complete Example

Here's a complete example showing multiple services using the modular architecture:

use anyhow::{Context, Result};
use std::env;
use xai_sdk::Request;
use xai_sdk::api::{
    Content, GetCompletionsRequest, GetModelRequest, Message, MessageRole, SampleTextRequest,
    content,
};
use xai_sdk::{chat, models, sample};

#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
    // Load API key for authentication
    let api_key = env::var("XAI_API_KEY").context("XAI_API_KEY environment variable must be set")?;

    // Create authenticated clients for different services
    let mut models_client = models::client::new(api_key.clone()).await?;
    let mut sample_client = sample::client::new(api_key.clone()).await?;
    let mut chat_client = chat::client::new(api_key).await?;

    // List available models
    let models_request = Request::new(());
    let models_response = models_client.list_language_models(models_request).await?;
    println!("Available models: {:?}", models_response.into_inner().models);

    // Get details for a specific model
    let model_request = Request::new(GetModelRequest {
        name: "grok-3-latest".to_string(),
    });
    let model_response = models_client.get_language_model(model_request).await?;
    println!("Model details: {:?}", model_response.into_inner());

    // Generate text
    let sample_request = Request::new(SampleTextRequest {
        prompt: vec!["Hello, world!".to_string()],
        model: "grok-3-latest".to_string(),
        ..Default::default()
    });
    let sample_response = sample_client.sample_text(sample_request).await?;
    println!("Generated: {}", sample_response.into_inner().choices[0].text);

    // Chat completion
    let message = Message {
        role: MessageRole::RoleUser.into(),
        content: vec![Content {
            content: Some(content::Content::Text("Explain Rust ownership".to_string())),
        }],
        ..Default::default()
    };
    let chat_request = Request::new(GetCompletionsRequest {
        model: "grok-3-latest".to_string(),
        messages: vec![message],
        ..Default::default()
    });
    let chat_response = chat_client.get_completion(chat_request).await?;
    println!(
        "Chat response: {}",
        chat_response.into_inner().outputs[0]
            .message
            .as_ref()
            .unwrap()
            .content
    );

    Ok(())
}

Streaming Utilities

The SDK provides powerful utilities for working with streaming responses:

Stream Consumer

A flexible callback system for processing streaming data:

on_chunk(chunk) - Called for each complete chunk received
on_reasoning_start(&OutputContext) - Called once per output when the reasoning phase starts (before the first reasoning token)
on_reasoning_token(&OutputContext, token: &str) - Called for each piece of reasoning content
on_reasoning_complete(&OutputContext) - Called once when the reasoning phase completes for an output
on_content_start(&OutputContext) - Called once per output when the content phase starts (before the first content token)
on_content_token(&OutputContext, token: &str) - Called for each piece of response content
on_content_complete(&OutputContext) - Called once when the content phase completes for an output
on_inline_citations(&OutputContext, &[InlineCitation]) - Called when inline citations are present in a delta
on_client_tool_calls(&OutputContext, &[ToolCall]) - Called when client-side tool calls are present
on_server_tool_calls(&OutputContext, &[ToolCall]) - Called when server-side tool calls are present
on_usage(&SamplingUsage) - Called once on the last chunk with usage statistics
on_citations(&[String]) - Called once on the last chunk with citations
on_complete() - Called once after a successful stream (after usage/citations callbacks when present)
on_error(&Status) - Called once when the stream yields an error, before process() returns Err

The OutputContext provides:

total_outputs - Total number of outputs in the stream
output_index - Index of the output this context belongs to
reasoning_status - Current status of the reasoning phase (Init, Start, Pending, or Complete)
content_status - Current status of the content phase (Init, Start, Pending, or Complete)

Stream Processing Functions

chat::stream::process - Process streaming responses with custom callbacks
chat::stream::assemble - Convert collected chunks into complete responses
chat::stream::Consumer::new_static() - Empty consumer with 'static lifetime for one-line chaining: Consumer::new_static().on_content_token(...).on_usage(...)
chat::stream::Consumer::with_sink(snk) - Consumer that forwards all stream activity as Events into any Sink<Event> (e.g. futures::channel::mpsc::unbounded()). Create the channel, pass the sender, then drain the receiver for a single event stream. Bounded senders apply backpressure.
chat::stream::Event - Enum of streaming events: Chunk, ReasoningStart/ReasoningToken/ReasoningComplete, ContentStart/ContentToken/ContentComplete, InlineCitations, ClientToolCalls, ServerToolCalls, Citations, Usage, Complete (successful end), Error(Status) (gRPC failure).
chat::stream::Consumer::with_stdout() - Pre-configured consumer for single-output real-time output
chat::stream::Consumer::with_buffered_stdout() - Pre-configured consumer for multi-output buffered output

Interceptors

The SDK provides a flexible interceptor system for customizing request handling:

Authentication Interceptor

The auth() function creates an interceptor that adds Bearer token authentication:

use xai_sdk::common::interceptor::auth;

let interceptor = auth("your-api-key");
let client = chat::client::with_interceptor(interceptor).await?;

Composing Interceptors

Combine multiple interceptors using compose():

use xai_sdk::common::interceptor::{auth, compose};

let interceptors: Vec<Box<dyn xai_sdk::export::service::Interceptor + Send + Sync>> = vec![
    Box::new(auth("your-api-key")),
    Box::new(|mut req| {
        req.metadata_mut().insert("x-custom-header", "value".parse().unwrap());
        Ok(req)
    }),
];
let composed = compose(interceptors);
let client = chat::client::with_interceptor(composed).await?;

Note: All interceptors must be Send + Sync to ensure thread safety when used in async contexts.

ClientInterceptor Type

All client functions return ClientInterceptor, a concrete type that can be:

Stored in structs
Used in trait implementations
Passed between functions without type erasure issues
Used across thread boundaries (Send + Sync) - safe to use in tokio::spawn and other async contexts

The ClientInterceptor can be created from any impl Interceptor + Send + Sync + 'static:

use xai_sdk::common::interceptor::ClientInterceptor;

let interceptor = ClientInterceptor::new(|mut req| {
    // Custom logic
    Ok(req)
});

// Can be used in spawned tasks
tokio::spawn(async move {
    let client = chat::client::with_interceptor(interceptor).await?;
    // Use client...
});

Configuration

The SDK supports comprehensive configuration options:

Temperature: Controls randomness (0.0 to 2.0)
Top-p: Nucleus sampling parameter (0.0 to 1.0)
Max tokens: Maximum tokens to generate
Log probabilities: Enable detailed token probability logging
Multiple completions: Generate multiple responses per request
Stop sequences: Custom stop conditions
Frequency/Presence penalties: Control repetition and topic diversity

Security

TLS Encryption: Automatic HTTPS with certificate validation
Authentication: Bearer token support for API key authentication
Secure by Default: No manual TLS configuration required

Error Handling

Comprehensive error handling for:

Connection errors and timeouts
Authentication failures
API rate limiting
Invalid parameters
Network issues

Development

This SDK is built using:

Protocol Buffers: Auto-generated Rust types from xAI's .proto definitions
Tonic: Modern gRPC framework for Rust with async/await support
Prost: High-performance Protocol Buffer implementation

The code is generated from xAI's official Protocol Buffer definitions, ensuring compatibility and type safety.

Changelog

See CHANGELOG.md for a detailed list of changes and new features.

License

This project is licensed under the MIT License.

Name		Name	Last commit message	Last commit date
examples		examples
sdk		sdk
.gitignore		.gitignore
.gitmodules		.gitmodules
CHANGELOG.md		CHANGELOG.md
Cargo.lock		Cargo.lock
Cargo.toml		Cargo.toml
README.md		README.md

Latest commit

History

Folders and files

Repository files navigation