cerebro 1.1.8

//! # Swarm Orchestration Patterns
//!
//! Defines the trait and built-in patterns for multi-agent coordination.

pub mod hierarchical;
pub mod parallel;
pub mod sequential;
pub mod templates;

use async_trait::async_trait;
use std::collections::HashMap;

use super::agent::AgentConfig;
use super::llm::LlmClient;
use super::memory_bus::CerebroMemoryBus;
use super::trace::ExecutionTracer;
use crate::traits::Result;

/// The result of a completed swarm execution.
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
pub struct SwarmResult {
    /// The final consolidated output from the swarm.
    pub final_output: String,
    /// The full execution trace.
    pub trace: super::trace::ExecutionTrace,
    /// Total tokens consumed across all agents.
    pub total_tokens: usize,
    /// Total wall-clock duration in milliseconds.
    pub total_duration_ms: u64,
}

/// Trait that all swarm orchestration patterns must implement.
#[async_trait]
pub trait SwarmPatternExecutor: Send + Sync {
    /// Execute the swarm pattern with the given agents, input, and shared resources.
    async fn execute(
        &self,
        agents: &HashMap<String, AgentConfig>,
        tools: &HashMap<String, std::sync::Arc<dyn super::tools::AgentTool>>,
        memory: &CerebroMemoryBus,
        tracer: &ExecutionTracer,
        llm: &LlmClient,
        input: &str,
    ) -> Result<SwarmResult>;
}

/// Helper to execute an agent to completion, handling tool calls within a loop.
#[allow(clippy::too_many_arguments)]
pub async fn execute_agent_loop(
    agent_id: &str,
    agent_name: &str,
    config: &super::agent::AgentConfig,
    tools: &HashMap<String, std::sync::Arc<dyn super::tools::AgentTool>>,
    memory: &CerebroMemoryBus,
    tracer: &ExecutionTracer,
    llm: &LlmClient,
    pipeline_step: usize,
    runtime: &mut super::agent::AgentRuntime,
) -> Result<(String, usize, u64)> {
    let mut total_tokens = 0;
    let mut total_duration = 0;

    loop {
        if runtime.is_exhausted() {
            return Ok((
                "Error: Max steps exceeded".into(),
                total_tokens,
                total_duration,
            ));
        }

        let start = std::time::Instant::now();
        let response = llm.chat(&config.model, &runtime.messages).await?;
        let duration = start.elapsed().as_millis() as u64;
        let step_tokens = response.input_tokens + response.output_tokens;
        total_tokens += step_tokens;
        total_duration += duration;

        tracer.record_llm_call(
            agent_id,
            agent_name,
            pipeline_step,
            "Internal step", // For concise trace
            &response.content,
            duration,
            step_tokens,
        );

        memory.push_message(
            agent_id,
            super::agent::ChatMessage::new(super::agent::Role::Assistant, &response.content)
                .with_metadata("tokens", step_tokens.to_string())
                .with_metadata("duration_ms", duration.to_string()),
        );

        if let Some(tool_req) = &response.tool_call {
            let tool_output = if let Some(tool) = tools.get(&tool_req.tool_name) {
                // Execute the tool
                tool.execute(tool_req.arguments.clone())
                    .await
                    .unwrap_or_else(|e| format!("Error executing tool: {}", e))
            } else {
                format!("Error: tool '{}' not found", tool_req.tool_name)
            };

            // Feed it back to the runtime and loop again
            let tool_msg = super::agent::ChatMessage::new(super::agent::Role::Tool, &tool_output);
            runtime.push_message(tool_msg.clone());
            memory.push_message(agent_id, tool_msg);

            // Advance internal step counter
            runtime.advance_step();
            continue;
        } else {
            return Ok((response.content, total_tokens, total_duration));
        }
    }
}