dirigent/crates/dirigent_archivist/examples/basic_usage.rs

//! Basic usage example for dirigent_archivist
//!
//! This example demonstrates:
//! - Creating a Archivist
//! - Registering a connector
//! - Registering a session
//! - Appending messages to a session
//! - Listing sessions for a connector
//! - Retrieving messages for a session

use chrono::Utc;
use dirigent_archivist::{
    Archivist, MessageRecord, RegisterConnectorRequest, RegisterSessionRequest,
    Result,
};
use std::path::PathBuf;
use uuid::Uuid;

#[tokio::main]
async fn main() -> Result<()> {
    // Create a temporary archive directory for this example
    let temp_dir = std::env::temp_dir().join(format!("dirigent_example_{}", Uuid::now_v7()));
    println!("Creating archive at: {}", temp_dir.display());

    // Step 1: Create a Archivist
    let archivist = Archivist::new_with_single_archive(temp_dir.clone()).await?;
    println!("Archivist created successfully");

    // Step 2: Register a connector
    println!("\n--- Registering Connector ---");
    let connector_req = RegisterConnectorRequest {
        r#type: "OpenCode".to_string(),
        title: "OpenCode Local".to_string(),
        client_native_id: "opencode@http://localhost:12225".to_string(),
        custom_uid: None, // Let archivist generate a UID
        metadata: serde_json::json!({
            "version": "0.1.0",
            "protocol": "OpenCode HTTP API"
        }),
        fingerprint: None,
    };

    let connector_resp = archivist.register_connector(connector_req, None).await?;
    println!("Connector registered: {:?}", connector_resp);
    let connector_uid = connector_resp.connector_uid;

    // Step 3: Register a session
    println!("\n--- Registering Session ---");
    let session_req = RegisterSessionRequest {
        connector_uid,
        native_session_id: "session-abc123".to_string(),
        title: Some("Example chat session".to_string()),
        custom_scroll_id: None, // Let archivist generate a scroll ID
        metadata: serde_json::json!({
            "project_path": "/home/user/projects/example",
            "model": "claude-3-5-sonnet"
        }),
        completeness: Default::default(),
        parent_scroll_id: None,
        is_subagent: false,
        continuation: None,
        agent_id: None,
        subagent_type: None,
        spawning_tool_use_id: None,
    };

    let session_resp = archivist.register_session(session_req, None).await?;
    println!("Session registered: {:?}", session_resp);
    let scroll_id = session_resp.scroll_id;

    // Step 4: Append messages to the session
    println!("\n--- Appending Messages ---");

    // User message
    let user_msg = MessageRecord {
        version: 1,
        message_id: Uuid::now_v7(),
        session: scroll_id,
        parent_id: None,
        ts: Utc::now(),
        role: "user".to_string(),
        author: Some("alice".to_string()),
        content_md: "Hello! Can you help me write a function to calculate fibonacci numbers?"
            .to_string(),
        content_parts: None,
        attachments: vec![],
        metadata: serde_json::json!({}),
    };

    // Assistant message
    let assistant_msg = MessageRecord {
        version: 1,
        message_id: Uuid::now_v7(),
        session: scroll_id,
        parent_id: Some(user_msg.message_id),
        ts: Utc::now(),
        role: "assistant".to_string(),
        author: Some("claude".to_string()),
        content_md: r#"Sure! Here's a recursive fibonacci function in Rust:

```rust
fn fibonacci(n: u32) -> u64 {
    match n {
        0 => 0,
        1 => 1,
        _ => fibonacci(n - 1) + fibonacci(n - 2),
    }
}
```

This is the classic recursive implementation, though it's not the most efficient for large values of n."#
            .to_string(),
        content_parts: None,
        attachments: vec![],
        metadata: serde_json::json!({
            "model": "claude-3-5-sonnet",
            "latency_ms": 1245
        }),
    };

    archivist
        .append_messages(scroll_id, vec![user_msg.clone(), assistant_msg.clone()], None)
        .await?;
    println!("Appended 2 messages to session");

    // Step 5: List all sessions for the connector
    println!("\n--- Listing 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) for connector:", sessions.len());
    for session in &sessions {
        println!(
            "  - {} ({}): {:?}",
            session.scroll_id,
            session.created_at.format("%Y-%m-%d %H:%M:%S"),
            session.title
        );
    }

    // Step 6: Retrieve all messages for the session
    println!("\n--- Retrieving Messages ---");
    let messages = archivist.get_messages(scroll_id, None).await?;
    println!("Retrieved {} message(s):", messages.len());
    for msg in &messages {
        println!("\n[{}] {}", msg.role, msg.ts.format("%Y-%m-%d %H:%M:%S"));
        println!("{}", msg.content_md);
    }

    // Step 7: Demonstrate session resolution
    println!("\n--- Resolving Session ---");
    let resolved_scroll_id = archivist
        .resolve_session(connector_uid, "session-abc123", None)
        .await?;
    println!(
        "Resolved native session 'session-abc123' to scroll_id: {}",
        resolved_scroll_id
    );
    assert_eq!(resolved_scroll_id, scroll_id);

    // Step 8: Show archive structure
    println!("\n--- Archive Structure ---");
    println!("Archive root: {}", temp_dir.display());
    println!("\nDirectory structure:");
    show_directory_tree(&temp_dir, 0)?;

    // Cleanup
    println!("\n--- Cleanup ---");
    std::fs::remove_dir_all(&temp_dir)?;
    println!("Removed temporary archive");

    Ok(())
}

/// Helper function to display directory tree
fn show_directory_tree(path: &PathBuf, depth: usize) -> Result<()> {
    let indent = "  ".repeat(depth);

    if path.is_dir() {
        println!("{}{}/", indent, path.file_name().unwrap().to_string_lossy());

        let mut entries: Vec<_> = std::fs::read_dir(path)?.filter_map(|e| e.ok()).collect();
        entries.sort_by_key(|e| e.path());

        for entry in entries {
            show_directory_tree(&entry.path(), depth + 1)?;
        }
    } else {
        println!("{}{}", indent, path.file_name().unwrap().to_string_lossy());
    }

    Ok(())
}