enki-runtime 0.1.4

//! RAG with TOML Configuration Example
//!
//! This example demonstrates how to implement RAG using TOML-configured agents
//! with **separate memory backends per agent**:
//! 1. Load agent and memory configuration from a TOML file
//! 2. Create **per-agent memory backends** from configuration
//! 3. Store documents in agent-specific memory
//! 4. Retrieve relevant documents based on user queries
//! 5. Augment the LLM prompt with retrieved context
//! 6. Generate grounded responses
//!
//! # Per-Agent Memory
//!
//! Each agent can have its own isolated memory backend:
//! - Agents without `[agents.memory]` use the mesh-level `[memory]` config
//! - Agents with `[agents.memory]` get their own isolated storage
//!
//! # Prerequisites
//!
//! 1. Install Ollama from https://ollama.ai
//! 2. Pull the gemma3 model: `ollama pull gemma3:latest`
//! 3. Ensure Ollama is running (default: http://127.0.0.1:11434)
//!
//! # Running
//!
//! ```bash
//! cargo run --example rag_toml_example
//! ```

use enki_runtime::config::MeshConfig;
use enki_runtime::core::agent::AgentContext;
use enki_runtime::core::error::Result;
use enki_runtime::core::memory::{Memory, MemoryEntry, MemoryQuery};
use enki_runtime::llm::LlmAgent;
use enki_runtime::memory::InMemoryBackend;
use enki_runtime::{LlmAgentFromConfig, MemoryFromConfig};
use serde_json::json;
use std::collections::HashMap;

/// A simple document for our knowledge base
struct Document {
    title: String,
    content: String,
    category: String,
}

impl Document {
    fn new(title: &str, content: &str, category: &str) -> Self {
        Self {
            title: title.to_string(),
            content: content.to_string(),
            category: category.to_string(),
        }
    }
}

/// Registry that manages separate memory backends for each agent.
///
/// This struct demonstrates how to give each agent its own isolated memory:
/// - Agents with `[agents.memory]` in TOML get their own configured backend
/// - Agents without it fall back to the mesh-level `[memory]` config
/// - If no config exists anywhere, a default InMemoryBackend is used
struct AgentMemoryRegistry {
    /// Per-agent memory backends, keyed by agent name
    memories: HashMap<String, InMemoryBackend>,
    /// Default memory for agents not in the registry
    default_memory: InMemoryBackend,
}

impl AgentMemoryRegistry {
    /// Create a registry from mesh configuration, creating separate memory for each agent.
    fn from_mesh_config(config: &MeshConfig) -> Self {
        println!("\n📦 Creating per-agent memory backends:");

        // Create default memory from mesh config (or use defaults)
        let default_memory = config
            .memory
            .as_ref()
            .map(|m| {
                println!(
                    "  • Default (mesh): {:?} (max_entries: {:?})",
                    m.backend, m.max_entries
                );
                InMemoryBackend::from_config(m)
            })
            .unwrap_or_else(|| {
                println!("  • Default (mesh): InMemory (defaults)");
                InMemoryBackend::new()
            });

        // Create per-agent memory backends
        let mut memories = HashMap::new();
        for agent in &config.agents {
            let effective_config = agent.get_effective_memory_config(config.memory.as_ref());
            let memory = effective_config
                .map(|m| InMemoryBackend::from_config(m))
                .unwrap_or_else(InMemoryBackend::new);

            // Log what each agent is using
            if agent.memory.is_some() {
                let m = agent.memory.as_ref().unwrap();
                println!(
                    "  • {}: {:?} (max_entries: {:?}) [AGENT-SPECIFIC]",
                    agent.name, m.backend, m.max_entries
                );
            } else if config.memory.is_some() {
                println!("  • {}: uses mesh default", agent.name);
            } else {
                println!("  • {}: InMemory (defaults)", agent.name);
            }

            memories.insert(agent.name.clone(), memory);
        }

        Self {
            memories,
            default_memory,
        }
    }

    /// Get memory for a specific agent.
    fn get(&self, agent_name: &str) -> &InMemoryBackend {
        self.memories
            .get(agent_name)
            .unwrap_or(&self.default_memory)
    }

    /// Get mutable memory for a specific agent.
    fn get_mut(&mut self, agent_name: &str) -> &mut InMemoryBackend {
        self.memories
            .get_mut(agent_name)
            .unwrap_or(&mut self.default_memory)
    }
}

