diff --git a/crates/dashboard/Cargo.toml b/crates/dashboard/Cargo.toml index 29892c6..021920e 100644 --- a/crates/dashboard/Cargo.toml +++ b/crates/dashboard/Cargo.toml @@ -26,6 +26,16 @@ rust-ai-agents-monitoring = { path = "../monitoring" } # Agents crate for real data integration rust-ai-agents-agents = { path = "../agents", optional = true } +# Providers crate for LLM streaming +rust-ai-agents-providers = { path = "../providers", optional = true } + +# Core types +rust-ai-agents-core = { path = "../core", optional = true } + +# Async streams for SSE +tokio-stream = "0.1" +async-stream = "0.3" + # Time chrono = { workspace = true } @@ -47,3 +57,4 @@ metrics-exporter-prometheus = "0.15" [features] default = [] agents = ["rust-ai-agents-agents"] +streaming = ["rust-ai-agents-providers", "rust-ai-agents-core"] diff --git a/crates/dashboard/src/lib.rs b/crates/dashboard/src/lib.rs index 158552e..01f448b 100644 --- a/crates/dashboard/src/lib.rs +++ b/crates/dashboard/src/lib.rs @@ -36,6 +36,9 @@ mod websocket; #[cfg(feature = "agents")] pub mod integration; +#[cfg(feature = "streaming")] +pub mod streaming; + pub use server::DashboardServer; pub use state::{ AgentStatus, DashboardMetrics, DashboardState, Session, SessionStatus, TraceEntry, @@ -46,3 +49,6 @@ pub use state::{ pub use integration::{ add_demo_data, AgentBridge, DashboardBridge, SessionBridge, TrajectoryBridge, }; + +#[cfg(feature = "streaming")] +pub use streaming::{streaming_routes, StreamEventOutput, StreamInferenceRequest, StreamingState}; diff --git a/crates/dashboard/src/streaming.rs b/crates/dashboard/src/streaming.rs new file mode 100644 index 0000000..ec78602 --- /dev/null +++ b/crates/dashboard/src/streaming.rs @@ -0,0 +1,294 @@ +//! Server-Sent Events (SSE) streaming for LLM responses. +//! +//! This module provides SSE endpoints for streaming LLM inference results +//! to clients in real-time. +//! +//! # Feature Flag +//! +//! This module requires the `streaming` feature: +//! +//! ```toml +//! rust-ai-agents-dashboard = { version = "0.1", features = ["streaming"] } +//! ``` + +use axum::{ + extract::State, + http::StatusCode, + response::{ + sse::{Event, KeepAlive, Sse}, + IntoResponse, + }, + Json, +}; +use rust_ai_agents_core::{tool::ToolSchema, LLMMessage, MessageRole}; +use rust_ai_agents_providers::{LLMBackend, StreamEvent}; +use serde::{Deserialize, Serialize}; +use std::{convert::Infallible, sync::Arc, time::Duration}; +use tokio_stream::StreamExt; + +use crate::state::DashboardState; + +/// Request body for streaming inference +#[derive(Debug, Clone, Deserialize)] +pub struct StreamInferenceRequest { + /// Messages to send to the LLM + pub messages: Vec, + /// Optional tools available to the model + #[serde(default)] + pub tools: Vec, + /// Temperature for sampling (0.0 - 1.0) + #[serde(default = "default_temperature")] + pub temperature: f32, + /// Model to use (optional, uses default if not specified) + pub model: Option, +} + +fn default_temperature() -> f32 { + 0.7 +} + +/// Input message format +#[derive(Debug, Clone, Deserialize)] +pub struct MessageInput { + pub role: String, + pub content: String, +} + +impl From for LLMMessage { + fn from(msg: MessageInput) -> Self { + let role = match msg.role.to_lowercase().as_str() { + "system" => MessageRole::System, + "assistant" => MessageRole::Assistant, + "tool" => MessageRole::Tool, + _ => MessageRole::User, + }; + LLMMessage { + role, + content: msg.content, + tool_calls: None, + tool_call_id: None, + name: None, + } + } +} + +/// SSE event types sent to clients +#[derive(Debug, Clone, Serialize)] +#[serde(tag = "type", content = "data")] +pub enum StreamEventOutput { + /// Incremental text chunk + #[serde(rename = "text_delta")] + TextDelta { text: String }, + /// Tool call started + #[serde(rename = "tool_call_start")] + ToolCallStart { id: String, name: String }, + /// Tool call arguments chunk + #[serde(rename = "tool_call_delta")] + ToolCallDelta { id: String, arguments_delta: String }, + /// Tool call completed + #[serde(rename = "tool_call_end")] + ToolCallEnd { id: String }, + /// Stream completed + #[serde(rename = "done")] + Done { + input_tokens: Option, + output_tokens: Option, + }, + /// Error occurred + #[serde(rename = "error")] + Error { message: String }, +} + +impl From for StreamEventOutput { + fn from(event: StreamEvent) -> Self { + match event { + StreamEvent::TextDelta(text) => StreamEventOutput::TextDelta { text }, + StreamEvent::ToolCallStart { id, name } => { + StreamEventOutput::ToolCallStart { id, name } + } + StreamEvent::ToolCallDelta { + id, + arguments_delta, + } => StreamEventOutput::ToolCallDelta { + id, + arguments_delta, + }, + StreamEvent::ToolCallEnd { id } => StreamEventOutput::ToolCallEnd { id }, + StreamEvent::Done { token_usage } => StreamEventOutput::Done { + input_tokens: token_usage.as_ref().map(|t| t.prompt_tokens as u32), + output_tokens: token_usage.as_ref().map(|t| t.completion_tokens as u32), + }, + StreamEvent::Error(message) => StreamEventOutput::Error { message }, + } + } +} + +/// Streaming state that holds the LLM backend +pub struct StreamingState { + pub backend: Arc, + pub dashboard: Arc, +} + +impl StreamingState { + pub fn new(backend: Arc, dashboard: Arc) -> Self { + Self { backend, dashboard } + } +} + +/// SSE streaming inference handler +/// +/// Returns a Server-Sent Events stream with LLM response chunks. +/// +/// # Example +/// +/// ```javascript +/// const eventSource = new EventSource('/api/inference/stream'); +/// eventSource.onmessage = (event) => { +/// const data = JSON.parse(event.data); +/// if (data.type === 'text_delta') { +/// appendText(data.data.text); +/// } +/// }; +/// ``` +pub async fn stream_inference_handler( + State(state): State>>, + Json(request): Json, +) -> impl IntoResponse { + let messages: Vec = request.messages.into_iter().map(Into::into).collect(); + let tools = request.tools; + let temperature = request.temperature; + let backend = state.backend.clone(); + + // Create the SSE stream + let stream = async_stream::stream! { + match backend.infer_stream(&messages, &tools, temperature).await { + Ok(mut response_stream) => { + while let Some(result) = response_stream.next().await { + match result { + Ok(event) => { + let output: StreamEventOutput = event.into(); + let json = serde_json::to_string(&output).unwrap_or_default(); + yield Ok::<_, Infallible>(Event::default().data(json)); + } + Err(e) => { + let error = StreamEventOutput::Error { + message: e.to_string(), + }; + let json = serde_json::to_string(&error).unwrap_or_default(); + yield Ok(Event::default().data(json)); + break; + } + } + } + } + Err(e) => { + let error = StreamEventOutput::Error { + message: e.to_string(), + }; + let json = serde_json::to_string(&error).unwrap_or_default(); + yield Ok(Event::default().data(json)); + } + } + }; + + Sse::new(stream).keep_alive( + KeepAlive::new() + .interval(Duration::from_secs(15)) + .text("keep-alive"), + ) +} + +/// POST handler for streaming inference (accepts JSON body) +pub async fn post_stream_inference_handler( + State(state): State>>, + Json(request): Json, +) -> impl IntoResponse { + stream_inference_handler(State(state), Json(request)).await +} + +/// Simple chat completion endpoint (non-streaming, for comparison) +pub async fn chat_completion_handler( + State(state): State>>, + Json(request): Json, +) -> impl IntoResponse { + let messages: Vec = request.messages.into_iter().map(Into::into).collect(); + let tools = request.tools; + let temperature = request.temperature; + + match state.backend.infer(&messages, &tools, temperature).await { + Ok(output) => Json(serde_json::json!({ + "content": output.content, + "tool_calls": output.tool_calls, + "token_usage": output.token_usage, + })) + .into_response(), + Err(e) => ( + StatusCode::INTERNAL_SERVER_ERROR, + Json(serde_json::json!({ "error": e.to_string() })), + ) + .into_response(), + } +} + +/// Create SSE routes for streaming +/// +/// # Example +/// +/// ```rust,ignore +/// use rust_ai_agents_dashboard::streaming::{streaming_routes, StreamingState}; +/// use rust_ai_agents_providers::AnthropicProvider; +/// +/// let backend = Arc::new(AnthropicProvider::new(api_key)); +/// let streaming_state = Arc::new(StreamingState::new(backend, dashboard_state)); +/// +/// let app = Router::new() +/// .merge(streaming_routes()) +/// .with_state(streaming_state); +/// ``` +pub fn streaming_routes() -> axum::Router>> { + use axum::routing::post; + + axum::Router::new() + .route( + "/api/inference/stream", + post(post_stream_inference_handler::), + ) + .route("/api/inference/chat", post(chat_completion_handler::)) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_message_input_conversion() { + let input = MessageInput { + role: "user".to_string(), + content: "Hello!".to_string(), + }; + + let llm_msg: LLMMessage = input.into(); + assert!(matches!(llm_msg.role, MessageRole::User)); + assert_eq!(llm_msg.content, "Hello!"); + } + + #[test] + fn test_stream_event_output_serialization() { + let event = StreamEventOutput::TextDelta { + text: "Hello".to_string(), + }; + let json = serde_json::to_string(&event).unwrap(); + assert!(json.contains("text_delta")); + assert!(json.contains("Hello")); + } + + #[test] + fn test_stream_event_conversion() { + let event = StreamEvent::TextDelta("test".to_string()); + let output: StreamEventOutput = event.into(); + match output { + StreamEventOutput::TextDelta { text } => assert_eq!(text, "test"), + _ => panic!("Wrong variant"), + } + } +} diff --git a/crates/studio/Cargo.toml b/crates/studio/Cargo.toml index b43d752..b3928f0 100644 --- a/crates/studio/Cargo.toml +++ b/crates/studio/Cargo.toml @@ -36,6 +36,11 @@ web-sys = { version = "0.3", features = [ "ErrorEvent", "Location", "Storage", + "KeyboardEvent", + "MouseEvent", + "ReadableStream", + "ReadableStreamDefaultReader", + "TextDecoder", ] } wasm-bindgen = "0.2" wasm-bindgen-futures = "0.4" diff --git a/crates/studio/src/components/sidebar.rs b/crates/studio/src/components/sidebar.rs index 7d9076e..29aa76c 100644 --- a/crates/studio/src/components/sidebar.rs +++ b/crates/studio/src/components/sidebar.rs @@ -31,6 +31,7 @@ pub fn Sidebar() -> impl IntoView { "Agents" + "Chat" diff --git a/crates/studio/src/lib.rs b/crates/studio/src/lib.rs index 99d8923..18e0fc8 100644 --- a/crates/studio/src/lib.rs +++ b/crates/studio/src/lib.rs @@ -129,6 +129,7 @@ pub fn App() -> impl IntoView { + diff --git a/crates/studio/src/pages/chat.rs b/crates/studio/src/pages/chat.rs new file mode 100644 index 0000000..1e75eca --- /dev/null +++ b/crates/studio/src/pages/chat.rs @@ -0,0 +1,326 @@ +//! Chat page with streaming LLM responses + +use leptos::prelude::*; +use leptos::task::spawn_local; +use serde::{Deserialize, Serialize}; +use wasm_bindgen::prelude::*; +use wasm_bindgen::JsCast; + +/// Message in the chat +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct ChatMessage { + pub role: String, + pub content: String, +} + +/// Stream event from SSE +#[derive(Debug, Clone, Deserialize)] +#[serde(tag = "type", content = "data")] +pub enum StreamEvent { + #[serde(rename = "text_delta")] + TextDelta { text: String }, + #[serde(rename = "tool_call_start")] + ToolCallStart { id: String, name: String }, + #[serde(rename = "tool_call_delta")] + ToolCallDelta { id: String, arguments_delta: String }, + #[serde(rename = "tool_call_end")] + ToolCallEnd { id: String }, + #[serde(rename = "done")] + Done { + input_tokens: Option, + output_tokens: Option, + }, + #[serde(rename = "error")] + Error { message: String }, +} + +/// Chat page component +#[component] +pub fn ChatPage() -> impl IntoView { + let (messages, set_messages) = signal(Vec::::new()); + let (input, set_input) = signal(String::new()); + let (is_streaming, set_streaming) = signal(false); + let (current_response, set_current_response) = signal(String::new()); + let (error, set_error) = signal(Option::::None); + + // Core send logic (extracted to avoid event type issues) + let do_send = move || { + let user_input = input.get(); + if user_input.trim().is_empty() || is_streaming.get() { + return; + } + + // Add user message + set_messages.update(|msgs| { + msgs.push(ChatMessage { + role: "user".to_string(), + content: user_input.clone(), + }); + }); + + // Clear input and start streaming + set_input.set(String::new()); + set_streaming.set(true); + set_current_response.set(String::new()); + set_error.set(None); + + // Build request body + let all_messages: Vec<_> = messages + .get() + .iter() + .map(|m| { + serde_json::json!({ + "role": m.role, + "content": m.content + }) + }) + .collect(); + + let request_body = serde_json::json!({ + "messages": all_messages, + "temperature": 0.7 + }); + + // Use fetch with ReadableStream for SSE + spawn_local(async move { + match fetch_stream(&request_body.to_string(), set_current_response, set_error).await { + Ok(final_content) => { + // Add assistant message when done + set_messages.update(|msgs| { + msgs.push(ChatMessage { + role: "assistant".to_string(), + content: final_content, + }); + }); + set_current_response.set(String::new()); + } + Err(e) => { + set_error.set(Some(e)); + } + } + set_streaming.set(false); + }); + }; + + // Button click handler + let on_click = move |_: web_sys::MouseEvent| { + do_send(); + }; + + // Handle Enter key + let on_keydown = move |e: web_sys::KeyboardEvent| { + if e.key() == "Enter" && !e.shift_key() { + e.prevent_default(); + do_send(); + } + }; + + view! { +
+
+

"Chat"

+
+ {move || if is_streaming.get() { "Streaming..." } else { "Ready" }} +
+
+ + // Error display + {move || error.get().map(|e| view! { +
+ {e} +
+ })} + + // Messages container +
+ {move || { + let msgs = messages.get(); + if msgs.is_empty() && current_response.get().is_empty() { + view! { +
+ "💬" +

"Start a conversation with the AI"

+
+ }.into_any() + } else { + view! { +
+ {msgs.into_iter().map(|msg| { + let is_user = msg.role == "user"; + view! { + + } + }).collect::>()} + + // Streaming response + {move || { + let response = current_response.get(); + if !response.is_empty() { + Some(view! { + + }) + } else { + None + } + }} +
+ }.into_any() + } + }} +
+ + // Input area +
+