ares_agent/
orchestrator.rs

1use crate::{Agent, AgentConfig, AgentRegistry, AgentResponse};
2use ares_llm::LLMClient;
3use ares_types::types::{AgentContext, AgentType, AppError, Result};
4use async_trait::async_trait;
5use std::future::Future;
6use std::sync::Arc;
7
8/// Orchestrator agent that coordinates multiple specialized agents.
9///
10/// This agent decomposes complex queries into subtasks and delegates
11/// them to appropriate specialized agents via the AgentRegistry.
12pub struct OrchestratorAgent {
13    llm: Box<dyn LLMClient>,
14    agents: Arc<std::collections::HashMap<String, AgentConfig>>,
15    agent_registry: Arc<AgentRegistry>,
16}
17
18impl OrchestratorAgent {
19    /// Creates a new OrchestratorAgent with the given dependencies.
20    pub fn new(
21        llm: Box<dyn LLMClient>,
22        agents: Arc<std::collections::HashMap<String, AgentConfig>>,
23        agent_registry: Arc<AgentRegistry>,
24    ) -> Self {
25        Self {
26            llm,
27            agents,
28            agent_registry,
29        }
30    }
31
32    /// Decompose a complex task into subtasks for specialized agents
33    async fn decompose_task(&self, input: &str) -> Result<Vec<(String, String)>> {
34        // Get available agents from registry
35        let available_agents = self.agent_registry.agent_names();
36        let agent_list = available_agents
37            .iter()
38            .filter(|name| **name != "orchestrator" && **name != "router")
39            .cloned()
40            .collect::<Vec<_>>()
41            .join(", ");
42
43        let system_prompt = format!(
44            r#"You are a task decomposition agent. Break down complex queries into subtasks for specialized agents.
45
46Available agents: {}
47
48Return a JSON array of tasks:
49[
50    {{"agent": "sales", "task": "Get Q1 revenue"}},
51    {{"agent": "product", "task": "List top products"}}
52]
53
54Only respond with valid JSON."#,
55            agent_list
56        );
57
58        let response = self.llm.generate_with_system(&system_prompt, input).await?;
59
60        // Parse JSON response
61        let tasks: Vec<serde_json::Value> = serde_json::from_str(&response)
62            .map_err(|e| AppError::LLM(format!("Failed to parse tasks: {}", e)))?;
63
64        let mut result = Vec::new();
65        for task in tasks {
66            let agent_name = task["agent"].as_str().unwrap_or("product").to_string();
67            let task_str = task["task"].as_str().unwrap_or("").to_string();
68
69            // Validate agent exists in registry
70            if self.agent_registry.has_agent(&agent_name) {
71                result.push((agent_name, task_str));
72            } else {
73                // Fall back to product agent if unknown
74                result.push(("product".to_string(), task_str));
75            }
76        }
77
78        Ok(result)
79    }
80
81    /// Execute a subtask using the appropriate agent from the registry
82    async fn execute_subtask(
83        &self,
84        agent_name: &str,
85        task: &str,
86        context: &AgentContext,
87    ) -> Result<String> {
88        // Create agent from registry (handles model and tool configuration)
89        let agent = self.agent_registry.create_agent(agent_name).await?;
90        let resp = agent.execute(task, context).await?;
91        Ok(resp.content)
92    }
93}
94
95/// Join fallible subtask futures concurrently, preserving input order.
96/// Returns the first error via `try_join_all`.
97pub(crate) async fn join_subtask_results<T, E, Fut>(
98    futs: Vec<Fut>,
99) -> std::result::Result<Vec<T>, E>
100where
101    Fut: Future<Output = std::result::Result<T, E>>,
102{
103    futures::future::try_join_all(futs).await
104}
105
106#[async_trait]
107impl Agent for OrchestratorAgent {
108    async fn execute(&self, input: &str, context: &AgentContext) -> Result<AgentResponse> {
109        // Decompose the task into subtasks
110        let subtasks = self.decompose_task(input).await?;
111
112        if subtasks.is_empty() {
113            let content = self.llm.generate(input).await?;
114            return Ok(AgentResponse { content, usage: None, metadata: None });
115        }
116
117        // Execute subtasks concurrently via try_join_all
118        let futs = subtasks
119            .into_iter()
120            .map(|(agent_name, task)| async move {
121                let result = self.execute_subtask(&agent_name, &task, context).await?;
122                Ok::<_, AppError>(format!("[{}] {}", agent_name, result))
123            })
124            .collect();
125        let results = join_subtask_results(futs).await?;
126
127        // Synthesize results into final response
128        let synthesis_prompt = format!(
129            "Original query: {}\n\nSubtask results:\n{}\n\nProvide a comprehensive answer:",
130            input,
131            results.join("\n\n")
132        );
133
134        let content = self.llm.generate(&synthesis_prompt).await?;
135        Ok(AgentResponse { content, usage: None, metadata: None })
136    }
137
138    fn system_prompt(&self) -> String {
139        // Get system prompt from config if available
140        self.agents
141            .get("orchestrator")
142            .and_then(|a| a.system_prompt.clone())
143            .unwrap_or_else(|| {
144                "You are an orchestrator agent that coordinates multiple specialized agents to answer complex queries.".to_string()
145            })
146    }
147
148    fn agent_type(&self) -> AgentType {
149        AgentType::Orchestrator
150    }
151}
152
153#[cfg(test)]
154mod tests {
155    use super::*;
156    use crate::AgentConfig;
157    use ares_llm::{LLMClient, LLMResponse, ModelConfig, ProviderConfig, ProviderRegistry};
158    use ares_tools::{Tool, Tools};
159    use ares_types::types::ToolDefinition;
160    use async_trait::async_trait;
161    use std::collections::{HashMap, VecDeque};
162    use std::sync::{Arc, Mutex};
163
164    #[derive(Clone)]
165    struct ScriptedLlm {
166        responses: Arc<Mutex<VecDeque<String>>>,
167        system_prompts: Arc<Mutex<Vec<String>>>,
168        generate_prompts: Arc<Mutex<Vec<String>>>,
169    }
170
171    impl ScriptedLlm {
172        fn new(responses: Vec<&str>) -> Self {
173            Self {
174                responses: Arc::new(Mutex::new(responses.into_iter().map(str::to_string).collect())),
175                system_prompts: Arc::new(Mutex::new(Vec::new())),
176                generate_prompts: Arc::new(Mutex::new(Vec::new())),
177            }
178        }
179
180        fn system_prompts(&self) -> Vec<String> {
181            self.system_prompts.lock().unwrap().clone()
182        }
183
184        fn generate_prompts(&self) -> Vec<String> {
185            self.generate_prompts.lock().unwrap().clone()
186        }
187
188        fn next_response(&self) -> String {
189            self.responses
190                .lock()
191                .unwrap()
192                .pop_front()
193                .unwrap_or_else(|| "fallback-response".to_string())
194        }
195    }
196
197    #[async_trait]
198    impl LLMClient for ScriptedLlm {
199        fn model_name(&self) -> &str {
200            "scripted-test"
201        }
202        async fn generate(&self, prompt: &str) -> Result<String> {
203            self.generate_prompts.lock().unwrap().push(prompt.to_string());
204            Ok(self.next_response())
205        }
206        async fn generate_with_system(&self, system: &str, _: &str) -> Result<String> {
207            self.system_prompts.lock().unwrap().push(system.to_string());
208            Ok(self.next_response())
209        }
210        async fn generate_with_history(&self, _: &[(String, String)]) -> Result<LLMResponse> {
211            Ok(LLMResponse {
212                content: self.next_response(),
213                tool_calls: vec![],
214                finish_reason: "stop".to_string(),
215                usage: None,
216            })
217        }
218        async fn generate_with_tools(&self, _: &str, _: &[ToolDefinition]) -> Result<LLMResponse> {
219            Ok(LLMResponse {
220                content: self.next_response(),
221                tool_calls: vec![],
222                finish_reason: "stop".to_string(),
223                usage: None,
224            })
225        }
226        async fn generate_with_tools_and_history(
227            &self,
228            _: &[ares_llm::coordinator::ConversationMessage],
229            _: &[ToolDefinition],
230        ) -> Result<LLMResponse> {
231            Ok(LLMResponse {
232                content: self.next_response(),
233                tool_calls: vec![],
234                finish_reason: "stop".to_string(),
235                usage: None,
236            })
237        }
238        async fn stream(&self, _: &str) -> Result<Box<dyn futures::Stream<Item = Result<String>> + Send + Unpin>> {
239            Ok(Box::new(futures::stream::empty()))
240        }
241        async fn stream_with_system(&self, _: &str, _: &str) -> Result<Box<dyn futures::Stream<Item = Result<String>> + Send + Unpin>> {
242            Ok(Box::new(futures::stream::empty()))
243        }
244        async fn stream_with_history(&self, _: &[(String, String)]) -> Result<Box<dyn futures::Stream<Item = Result<String>> + Send + Unpin>> {
245            Ok(Box::new(futures::stream::empty()))
246        }
247    }
248
249    fn test_context() -> AgentContext {
250        AgentContext {
251            user_id: "test-user".to_string(),
252            session_id: "test-session".to_string(),
253            conversation_history: vec![],
254            user_memory: None,
255        }
256    }
257
258    fn sample_agent_config() -> AgentConfig {
259        AgentConfig {
260            model: "default".to_string(),
261            system_prompt: None,
262            tools: vec![],
263            max_tool_iterations: 10,
264            parallel_tools: false,
265            allowed_tools: None,
266            extra: HashMap::new(),
267            compaction_enabled: None,
268        }
269    }
270
271    fn create_test_agent_map(agents: HashMap<String, AgentConfig>) -> HashMap<String, AgentConfig> {
272        agents
273    }
274
275    fn create_test_provider_registry() -> Arc<ProviderRegistry> {
276        let mut registry = ProviderRegistry::new();
277        registry.register_provider(
278            "ollama-local",
279            ProviderConfig::Ollama {
280                api_key_env: "TEST_KEY".to_string(),
281                base_url: "https://test.example.com".to_string(),
282                default_model: "ministral-3:3b".to_string(),
283            },
284        );
285        registry.register_model(
286            "default",
287            ModelConfig {
288                provider: "ollama-local".to_string(),
289                model: "ministral-3:3b".to_string(),
290                temperature: 0.7,
291                max_tokens: 512,
292            },
293        );
294        Arc::new(registry)
295    }
296
297
298    fn create_test_provider_registry_with_base_url(base_url: &str) -> Arc<ProviderRegistry> {
299        let mut registry = ProviderRegistry::new();
300        registry.register_provider(
301            "ollama-local",
302            ProviderConfig::Ollama {
303                api_key_env: "TEST_KEY".to_string(),
304                base_url: base_url.to_string(),
305                default_model: "ministral-3:3b".to_string(),
306            },
307        );
308        registry.register_model(
309            "default",
310            ModelConfig {
311                provider: "ollama-local".to_string(),
312                model: "ministral-3:3b".to_string(),
313                temperature: 0.7,
314                max_tokens: 512,
315            },
316        );
317        Arc::new(registry)
318    }
319
320    fn build_registry_with_provider(agent_names: &[&str], base_url: &str) -> Arc<AgentRegistry> {
321        let provider_registry = create_test_provider_registry_with_base_url(base_url);
322        let tools = Arc::new(Tools::from_static(Vec::<Arc<dyn Tool>>::new()));
323        let mut registry = AgentRegistry::new(provider_registry, tools);
324        for name in agent_names {
325            registry.register(name, sample_agent_config());
326        }
327        registry.register("orchestrator", sample_agent_config());
328        registry.register("router", sample_agent_config());
329        Arc::new(registry)
330    }
331
332    fn chat_done_json(content: &str) -> String {
333        serde_json::json!({
334            "model": "test-model",
335            "created_at": "2024-01-01T00:00:00Z",
336            "message": { "role": "assistant", "content": content },
337            "done": true
338        })
339        .to_string()
340    }
341
342    fn build_registry(agent_names: &[&str]) -> Arc<AgentRegistry> {
343        let provider_registry = create_test_provider_registry();
344        let tools = Arc::new(Tools::from_static(Vec::<Arc<dyn Tool>>::new()));
345        let mut registry = AgentRegistry::new(provider_registry, tools);
346        for name in agent_names {
347            registry.register(name, sample_agent_config());
348        }
349        registry.register("orchestrator", sample_agent_config());
350        registry.register("router", sample_agent_config());
351        Arc::new(registry)
352    }
353
354    fn build_orchestrator(llm: ScriptedLlm, agents: &[&str], config: HashMap<String, AgentConfig>) -> OrchestratorAgent {
355        OrchestratorAgent::new(
356            Box::new(llm),
357            Arc::new(config),
358            build_registry(agents),
359        )
360    }
361
362    #[tokio::test]
363    async fn test_decompose_task_parses_valid_json() {
364        let orch = build_orchestrator(
365            ScriptedLlm::new(vec![r#"[{"agent":"sales","task":"Get Q1 revenue"},{"agent":"product","task":"Top SKUs"}]"#]),
366            &["sales", "product"],
367            create_test_agent_map(HashMap::new()),
368        );
369        let tasks = orch.decompose_task("Quarterly business review").await.expect("decompose");
370        assert_eq!(tasks.len(), 2);
371        assert_eq!(tasks[0], ("sales".to_string(), "Get Q1 revenue".to_string()));
372        assert_eq!(tasks[1], ("product".to_string(), "Top SKUs".to_string()));
373    }
374
375    #[tokio::test]
376    async fn test_decompose_task_unknown_agent_falls_back_to_product() {
377        let orch = build_orchestrator(
378            ScriptedLlm::new(vec![r#"[{"agent":"unknown-agent","task":"Do something"}]"#]),
379            &["product"],
380            create_test_agent_map(HashMap::new()),
381        );
382        let tasks = orch.decompose_task("task").await.expect("decompose");
383        assert_eq!(tasks, vec![("product".to_string(), "Do something".to_string())]);
384    }
385
386    #[tokio::test]
387    async fn test_decompose_task_invalid_json_errors() {
388        let orch = build_orchestrator(
389            ScriptedLlm::new(vec!["not-json"]),
390            &["product"],
391            create_test_agent_map(HashMap::new()),
392        );
393        let err = orch.decompose_task("task").await.unwrap_err();
394        assert!(matches!(err, AppError::LLM(_)));
395    }
396
397    #[tokio::test]
398    async fn test_decompose_task_excludes_orchestrator_and_router_from_prompt() {
399        let llm = ScriptedLlm::new(vec![r#"[]"#]);
400        let llm_clone = llm.clone();
401        let orch = build_orchestrator(llm, &["sales", "product"], create_test_agent_map(HashMap::new()));
402        orch.decompose_task("plan").await.expect("decompose");
403        let system = llm_clone.system_prompts().into_iter().next().expect("system prompt");
404        assert!(system.contains("sales"));
405        assert!(system.contains("product"));
406        assert!(!system.contains("orchestrator"));
407        assert!(!system.contains("router"));
408    }
409
410    #[tokio::test]
411    async fn test_execute_with_no_subtasks_uses_direct_generation() {
412        let orch = build_orchestrator(
413            ScriptedLlm::new(vec![r#"[]"#, "direct-answer"]),
414            &["product"],
415            create_test_agent_map(HashMap::new()),
416        );
417        let resp = orch.execute("simple question", &test_context()).await.expect("execute");
418        assert_eq!(resp.content, "direct-answer");
419    }
420
421    #[tokio::test]
422    async fn test_execute_subtasks_run_concurrently() {
423        use std::time::{Duration, Instant};
424
425        async fn sleep_ok(label: &'static str) -> std::result::Result<&'static str, &'static str> {
426            tokio::time::sleep(Duration::from_millis(80)).await;
427            Ok(label)
428        }
429
430        let start = Instant::now();
431        let results = join_subtask_results(vec![sleep_ok("a"), sleep_ok("b")])
432            .await
433            .expect("join");
434        let elapsed = start.elapsed();
435        assert_eq!(results, vec!["a", "b"]);
436        assert!(
437            elapsed < Duration::from_millis(140),
438            "expected concurrent join under 140ms, got {elapsed:?}"
439        );
440    }
441
442    #[test]
443    fn test_system_prompt_from_config() {
444        let mut agents = HashMap::new();
445        agents.insert(
446            "orchestrator".to_string(),
447            AgentConfig {
448                model: "default".to_string(),
449                system_prompt: Some("Custom orchestrator prompt".to_string()),
450                tools: vec![],
451                max_tool_iterations: 10,
452                parallel_tools: false,
453            allowed_tools: None,
454            extra: HashMap::new(),
455            compaction_enabled: None,
456            },
457        );
458        let orch = build_orchestrator(ScriptedLlm::new(vec![]), &[], create_test_agent_map(agents));
459        assert_eq!(orch.system_prompt(), "Custom orchestrator prompt");
460    }
461
462    #[test]
463    fn test_system_prompt_default_when_missing() {
464        let orch = build_orchestrator(ScriptedLlm::new(vec![]), &[], create_test_agent_map(HashMap::new()));
465        assert!(orch.system_prompt().contains("orchestrator agent"));
466    }
467
468    #[test]
469    fn test_agent_type_is_orchestrator() {
470        let orch = build_orchestrator(ScriptedLlm::new(vec![]), &[], create_test_agent_map(HashMap::new()));
471        assert_eq!(orch.agent_type(), AgentType::Orchestrator);
472    }
473
474    #[tokio::test]
475    async fn test_execute_with_subtasks_delegates_and_synthesizes() {
476        use wiremock::matchers::{method, path};
477        use wiremock::{Mock, MockServer, ResponseTemplate};
478
479        let server = MockServer::start().await;
480        Mock::given(method("POST"))
481            .and(path("/api/chat"))
482            .respond_with(
483                ResponseTemplate::new(200).set_body_string(chat_done_json("sales-agent-output")),
484            )
485            .mount(&server)
486            .await;
487
488        let llm = ScriptedLlm::new(vec![
489            r#"[{"agent":"sales","task":"Get Q1 revenue"}]"#,
490            "synthesized final answer",
491        ]);
492        let registry = build_registry_with_provider(&["sales"], &server.uri());
493        let orch = OrchestratorAgent::new(
494            Box::new(llm),
495            Arc::new(create_test_agent_map(HashMap::new())),
496            registry,
497        );
498
499        let resp = orch
500            .execute("quarterly business review", &test_context())
501            .await
502            .expect("execute with subtasks");
503        assert_eq!(resp.content, "synthesized final answer");
504    }
505
506    #[tokio::test]
507    async fn test_decompose_task_defaults_missing_json_fields() {
508        let orch = build_orchestrator(
509            ScriptedLlm::new(vec![r#"[{"task":"task-only"},{"agent":"sales"}]"#]),
510            &["sales", "product"],
511            create_test_agent_map(HashMap::new()),
512        );
513        let tasks = orch.decompose_task("plan").await.expect("decompose");
514        assert_eq!(
515            tasks,
516            vec![
517                ("product".to_string(), "task-only".to_string()),
518                ("sales".to_string(), String::new()),
519            ]
520        );
521    }
522
523    #[tokio::test]
524    async fn test_decompose_task_mixed_known_and_unknown_agents() {
525        let orch = build_orchestrator(
526            ScriptedLlm::new(vec![
527                r#"[{"agent":"sales","task":"Revenue"},{"agent":"ghost","task":"Haunt"},{"agent":"finance","task":"Budget"}]"#,
528            ]),
529            &["sales", "finance", "product"],
530            create_test_agent_map(HashMap::new()),
531        );
532        let tasks = orch.decompose_task("mixed").await.expect("decompose");
533        assert_eq!(
534            tasks,
535            vec![
536                ("sales".to_string(), "Revenue".to_string()),
537                ("product".to_string(), "Haunt".to_string()),
538                ("finance".to_string(), "Budget".to_string()),
539            ]
540        );
541    }
542
543    #[tokio::test]
544    async fn test_decompose_task_only_orchestrator_router_yields_empty_agent_list() {
545        let llm = ScriptedLlm::new(vec![r#"[]"#]);
546        let llm_clone = llm.clone();
547        let orch = build_orchestrator(llm, &[], create_test_agent_map(HashMap::new()));
548        orch.decompose_task("plan").await.expect("decompose");
549        let system = llm_clone.system_prompts().into_iter().next().expect("system prompt");
550        assert!(system.contains("Available agents: "));
551        assert!(!system.contains("orchestrator"));
552        assert!(!system.contains("router"));
553    }
554
555    #[tokio::test]
556    async fn test_execute_subtask_returns_registered_agent_content() {
557        use wiremock::matchers::{method, path};
558        use wiremock::{Mock, MockServer, ResponseTemplate};
559
560        let server = MockServer::start().await;
561        Mock::given(method("POST"))
562            .and(path("/api/chat"))
563            .respond_with(ResponseTemplate::new(200).set_body_string(chat_done_json("subtask-body")))
564            .mount(&server)
565            .await;
566
567        let orch = OrchestratorAgent::new(
568            Box::new(ScriptedLlm::new(vec![])),
569            Arc::new(create_test_agent_map(HashMap::new())),
570            build_registry_with_provider(&["sales"], &server.uri()),
571        );
572        let content = orch
573            .execute_subtask("sales", "Get Q1 revenue", &test_context())
574            .await
575            .expect("subtask");
576        assert_eq!(content, "subtask-body");
577    }
578
579    #[tokio::test]
580    async fn test_execute_subtask_unknown_agent_errors() {
581        let orch = build_orchestrator(
582            ScriptedLlm::new(vec![]),
583            &["product"],
584            create_test_agent_map(HashMap::new()),
585        );
586        let err = orch
587            .execute_subtask("missing-agent", "task", &test_context())
588            .await
589            .unwrap_err();
590        assert!(!matches!(err, AppError::LLM(_)));
591    }
592
593    #[tokio::test]
594    async fn test_execute_multiple_subtasks_formats_results_for_synthesis() {
595        use wiremock::matchers::{method, path};
596        use wiremock::{Mock, MockServer, ResponseTemplate};
597
598        let server = MockServer::start().await;
599        Mock::given(method("POST"))
600            .and(path("/api/chat"))
601            .respond_with(ResponseTemplate::new(200).set_body_string(chat_done_json("agent-output")))
602            .mount(&server)
603            .await;
604
605        let llm = ScriptedLlm::new(vec![
606            r#"[{"agent":"sales","task":"Revenue"},{"agent":"finance","task":"Budget"}]"#,
607            "combined answer",
608        ]);
609        let llm_clone = llm.clone();
610        let orch = OrchestratorAgent::new(
611            Box::new(llm),
612            Arc::new(create_test_agent_map(HashMap::new())),
613            build_registry_with_provider(&["sales", "finance"], &server.uri()),
614        );
615
616        let resp = orch
617            .execute("annual review", &test_context())
618            .await
619            .expect("execute");
620        assert_eq!(resp.content, "combined answer");
621        let synthesis = llm_clone.generate_prompts().into_iter().next().expect("synthesis prompt");
622        assert!(synthesis.contains("Original query: annual review"));
623        assert!(synthesis.contains("[sales] agent-output"));
624        assert!(synthesis.contains("[finance] agent-output"));
625        assert!(synthesis.contains("Subtask results:"));
626    }
627
628    #[test]
629    fn test_agent_type_orchestrator_serde_roundtrip() {
630        let agent = AgentType::Orchestrator;
631        let json = serde_json::to_string(&agent).expect("serialize");
632        assert_eq!(json, "\"orchestrator\"");
633        let parsed: AgentType = serde_json::from_str(&json).expect("deserialize");
634        assert_eq!(parsed, agent);
635        assert_eq!(format!("{parsed:?}"), "Orchestrator");
636        assert_eq!(agent.clone(), AgentType::Orchestrator);
637    }
638
639    #[test]
640    fn test_agent_context_debug_clone() {
641        let ctx = test_context();
642        let cloned = ctx.clone();
643        assert_eq!(cloned.user_id, "test-user");
644        assert_eq!(cloned.session_id, "test-session");
645        assert!(format!("{ctx:?}").contains("AgentContext"));
646    }
647
648    #[test]
649    fn test_agent_response_none_usage_and_metadata() {
650        let resp = AgentResponse {
651            content: "ok".to_string(),
652            usage: None,
653            metadata: None,
654        };
655        assert_eq!(resp.content, "ok");
656        assert!(resp.usage.is_none());
657        assert!(resp.metadata.is_none());
658    }
659
660
661}
ares_agent/orchestrator.rs

ares_agent/
orchestrator.rs