/// A RAG system that uses TOML-configured agents with per-agent memory
struct RagSystem {
    memory_registry: AgentMemoryRegistry,
    agent: LlmAgent,
    agent_name: String,
}

impl RagSystem {
    /// Add a document to this agent's knowledge base
    async fn add_document(&self, doc: Document) -> Result<String> {
        let entry = MemoryEntry::new(&doc.content)
            .with_metadata("title", json!(doc.title))
            .with_metadata("category", json!(doc.category));

        let memory = self.memory_registry.get(&self.agent_name);
        let id = memory.store(entry).await?;
        println!(
            "📄 Stored in {}'s memory: \"{}\" (category: {})",
            self.agent_name, doc.title, doc.category
        );
        Ok(id)
    }

    /// Search for relevant documents in this agent's memory
    async fn search(&self, query: &str, limit: usize) -> Result<Vec<MemoryEntry>> {
        let memory_query = MemoryQuery::new()
            .with_semantic_query(query)
            .with_limit(limit);

        let memory = self.memory_registry.get(&self.agent_name);
        memory.search(memory_query).await
    }

    /// Perform RAG: retrieve context and generate a response
    async fn query(&mut self, question: &str, ctx: &mut AgentContext) -> Result<String> {
        println!("\n🔍 Query: {}", question);
        println!("   Using {}'s memory backend", self.agent_name);

        // Step 1: Retrieve relevant documents from this agent's memory
        let results = self.search(question, 3).await?;

        if results.is_empty() {
            println!("⚠️  No relevant documents found in knowledge base.");
            return Ok("I don't have enough information to answer that question.".to_string());
        }

        // Step 2: Build context from retrieved documents
        println!("\n📚 Retrieved {} relevant document(s):", results.len());
        let mut context = String::new();
        for (i, entry) in results.iter().enumerate() {
            let title = entry
                .metadata
                .get("title")
                .and_then(|v| v.as_str())
                .unwrap_or("Untitled");
            println!("   {}. \"{}\"", i + 1, title);
            context.push_str(&format!(
                "\n--- Document: {} ---\n{}\n",
                title, entry.content
            ));
        }

        // Step 3: Augment the prompt with retrieved context
        let augmented_prompt = format!(
            "Based on the following context, please answer the question.\n\n\
             CONTEXT:\n{}\n\n\
             QUESTION: {}\n\n\
             Please provide a helpful answer based primarily on the context above.",
            context, question
        );

        // Step 4: Generate response using LLM
        println!("\n🤖 Generating response...\n");
        let response = self
            .agent
            .send_message_and_get_response(&augmented_prompt, ctx)
            .await?;

        Ok(response)
    }
}

/// Build a sample knowledge base about Enki Runtime
fn sample_knowledge_base() -> Vec<Document> {
    vec![
        Document::new(
            "Enki Runtime Overview",
            "Enki Runtime is a Rust-based agent mesh framework for building local and distributed \
             AI agent systems. It provides core abstractions like Agent, Memory, Mesh, and Message \
             for building autonomous AI applications. The framework is modular and split into \
             focused sub-crates: Enki-core, Enki-llm, Enki-local, Enki-memory, and \
             Enki-observability.",
            "overview",
        ),
        Document::new(
            "LLM Integration",
            "Enki Runtime supports 13+ LLM providers through a unified interface. Supported \
             providers include OpenAI (GPT-4, GPT-4o), Anthropic (Claude 3, Claude 3.5), \
             Ollama (local models like Llama, Mistral, Gemma), Google (Gemini), DeepSeek, \
             xAI (Grok), Groq, Mistral, Cohere, Phind, and OpenRouter. You can create an \
             LlmAgent using the builder pattern: LlmAgent::builder(name, model).build().",
            "llm",
        ),
        Document::new(
            "Memory System",
            "Enki Runtime provides a pluggable memory system with multiple backends: \
             InMemoryBackend (default, fast, no persistence), SqliteBackend (persistent, \
             requires 'sqlite' feature), and RedisBackend (distributed, requires 'redis' feature). \
             All backends implement the Memory trait with methods like store(), get(), search(), \
             delete(), and clear(). The VectorMemory trait extends Memory with embed() and \
             similarity_search() for semantic search capabilities.",
            "memory",
        ),
        Document::new(
            "Agent Communication",
            "Agents in Enki Runtime communicate through a Mesh. The LocalMesh provides \
             in-process communication between agents. Messages have a topic, payload, and sender. \
             Agents implement the Agent trait with on_start() and on_message() lifecycle methods. \
             You can send messages to specific agents or broadcast to all agents on the mesh.",
            "mesh",
        ),
        Document::new(
            "TOML Configuration",
            "Enki Runtime supports TOML-based agent and memory configuration. Create agents from config \
             files using AgentConfig::from_file() and LlmAgent::from_config(). Memory backends can be \
             configured using MemoryConfig with options for backend type, path (sqlite), url (redis), \
             max_entries, and ttl_seconds. Each agent can have its own [agents.memory] section for \
             isolated storage, or fall back to the mesh-level [memory] configuration.",
            "config",
        ),
    ]
}

