car_multi/patterns/

map_reduce.rs

//! Map-Reduce — split a task into N items, run agents in parallel, reduce results.
//!
//! The mapper spec is cloned for each item. Each mapper processes one item.
//! The reducer combines all mapper outputs into a single result.

use crate::error::MultiError;
use crate::mailbox::Mailbox;
use crate::runner::AgentRunner;
use crate::shared::SharedInfra;
use crate::types::{AgentOutput, AgentSpec};
use serde::{Deserialize, Serialize};
use std::sync::Arc;

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct MapReduceResult {
    pub task: String,
    pub map_outputs: Vec<AgentOutput>,
    pub reduced_answer: String,
}

impl MapReduceResult {
    pub fn all_succeeded(&self) -> bool {
        self.map_outputs.iter().all(|o| o.succeeded())
    }
}

pub struct MapReduce {
    pub mapper: AgentSpec,
    pub reducer: AgentSpec,
    pub max_concurrent: usize,
}

impl MapReduce {
    pub fn new(mapper: AgentSpec, reducer: AgentSpec) -> Self {
        Self {
            mapper,
            reducer,
            max_concurrent: 5,
        }
    }

    pub fn with_max_concurrent(mut self, n: usize) -> Self {
        self.max_concurrent = n;
        self
    }

    pub async fn run(
        &self,
        task: &str,
        items: &[String],
        runner: &Arc<dyn AgentRunner>,
        infra: &SharedInfra,
    ) -> Result<MapReduceResult, MultiError> {
        // Map phase: one agent per item, bounded concurrency
        let semaphore = Arc::new(tokio::sync::Semaphore::new(self.max_concurrent));
        let mut handles = Vec::new();

        for (i, item) in items.iter().enumerate() {
            let sem = Arc::clone(&semaphore);
            let runner = Arc::clone(runner);
            let rt = infra.make_runtime();
            let mailbox = Mailbox::default();

            let mut spec = self.mapper.clone();
            spec.name = format!("{}_{}", self.mapper.name, i);

            for tool in &spec.tools {
                rt.register_tool(tool).await;
            }

            let subtask = format!("{}\n\nProcess this item: {}", task, item);

            handles.push(tokio::spawn(async move {
                let _permit = sem.acquire().await.unwrap();
                (i, runner.run(&spec, &subtask, &rt, &mailbox).await)
            }));
        }

        let results = futures::future::join_all(handles).await;
        let mut indexed: Vec<(usize, AgentOutput)> = Vec::new();

        for result in results {
            match result {
                Ok((i, Ok(output))) => indexed.push((i, output)),
                Ok((i, Err(e))) => {
                    indexed.push((
                        i,
                        AgentOutput {
                            name: format!("{}_{}", self.mapper.name, i),
                            answer: String::new(),
                            turns: 0,
                            tool_calls: 0,
                            duration_ms: 0.0,
                            error: Some(e.to_string()),
                        },
                    ));
                }
                Err(e) => {
                    indexed.push((
                        indexed.len(),
                        AgentOutput {
                            name: "unknown".to_string(),
                            answer: String::new(),
                            turns: 0,
                            tool_calls: 0,
                            duration_ms: 0.0,
                            error: Some(format!("join error: {}", e)),
                        },
                    ));
                }
            }
        }

        indexed.sort_by_key(|(i, _)| *i);
        let map_outputs: Vec<AgentOutput> = indexed.into_iter().map(|(_, o)| o).collect();

        // Reduce phase
        let summaries: Vec<String> = map_outputs
            .iter()
            .filter(|o| o.succeeded())
            .map(|o| format!("- [{}] {}", o.name, truncate(&o.answer, 300)))
            .collect();

        let reduce_task = format!(
            "Original task: {}\n\nResults from {} sub-agents:\n{}\n\n\
             Combine these into a single coherent result.",
            task,
            map_outputs.len(),
            summaries.join("\n")
        );

        let rt = infra.make_runtime();
        let mailbox = Mailbox::default();
        let reduced = runner
            .run(&self.reducer, &reduce_task, &rt, &mailbox)
            .await
            .map(|o| o.answer)
            .unwrap_or_default();

        Ok(MapReduceResult {
            task: task.to_string(),
            map_outputs,
            reduced_answer: reduced,
        })
    }
}

fn truncate(s: &str, max_len: usize) -> &str {
    if s.len() <= max_len {
        return s;
    }
    let mut end = max_len;
    while end > 0 && !s.is_char_boundary(end) {
        end -= 1;
    }
    &s[..end]
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::types::{AgentOutput, AgentSpec};
    use car_engine::Runtime;

    struct CountRunner;

    #[async_trait::async_trait]
    impl crate::runner::AgentRunner for CountRunner {
        async fn run(
            &self,
            spec: &AgentSpec,
            _task: &str,
            _runtime: &Runtime,
            _mailbox: &Mailbox,
        ) -> Result<AgentOutput, MultiError> {
            Ok(AgentOutput {
                name: spec.name.clone(),
                answer: format!("{} processed", spec.name),
                turns: 1,
                tool_calls: 0,
                duration_ms: 5.0,
                error: None,
            })
        }
    }

    #[tokio::test]
    async fn test_map_reduce() {
        let mapper = AgentSpec::new("summarizer", "Summarize the file");
        let reducer = AgentSpec::new("combiner", "Combine summaries");
        let items: Vec<String> = vec!["file_a.rs", "file_b.rs", "file_c.rs"]
            .into_iter()
            .map(String::from)
            .collect();

        let runner: Arc<dyn crate::runner::AgentRunner> = Arc::new(CountRunner);
        let infra = SharedInfra::new();

        let result = MapReduce::new(mapper, reducer)
            .run("summarize codebase", &items, &runner, &infra)
            .await
            .unwrap();

        assert_eq!(result.map_outputs.len(), 3);
        assert!(!result.reduced_answer.is_empty());
    }
}