Skip to main content

Crate rig_ai_sdk

Crate rig_ai_sdk 

Source
Expand description

AI SDK Data Stream Protocol adapter for rig

Provides types and utilities to work with AI SDK’s streaming protocol in rig-based applications.

§rig-ai-sdk

AI SDK Data Stream Protocol adapter for rig-core.

This crate provides types and utilities to work with AI SDK’s streaming protocol in rig-based applications, enabling seamless integration between frontend AI SDK clients (like assistant-ui) and rig-powered backends.

§Overview

The AI SDK Data Stream Protocol is a standardized format for streaming AI responses. This crate implements the protocol for use with the rig agent framework:

  • Event types - Complete set of events for streaming responses
  • Message types - Frontend message format parsing
  • Conversion utilities - Transform AI SDK messages to rig format
  • Stream adapters - Convert rig streams to AI SDK events
  • Axum SSE support - Direct integration with Axum handlers (optional feature)

§Features

FeatureDefaultDescription
axum-sseNoEnables SSE stream conversion for Axum handlers

§Installation

[dependencies]
rig-ai-sdk = { version = "0.1", features = ["axum-sse"] }

§Quick Start

§Basic Streaming

use rig_ai_sdk::{AISdkStreamBuilder, AISdkEvent};
use futures::StreamExt;

let mut builder = AISdkStreamBuilder::new();
let mut events = vec![builder.start()];

// Stream text
for chunk in ["Hello", ", ", "world", "!"] {
    events.push(builder.text_delta(chunk).unwrap());
}
events.push(builder.text_end().unwrap());
events.push(builder.finish());
events.push(builder.done());

§Axum SSE Handler

use axum::response::sse::{Event, Sse};
use rig_ai_sdk::adapt_rig_stream_sse;

async fn chat_handler() -> Sse<impl Stream<Item = Result<Event, Infallible>>> {
    let rig_stream = agent.stream_chat(prompt, history).await;
    Sse::new(adapt_rig_stream_sse(rig_stream))
}

§Converting Frontend Messages

use rig_ai_sdk::{UIMessage, extract_prompt_and_history};

let messages: Vec<UIMessage> = serde_json::from_str(&body)?;
let (prompt, history) = extract_prompt_and_history(&messages)?;
// Now use with rig agent
let response = agent.run_chat(prompt, history).await?;

§Modules

§event

Defines AISdkEvent enum representing all protocol event types:

  • Lifecycle: start, finish, done, abort, error
  • Text: text-start, text-delta, text-end
  • Reasoning: reasoning-start, reasoning-delta, reasoning-end
  • Tools: tool-input-start, tool-input-delta, tool-input-available, tool-output-available
  • Sources: source-url, source-document, file
  • Custom: data-{name} events for any custom data

§message

Types for receiving AI SDK messages from the frontend:

  • UIMessage - Individual message with role and parts
  • UIMessagePart - Message parts (text, tool, file, reasoning, etc.)
  • [PartState] - Streaming state (streaming or done)
  • [DynamicToolState] - Tool lifecycle states (5.x format)

§convert

Conversion utilities between AI SDK and rig formats:

§stream

AISdkStreamBuilder for constructing streaming responses with proper state management:

let mut builder = AISdkStreamBuilder::new();

// Proper event sequencing
builder.start();           // Initialize
builder.text_start();       // Begin text block
builder.text_delta("...");  // Stream content
builder.text_end();         // End text block
builder.finish();           // Complete stream
builder.done();             // Final marker

§adapter (requires axum-sse)

Stream adapters for integrating rig with Axum:

  • [adapt_rig_stream] - Convert rig stream to AI SDK events
  • [adapt_rig_stream_sse] - Direct Axum SSE conversion

§Event Flow

A typical AI SDK stream follows this pattern:

[start]
  ↓
[text-start] → [text-delta]* → [text-end]
  ↓
[reasoning-start] → [reasoning-delta]* → [reasoning-end]
  ↓
[tool-input-start] → [tool-input-delta]* → [tool-input-available]
  ↓
[tool-output-available]
  ↓
[finish] → [done]

§Conversion Mapping

AI SDK PartRig Content
TextUserContent::Text / AssistantContent::Text
ReasoningAssistantContent::Reasoning
File (image)UserContent::Image
ToolCallAssistantContent::ToolCall
DynamicTool (InputAvailable)AssistantContent::ToolCall
DynamicTool (OutputAvailable)UserContent::ToolResult
DynamicTool (OutputError)UserContent::ToolResult
ToolResultUserContent::ToolResult

§References

§License

MIT

Re-exports§

pub use convert::convert_message;
pub use convert::convert_messages;
pub use convert::extract_prompt_and_history;
pub use event::AISdkEvent;
pub use message::UIMessage;
pub use message::UIMessagePart;
pub use stream::AISdkStreamBuilder;

Modules§

convert
AI SDK UIMessage to rig::message::Message conversion
event
AI SDK Data Stream Protocol events
message
AI SDK UIMessage type definitions
prelude
Prelude for common imports
stream
AI SDK stream builder utilities