#[tokio::main]
async fn main() -> Result<()> {
    println!("=== Enki Runtime - RAG with Per-Agent Memory ===\n");
    println!("This example demonstrates RAG using TOML-configured agents");
    println!("with SEPARATE memory backends per agent:");
    println!("  1. Load agent and memory configuration from TOML file");
    println!("  2. Create per-agent memory backends from configuration");
    println!("  3. Store documents in agent-specific memory");
    println!("  4. Retrieve relevant documents for a query");
    println!("  5. Augment the LLM prompt with context");
    println!("  6. Generate grounded responses\n");

    // 1. Load configuration from TOML file
    let config_path = concat!(env!("CARGO_MANIFEST_DIR"), "/examples/rag_agents.toml");
    println!("📋 Loading configuration from: {}\n", config_path);

    let mesh_config = match MeshConfig::from_file(config_path) {
        Ok(config) => config,
        Err(e) => {
            eprintln!("Failed to load configuration: {}", e);
            eprintln!("Make sure rag_agents.toml exists in the examples directory.");
            return Err(e);
        }
    };

    println!("Mesh name: {}", mesh_config.name);
    println!(
        "Available agents: {:?}",
        mesh_config
            .agents
            .iter()
            .map(|a| &a.name)
            .collect::<Vec<_>>()
    );

    // 2. Create per-agent memory registry from configuration
    let memory_registry = AgentMemoryRegistry::from_mesh_config(&mesh_config);

    // 3. Get the RAG agent configuration
    let rag_config = match mesh_config.get_agent("rag_agent") {
        Some(config) => config.clone(),
        None => {
            eprintln!("RAG agent not found in configuration!");
            return Err(enki_runtime::core::error::Error::ConfigError(
                "RAG agent 'rag_agent' not found in configuration".to_string(),
            ));
        }
    };

    println!("\n🤖 Agent configuration:");
    println!("  Name: {}", rag_config.name);
    println!("  Model: {}", rag_config.model);
    println!("  Temperature: {:?}", rag_config.temperature);
    println!("  Max tokens: {:?}", rag_config.max_tokens);
    println!(
        "  Memory: {}",
        if rag_config.memory.is_some() {
            "agent-specific"
        } else {
            "mesh default"
        }
    );

    // 4. Create the agent from config
    let agent = match LlmAgent::from_config(rag_config.clone()) {
        Ok(agent) => agent,
        Err(e) => {
            eprintln!("Failed to create agent: {}", e);
            eprintln!("Make sure Ollama is running and gemma3:latest is available.");
            eprintln!("Pull the model with: ollama pull gemma3:latest");
            return Err(e);
        }
    };

    // 5. Create the RAG system with per-agent memory
    let mut rag = RagSystem {
        memory_registry,
        agent,
        agent_name: rag_config.name.clone(),
    };

    println!("\n✓ RAG system created with per-agent memory\n");

    // 6. Build knowledge base in the RAG agent's memory
    println!(
        "📖 Building knowledge base in {}'s memory...\n",
        rag.agent_name
    );
    for doc in sample_knowledge_base() {
        rag.add_document(doc).await?;
    }

    println!("\n✓ Knowledge base ready with {} documents\n", 5);

    // 7. Create agent context
    let mut ctx = AgentContext::new(mesh_config.name.clone(), None);

    // 8. Example queries that demonstrate RAG
    let queries = vec![
        "What LLM providers does Enki Runtime support?",
        "How do agents communicate with each other?",
        "What are the available memory backends?",
    ];

    for query in queries {
        println!("\n{}", "=".repeat(60));
        match rag.query(query, &mut ctx).await {
            Ok(response) => {
                println!("💬 Response:\n{}", response);
            }
            Err(e) => {
                eprintln!("❌ Error: {}", e);
            }
        }
        println!("{}\n", "=".repeat(60));
    }

    println!("\n=== RAG with Per-Agent Memory Complete ===");
    Ok(())
}