Building Multi-Agent Systems with ADK-Rust

use adk_rust::prelude::*;
use std::sync::Arc;

#[tokio::main]
async fn main() -> anyhow::Result<()> {
    let api_key = std::env::var("GOOGLE_API_KEY")?;
    let model = Arc::new(GeminiModel::new(&api_key, "gemini-2.0-flash")?);

    // Specialist: Billing Agent
    // Clear description helps the coordinator know when to transfer
    let billing_agent = LlmAgentBuilder::new("billing_agent")
        .description(
            "Handles all billing questions: invoices, payments, refunds,              subscription plans, and account charges. Transfer here for              any money-related questions."
        )
        .instruction(
            "You are a billing specialist. Answer questions about invoices,              payments, and subscription plans. Be concise and accurate.              If asked about technical issues, suggest transferring to support."
        )
        .model(model.clone())
        .build()?;

    // Specialist: Technical Support Agent
    let support_agent = LlmAgentBuilder::new("support_agent")
        .description(
            "Provides technical support: troubleshooting, bug reports,              feature questions, and integration help. Transfer here for              any technical problems."
        )
        .instruction(
            "You are a technical support specialist. Help users troubleshoot              issues step by step. Ask clarifying questions when needed.              Be patient and thorough."
        )
        .model(model.clone())
        .build()?;

    // Coordinator: Routes to specialists
    let coordinator = LlmAgentBuilder::new("coordinator")
        .description("Customer service coordinator")
        .instruction(
            "You are a friendly customer service coordinator. Your job is to:\n             1. Greet users warmly\n             2. Understand their needs\n             3. Route to the right specialist:\n                - Billing questions → transfer to billing_agent\n                - Technical issues → transfer to support_agent\n             4. Handle general questions yourself\n\n             Always explain who you're connecting them with."
        )
        .model(model.clone())
        .sub_agent(Arc::new(billing_agent))
        .sub_agent(Arc::new(support_agent))
        .build()?;

    // Run with the built-in launcher
    Launcher::new(Arc::new(coordinator))
        .run()
        .await?;

    Ok(())
}

use adk_agent::LlmAgentBuilder;
use adk_tool::{AgentTool, FunctionTool};
use adk_core::ToolContext;
use serde_json::{json, Value};
use std::sync::Arc;

// Calculator tool for the math agent
async fn calculator(
    _ctx: Arc<dyn ToolContext>,
    args: Value
) -> Result<Value, adk_core::AdkError> {
    let operation = args["operation"].as_str().unwrap_or("add");
    let a = args["a"].as_f64().unwrap_or(0.0);
    let b = args["b"].as_f64().unwrap_or(0.0);

    let result = match operation {
        "add" => a + b,
        "multiply" => a * b,
        "percent" => a * (b / 100.0),
        _ => return Err(adk_core::AdkError::Tool(
            format!("Unknown operation: {}", operation)
        )),
    };

    Ok(json!({ "result": result }))
}

#[tokio::main]
async fn main() -> anyhow::Result<()> {
    let api_key = std::env::var("GOOGLE_API_KEY")?;
    let model = Arc::new(GeminiModel::new(&api_key, "gemini-2.5-flash")?);

    // Create calculator tool
    let calc_tool = FunctionTool::new(
        "calculator",
        "Performs arithmetic: add, multiply, percent.          Args: operation (string), a (number), b (number)",
        calculator,
    );

    // Math Expert agent - has its own tools
    let math_agent = LlmAgentBuilder::new("math_expert")
        .description(
            "A math expert that performs calculations. Use for any              math-related questions, percentages, or numerical analysis."
        )
        .instruction(
            "You are a math expert. Use the calculator tool for              calculations. Show your work step by step. Be precise."
        )
        .model(model.clone())
        .tool(Arc::new(calc_tool))
        .build()?;

    // Trivia Expert agent - uses LLM knowledge
    let trivia_agent = LlmAgentBuilder::new("trivia_expert")
        .description(
            "A trivia and history expert. Use for questions about              historical facts, pop culture, science facts, and trivia."
        )
        .instruction(
            "You are a trivia expert with vast knowledge across domains.              Answer questions accurately and include interesting related facts."
        )
        .model(model.clone())
        .build()?;

    // Wrap agents as tools with configuration
    let math_tool = AgentTool::new(Arc::new(math_agent))
        .skip_summarization(false)   // Summarize lengthy responses
        .forward_artifacts(true);     // Pass through any generated files

    let trivia_tool = AgentTool::new(Arc::new(trivia_agent))
        .skip_summarization(false);

    // Coordinator uses agents as tools
    let coordinator = LlmAgentBuilder::new("coordinator")
        .description("Smart assistant that combines expert knowledge")
        .instruction(
            "You are a helpful assistant with access to expert agents:\n             - math_expert: For calculations and math problems\n             - trivia_expert: For facts, history, and trivia\n\n             When questions span multiple domains, call multiple experts              and synthesize their responses into a cohesive answer."
        )
        .model(model)
        .tool(Arc::new(math_tool))
        .tool(Arc::new(trivia_tool))
        .build()?;

    Launcher::new(Arc::new(coordinator)).run().await?;
    Ok(())
}

