enki-runtime 0.1.4

//! LLM Multi-Agent Mesh Example - Two LLM agents communicating via LocalMesh
//!
//! This example demonstrates how to create two LLM agents that communicate
//! with each other through Enki's LocalMesh infrastructure.
//!
//! # Architecture
//!
//! - **Researcher Agent**: Uses Ollama gemma3:latest to research topics
//! - **Summarizer Agent**: Uses Ollama gemma3:latest to summarize research
//!
//! The Researcher sends its findings to the Summarizer, which then creates
//! a concise summary.
//!
//! # 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 llm_multi_agent_mesh
//! ```

use async_trait::async_trait;
use enki_runtime::core::agent::{Agent, AgentContext};
use enki_runtime::core::error::Result;
use enki_runtime::core::mesh::Mesh;
use enki_runtime::core::message::Message;
use enki_runtime::llm::LlmAgent;
use enki_runtime::local::LocalMesh;
use std::sync::Arc;
use std::time::Duration;
use tokio::sync::Notify;

// --- Researcher Agent Wrapper ---
// Wraps an LlmAgent to add mesh communication capabilities
struct ResearcherAgent {
    llm_agent: LlmAgent,
    mesh: Arc<LocalMesh>,
}

impl ResearcherAgent {
    fn new(mesh: Arc<LocalMesh>) -> Result<Self> {
        let llm_agent = LlmAgent::builder("researcher", "ollama::gemma3:latest")
            .with_system_prompt(
                "You are a research assistant. When given a topic, provide detailed \
                 information and key facts about it. Be thorough but focused. \
                 Limit your response to 3-4 paragraphs.",
            )
            .with_temperature(0.7)
            .with_max_tokens(1024)
            .build()?;

        Ok(Self { llm_agent, mesh })
    }
}

#[async_trait]
impl Agent for ResearcherAgent {
    fn name(&self) -> String {
        "researcher".to_string()
    }

    async fn on_start(&mut self, _ctx: &mut AgentContext) -> Result<()> {
        println!("[Researcher] Started and ready for research tasks.");
        Ok(())
    }

    async fn on_message(&mut self, msg: Message, ctx: &mut AgentContext) -> Result<()> {
        let topic = String::from_utf8_lossy(&msg.payload);
        println!("\n[Researcher] Received research request: {}", topic);

        // Use LLM to research the topic
        let research_prompt = format!(
            "Research the following topic and provide key information: {}",
            topic
        );

        match self
            .llm_agent
            .send_message_and_get_response(&research_prompt, ctx)
            .await
        {
            Ok(research_result) => {
                println!("[Researcher] Research complete. Sending to summarizer...");

                // Send research results to the summarizer
                let research_msg =
                    Message::new("research_result", research_result.into_bytes(), self.name());
                self.mesh.send(research_msg, "summarizer").await?;
            }
            Err(e) => {
                eprintln!("[Researcher] Error during research: {}", e);
                // Send error message to summarizer
                let error_msg = Message::new(
                    "error",
                    format!("Research failed: {}", e).into_bytes(),
                    self.name(),
                );
                self.mesh.send(error_msg, "summarizer").await?;
            }
        }

        Ok(())
    }
}

// --- Summarizer Agent Wrapper ---
// Wraps an LlmAgent to summarize research findings
struct SummarizerAgent {
    llm_agent: LlmAgent,
    completion_notify: Arc<Notify>,
}

impl SummarizerAgent {
    fn new(completion_notify: Arc<Notify>) -> Result<Self> {
        let llm_agent = LlmAgent::builder("summarizer", "ollama::gemma3:latest")
            .with_system_prompt(
                "You are a summarization expert. When given research content, \
                 create a clear and concise summary with bullet points highlighting \
                 the most important facts. Keep the summary to 5-7 bullet points.",
            )
            .with_temperature(0.5)
            .with_max_tokens(512)
            .build()?;

        Ok(Self {
            llm_agent,
            completion_notify,
        })
    }
}

#[async_trait]
impl Agent for SummarizerAgent {
    fn name(&self) -> String {
        "summarizer".to_string()
    }

    async fn on_start(&mut self, _ctx: &mut AgentContext) -> Result<()> {
        println!("[Summarizer] Started and ready to summarize.");
        Ok(())
    }

