278 lines
9.7 KiB
Rust
278 lines
9.7 KiB
Rust
//! Event handling example for dirigent_archivist
|
|
//!
|
|
//! This example demonstrates:
|
|
//! - Creating an EventHandler
|
|
//! - Subscribing to dirigent_protocol events
|
|
//! - Accumulating streaming message chunks
|
|
//! - Finalizing complete messages
|
|
//! - Automatic archival via event stream
|
|
|
|
use chrono::Utc;
|
|
use dirigent_archivist::{Archivist, EventHandler, Result};
|
|
use dirigent_protocol::streaming::{BusEvent, BusReceiver, EventOrigin, EventRouting};
|
|
use dirigent_protocol::{
|
|
ContentBlock, Event, Message, MessageMetadata, MessagePart, MessageRole, MessageStatus,
|
|
Session, SessionMetadata, SessionUpdate, ToolCall, ToolCallStatus,
|
|
};
|
|
use std::sync::Arc;
|
|
use std::sync::atomic::AtomicU64;
|
|
use tokio::sync::mpsc;
|
|
use uuid::Uuid;
|
|
|
|
/// Wrap a raw `Event` in a `BusEvent` with default routing.
|
|
fn wrap(event: Event) -> BusEvent {
|
|
BusEvent {
|
|
routing: EventRouting::default(),
|
|
origin: EventOrigin::Runtime,
|
|
event: Arc::new(event),
|
|
}
|
|
}
|
|
|
|
#[tokio::main]
|
|
async fn main() -> Result<()> {
|
|
// Create a temporary archive directory for this example
|
|
let temp_dir = std::env::temp_dir().join(format!("dirigent_event_example_{}", Uuid::now_v7()));
|
|
println!("Creating archive at: {}", temp_dir.display());
|
|
|
|
// Step 1: Create archivist and event handler
|
|
let archivist = Archivist::new_with_single_archive(temp_dir.clone()).await?;
|
|
let archivist = Arc::new(archivist);
|
|
let handler = EventHandler::new(archivist.clone());
|
|
|
|
println!("EventHandler created successfully");
|
|
|
|
// Step 2: Create a mock event stream. In production this is built
|
|
// by `SharingBus::subscribe_all()`; here we fabricate a `BusReceiver`
|
|
// directly so the example stays self-contained.
|
|
let (tx, rx) = mpsc::channel::<BusEvent>(100);
|
|
let bus_rx = BusReceiver {
|
|
id: 0,
|
|
rx,
|
|
lagged: Arc::new(AtomicU64::new(0)),
|
|
};
|
|
|
|
// Step 3: Spawn event handler task
|
|
let handler_task = tokio::spawn(async move {
|
|
handler.run(bus_rx).await;
|
|
});
|
|
|
|
// Step 4: Simulate event flow
|
|
println!("\n--- Simulating Event Stream ---");
|
|
|
|
// Generate connector and session IDs
|
|
let connector_id = Uuid::now_v7().to_string();
|
|
let session_id = Uuid::now_v7().to_string();
|
|
let message_id = Uuid::now_v7().to_string();
|
|
|
|
// Event 1: SessionCreated
|
|
println!("\n1. Sending SessionCreated event...");
|
|
let session_created = Event::SessionCreated {
|
|
connector_id: connector_id.clone(),
|
|
session: Session {
|
|
id: session_id.clone(),
|
|
title: "Example streaming session".to_string(),
|
|
created_at: Utc::now(),
|
|
updated_at: Utc::now(),
|
|
metadata: SessionMetadata {
|
|
project_path: "/home/user/project".to_string(),
|
|
model: Some("claude-3-5-sonnet".to_string()),
|
|
total_messages: 0,
|
|
system_message: None,
|
|
current_mode_id: None,
|
|
_meta: None,
|
|
project_id: None,
|
|
},
|
|
cwd: None,
|
|
models: None,
|
|
modes: None,
|
|
config_options: None,
|
|
acp_client_id: None,
|
|
},
|
|
};
|
|
tx.send(wrap(session_created)).await.unwrap();
|
|
tokio::time::sleep(tokio::time::Duration::from_millis(100)).await;
|
|
|
|
// Event 2-5: Streaming message chunks (AgentMessageChunk)
|
|
println!("2. Sending streaming message chunks...");
|
|
let chunks = vec!["Hello! ", "I'm here to ", "help you with ", "your code."];
|
|
|
|
for (i, chunk) in chunks.iter().enumerate() {
|
|
let chunk_event = Event::SessionUpdate {
|
|
connector_id: connector_id.clone(),
|
|
session_id: session_id.clone(),
|
|
update: SessionUpdate::AgentMessageChunk {
|
|
message_id: message_id.clone(),
|
|
content: ContentBlock::Text {
|
|
text: chunk.to_string(),
|
|
},
|
|
_meta: None,
|
|
},
|
|
};
|
|
tx.send(wrap(chunk_event)).await.unwrap();
|
|
println!(" Chunk {}: {:?}", i + 1, chunk);
|
|
tokio::time::sleep(tokio::time::Duration::from_millis(50)).await;
|
|
}
|
|
|
|
// Event 6: Thinking chunk
|
|
println!("3. Sending thinking chunk...");
|
|
let thinking_event = Event::SessionUpdate {
|
|
connector_id: connector_id.clone(),
|
|
session_id: session_id.clone(),
|
|
update: SessionUpdate::AgentThoughtChunk {
|
|
message_id: message_id.clone(),
|
|
content: ContentBlock::Text {
|
|
text: "Let me consider the best approach...".to_string(),
|
|
},
|
|
_meta: None,
|
|
},
|
|
};
|
|
tx.send(wrap(thinking_event)).await.unwrap();
|
|
tokio::time::sleep(tokio::time::Duration::from_millis(100)).await;
|
|
|
|
// Event 7: Tool call
|
|
println!("4. Sending tool call event...");
|
|
let tool_call_event = Event::SessionUpdate {
|
|
connector_id: connector_id.clone(),
|
|
session_id: session_id.clone(),
|
|
update: SessionUpdate::ToolCall {
|
|
message_id: message_id.clone(),
|
|
tool_call: ToolCall {
|
|
id: "tool_call_123".to_string(),
|
|
tool_name: "read_file".to_string(),
|
|
status: ToolCallStatus::Completed,
|
|
content: vec![],
|
|
raw_input: Some(serde_json::json!({
|
|
"path": "/home/user/project/main.rs"
|
|
})),
|
|
raw_output: Some(serde_json::json!({
|
|
"content": "fn main() { println!(\"Hello\"); }"
|
|
})),
|
|
title: None,
|
|
error: None,
|
|
metadata: None,
|
|
origin: None,
|
|
},
|
|
_meta: None,
|
|
},
|
|
};
|
|
tx.send(wrap(tool_call_event)).await.unwrap();
|
|
tokio::time::sleep(tokio::time::Duration::from_millis(100)).await;
|
|
|
|
// Event 8: MessageCompleted (triggers finalization)
|
|
println!("5. Sending MessageCompleted event...");
|
|
let message_completed = Event::MessageCompleted {
|
|
connector_id: connector_id.clone(),
|
|
message: Message {
|
|
id: message_id.clone(),
|
|
session_id: session_id.clone(),
|
|
role: MessageRole::Assistant,
|
|
created_at: Utc::now(),
|
|
content: vec![MessagePart::Text {
|
|
text: chunks.concat(),
|
|
}],
|
|
status: MessageStatus::Completed,
|
|
metadata: Some(MessageMetadata {
|
|
cost: None,
|
|
tokens_input: None,
|
|
tokens_output: None,
|
|
response_time_ms: None,
|
|
latency_ms: Some(1500),
|
|
model: Some("claude-3-5-sonnet".to_string()),
|
|
other: None,
|
|
}),
|
|
},
|
|
};
|
|
tx.send(wrap(message_completed)).await.unwrap();
|
|
tokio::time::sleep(tokio::time::Duration::from_millis(200)).await;
|
|
|
|
// Event 9: Second message (user response)
|
|
println!("6. Sending user message...");
|
|
let user_message_id = Uuid::now_v7().to_string();
|
|
let user_chunks = vec!["Thanks! ", "Can you explain ", "the code?"];
|
|
|
|
for (i, chunk) in user_chunks.iter().enumerate() {
|
|
let chunk_event = Event::SessionUpdate {
|
|
connector_id: connector_id.clone(),
|
|
session_id: session_id.clone(),
|
|
update: SessionUpdate::UserMessageChunk {
|
|
message_id: user_message_id.clone(),
|
|
content: ContentBlock::Text {
|
|
text: chunk.to_string(),
|
|
},
|
|
_meta: None,
|
|
},
|
|
};
|
|
tx.send(wrap(chunk_event)).await.unwrap();
|
|
println!(" User chunk {}: {:?}", i + 1, chunk);
|
|
tokio::time::sleep(tokio::time::Duration::from_millis(50)).await;
|
|
}
|
|
|
|
let user_message_completed = Event::MessageCompleted {
|
|
connector_id: connector_id.clone(),
|
|
message: Message {
|
|
id: user_message_id.clone(),
|
|
session_id: session_id.clone(),
|
|
role: MessageRole::User,
|
|
created_at: Utc::now(),
|
|
content: vec![MessagePart::Text {
|
|
text: user_chunks.concat(),
|
|
}],
|
|
status: MessageStatus::Completed,
|
|
metadata: None,
|
|
},
|
|
};
|
|
tx.send(wrap(user_message_completed)).await.unwrap();
|
|
tokio::time::sleep(tokio::time::Duration::from_millis(200)).await;
|
|
|
|
// Step 5: Verify archived data
|
|
println!("\n--- Verifying Archived Data ---");
|
|
|
|
// Parse connector_uid from connector_id string
|
|
let connector_uid =
|
|
Uuid::parse_str(&connector_id).expect("connector_id should be a valid UUID");
|
|
|
|
// List sessions
|
|
let page = archivist
|
|
.list_sessions_paged(
|
|
dirigent_archivist::SessionListQuery::default()
|
|
.with_connector(connector_uid)
|
|
.with_limit(100),
|
|
)
|
|
.await?;
|
|
let sessions = page.items;
|
|
println!("Found {} session(s) in archive", sessions.len());
|
|
for session in &sessions {
|
|
println!(" Session: {} - {:?}", session.scroll_id, session.title);
|
|
}
|
|
|
|
// Get messages
|
|
if let Some(session) = sessions.first() {
|
|
let messages = archivist.get_messages(session.scroll_id, None).await?;
|
|
println!("\nFound {} message(s):", messages.len());
|
|
for msg in &messages {
|
|
println!("\n[{}] {} chars", msg.role, msg.content_md.len());
|
|
println!(
|
|
"Content preview: {}",
|
|
&msg.content_md.chars().take(100).collect::<String>()
|
|
);
|
|
}
|
|
}
|
|
|
|
// Step 6: Cleanup
|
|
println!("\n--- Cleanup ---");
|
|
|
|
// Drop the event sender to close the channel
|
|
drop(tx);
|
|
|
|
// Wait for handler to finish
|
|
handler_task.await.expect("Handler task failed");
|
|
|
|
// Remove temporary archive
|
|
std::fs::remove_dir_all(&temp_dir)?;
|
|
println!("Removed temporary archive");
|
|
|
|
println!("\nExample completed successfully!");
|
|
|
|
Ok(())
|
|
}
|