use adk_agent::LlmAgentBuilder;
use adk_graph::{
    StateGraph,
    edge::{START, END},
    node::{AgentNode, ExecutionConfig, NodeOutput},
    state::State,
};
use adk_model::GeminiModel;
use serde_json::json;
use std::sync::Arc;

#[tokio::main]
async fn main() -> anyhow::Result<()> {
    let api_key = std::env::var("GOOGLE_API_KEY")?;
    let model = Arc::new(GeminiModel::new(&api_key, "gemini-2.0-flash")?);

    // Supervisor: Decides which worker should act next
    let supervisor = LlmAgentBuilder::new("supervisor")
        .description("Routes tasks to specialized workers")
        .instruction(
            "You are a task supervisor. Based on the task and work done so far,              decide who should work next.\n\n             Workers available:\n             - researcher: Gathers information and facts\n             - writer: Writes content based on research\n             - coder: Creates code examples\n\n             Respond with ONLY one word: 'researcher', 'writer', 'coder', or 'done'."
        )
        .model(model.clone())
        .build()?;

    // Workers with specialized roles
    let researcher = LlmAgentBuilder::new("researcher")
        .instruction("Research the topic. Provide key facts as bullet points.")
        .model(model.clone())
        .build()?;

    let writer = LlmAgentBuilder::new("writer")
        .instruction("Write engaging content based on the research provided.")
        .model(model.clone())
        .build()?;

    let coder = LlmAgentBuilder::new("coder")
        .instruction("Write clean, documented code examples for the topic.")
        .model(model.clone())
        .build()?;

    // Create AgentNodes with input/output mappers
    let supervisor_node = AgentNode::new(Arc::new(supervisor))
        .with_input_mapper(|state| {
            let task = state.get("task").and_then(|v| v.as_str()).unwrap_or("");
            let history = state.get("history")
                .and_then(|v| v.as_array())
                .map(|arr| arr.iter()
                    .filter_map(|h| h.get("agent").and_then(|a| a.as_str()))
                    .map(|s| format!("- {} completed", s))
                    .collect::<Vec<_>>()
                    .join("\n"))
                .unwrap_or_default();

            adk_core::Content::new("user").with_text(format!(
                "Task: {}\n\nWork completed:\n{}\n\nWho next?",
                task,
                if history.is_empty() { "None yet" } else { &history }
            ))
        })
        .with_output_mapper(|events| {
            let mut updates = std::collections::HashMap::new();
            for event in events {
                if let Some(content) = event.content() {
                    let text: String = content.parts.iter()
                        .filter_map(|p| p.text())
                        .collect();
                    let next = if text.to_lowercase().contains("researcher") {
                        "researcher"
                    } else if text.to_lowercase().contains("writer") {
                        "writer"
                    } else if text.to_lowercase().contains("coder") {
                        "coder"
                    } else {
                        "done"
                    };
                    updates.insert("next_agent".to_string(), json!(next));
                }
            }
            updates
        });

    // Build the graph
    let graph = StateGraph::with_channels(&[
        "task", "next_agent", "history",
        "research_output", "written_content", "code_output"
    ])
    .add_node(supervisor_node)
    .add_node(AgentNode::new(Arc::new(researcher)))
    .add_node(AgentNode::new(Arc::new(writer)))
    .add_node(AgentNode::new(Arc::new(coder)))
    // Finalize node compiles all outputs
    .add_node_fn("finalize", |ctx| async move {
        let research = ctx.get("research_output").and_then(|v| v.as_str());
        let content = ctx.get("written_content").and_then(|v| v.as_str());
        let code = ctx.get("code_output").and_then(|v| v.as_str());

        let result = format!(
            "=== FINAL OUTPUT ===\n\n{}\n\n{}\n\n{}",
            research.unwrap_or("No research"),
            content.unwrap_or("No content"),
            code.unwrap_or("No code")
        );
        Ok(NodeOutput::new().with_update("final_result", json!(result)))
    })
    // Graph structure
    .add_edge(START, "supervisor")
    .add_conditional_edges(
        "supervisor",
        |state| state.get("next_agent")
            .and_then(|v| v.as_str())
            .unwrap_or("done")
            .to_string(),
        [
            ("researcher", "researcher"),
            ("writer", "writer"),
            ("coder", "coder"),
            ("done", "finalize"),
        ],
    )
    // Workers cycle back to supervisor
    .add_edge("researcher", "supervisor")
    .add_edge("writer", "supervisor")
    .add_edge("coder", "supervisor")
    .add_edge("finalize", END)
    .compile()?
    .with_recursion_limit(15);  // Prevent infinite loops

    // Execute
    let mut input = State::new();
    input.insert("task".to_string(), json!("Create a guide about Rust error handling"));
    input.insert("history".to_string(), json!([]));

    let result = graph.invoke(input, ExecutionConfig::new("content-thread")).await?;
    println!("{}", result.get("final_result").and_then(|v| v.as_str()).unwrap_or(""));

    Ok(())
}