    async fn on_message(&mut self, msg: Message, ctx: &mut AgentContext) -> Result<()> {
        let content = String::from_utf8_lossy(&msg.payload);
        println!("\n[Summarizer] Received content from {}", msg.sender);

        if msg.topic == "error" {
            println!("[Summarizer] Received error: {}", content);
            self.completion_notify.notify_one();
            return Ok(());
        }

        // Use LLM to summarize the research
        let summary_prompt = format!(
            "Please summarize the following research content into clear bullet points:\n\n{}",
            content
        );

        match self
            .llm_agent
            .send_message_and_get_response(&summary_prompt, ctx)
            .await
        {
            Ok(summary) => {
                println!("\n========================================");
                println!("           FINAL SUMMARY");
                println!("========================================\n");
                println!("{}", summary);
                println!("\n========================================\n");
            }
            Err(e) => {
                eprintln!("[Summarizer] Error during summarization: {}", e);
            }
        }

        // Signal completion
        self.completion_notify.notify_one();
        Ok(())
    }
}

// --- Coordinator Agent ---
// Initiates the workflow by sending a research topic
struct CoordinatorAgent {
    mesh: Arc<LocalMesh>,
    topic: String,
}

#[async_trait]
impl Agent for CoordinatorAgent {
    fn name(&self) -> String {
        "coordinator".to_string()
    }

    async fn on_start(&mut self, _ctx: &mut AgentContext) -> Result<()> {
        println!("[Coordinator] Starting multi-agent workflow...");
        println!("[Coordinator] Topic: {}\n", self.topic);

        // Small delay to ensure all agents are ready
        tokio::time::sleep(Duration::from_millis(100)).await;

        // Send topic to researcher
        let task_msg = Message::new(
            "research_topic",
            self.topic.as_bytes().to_vec(),
            self.name(),
        );
        self.mesh.send(task_msg, "researcher").await?;

        Ok(())
    }
}

#[tokio::main]
async fn main() -> Result<()> {
    println!("=== Enki Runtime - LLM Multi-Agent Mesh Example ===\n");
    println!("This example demonstrates two LLM agents communicating via LocalMesh:");
    println!("  1. Researcher Agent - Researches a given topic");
    println!("  2. Summarizer Agent - Summarizes the research findings\n");

    // 1. Create the Mesh
    let mesh = Arc::new(LocalMesh::new("llm_mesh"));

    // 2. Create completion notifier
    let completion = Arc::new(Notify::new());

    // 3. Create LLM-powered agents
    let researcher = match ResearcherAgent::new(mesh.clone()) {
        Ok(agent) => agent,
        Err(e) => {
            eprintln!("Failed to create Researcher agent: {}", e);
            eprintln!("Make sure Ollama is running and gemma3:latest is available.");
            return Err(e);
        }
    };

    let summarizer = match SummarizerAgent::new(completion.clone()) {
        Ok(agent) => agent,
        Err(e) => {
            eprintln!("Failed to create Summarizer agent: {}", e);
            eprintln!("Make sure Ollama is running and gemma3:latest is available.");
            return Err(e);
        }
    };

    let coordinator = CoordinatorAgent {
        mesh: mesh.clone(),
        topic: "The benefits of Rust programming language for system development".to_string(),
    };

    // 4. Register agents with the mesh
    mesh.add_agent(Box::new(researcher)).await?;
    mesh.add_agent(Box::new(summarizer)).await?;
    mesh.add_agent(Box::new(coordinator)).await?;

    println!("✓ All agents registered successfully\n");

    // 5. Start the mesh (this triggers on_start for all agents)
    mesh.start().await?;

    // 6. Wait for the workflow to complete
    println!("Waiting for multi-agent workflow to complete...\n");
    tokio::select! {
        _ = completion.notified() => {
            println!("Multi-agent workflow completed successfully!");
        }
        _ = tokio::time::sleep(Duration::from_secs(120)) => {
            println!("Timeout waiting for workflow completion.");
        }
    }

    // 7. Stop the mesh
    mesh.stop().await?;

    println!("\n=== Example finished ===");
    Ok(())
}