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                reasoning_content: None,
217                response_id: None,
218            })
219        }
220        async fn generate_with_tools(&self, _: &str, _: &[ToolDefinition]) -> Result<LLMResponse> {
221            Ok(LLMResponse {
222                content: self.next_response(),
223                tool_calls: vec![],
224                finish_reason: "stop".to_string(),
225                usage: None,
226                reasoning_content: None,
227                response_id: None,
228            })
229        }
230        async fn generate_with_tools_and_history(
231            &self,
232            _: &[ares_llm::coordinator::ConversationMessage],
233            _: &[ToolDefinition],
234        ) -> Result<LLMResponse> {
235            Ok(LLMResponse {
236                content: self.next_response(),
237                tool_calls: vec![],
238                finish_reason: "stop".to_string(),
239                usage: None,
240                reasoning_content: None,
241                response_id: None,
242            })
243        }
244        async fn stream(&self, _: &str) -> Result<Box<dyn futures::Stream<Item = Result<String>> + Send + Unpin>> {
245            Ok(Box::new(futures::stream::empty()))
246        }
247        async fn stream_with_system(&self, _: &str, _: &str) -> Result<Box<dyn futures::Stream<Item = Result<String>> + Send + Unpin>> {
248            Ok(Box::new(futures::stream::empty()))
249        }
250        async fn stream_with_history(&self, _: &[(String, String)]) -> Result<Box<dyn futures::Stream<Item = Result<String>> + Send + Unpin>> {
251            Ok(Box::new(futures::stream::empty()))
252        }
253    }
254
255    fn test_context() -> AgentContext {
256        AgentContext {
257            user_id: "test-user".to_string(),
258            session_id: "test-session".to_string(),
259            conversation_history: vec![],
260            user_memory: None,
261        }
262    }
263
264    fn sample_agent_config() -> AgentConfig {
265        AgentConfig {
266            model: "default".to_string(),
267            system_prompt: None,
268            tools: vec![],
269            max_tool_iterations: 10,
270            parallel_tools: false,
271            allowed_tools: None,
272            extra: HashMap::new(),
273            compaction_enabled: None,
274        }
275    }
276
277    fn create_test_agent_map(agents: HashMap<String, AgentConfig>) -> HashMap<String, AgentConfig> {
278        agents
279    }
280
281    fn create_test_provider_registry() -> Arc<ProviderRegistry> {
282        let mut registry = ProviderRegistry::new();
283        registry.register_provider(
284            "ollama-local",
285            ProviderConfig::Ollama {
286                api_key_env: "TEST_KEY".to_string(),
287                base_url: "https://test.example.com".to_string(),
288                default_model: "ministral-3:3b".to_string(),
289            },
290        );
291        registry.register_model(
292            "default",
293            ModelConfig {
294                provider: "ollama-local".to_string(),
295                model: "ministral-3:3b".to_string(),
296                temperature: 0.7,
297                max_tokens: 512,
298            },
299        );
300        Arc::new(registry)
301    }
302
303
304    fn create_test_provider_registry_with_base_url(base_url: &str) -> Arc<ProviderRegistry> {
305        let mut registry = ProviderRegistry::new();
306        registry.register_provider(
307            "ollama-local",
308            ProviderConfig::Ollama {
309                api_key_env: "TEST_KEY".to_string(),
310                base_url: base_url.to_string(),
311                default_model: "ministral-3:3b".to_string(),
312            },
313        );
314        registry.register_model(
315            "default",
316            ModelConfig {
317                provider: "ollama-local".to_string(),
318                model: "ministral-3:3b".to_string(),
319                temperature: 0.7,
320                max_tokens: 512,
321            },
322        );
323        Arc::new(registry)
324    }
325
326    fn build_registry_with_provider(agent_names: &[&str], base_url: &str) -> Arc<AgentRegistry> {
327        let provider_registry = create_test_provider_registry_with_base_url(base_url);
328        let tools = Arc::new(Tools::from_static(Vec::<Arc<dyn Tool>>::new()));
329        let mut registry = AgentRegistry::new(provider_registry, tools);
330        for name in agent_names {
331            registry.register(name, sample_agent_config());
332        }
333        registry.register("orchestrator", sample_agent_config());
334        registry.register("router", sample_agent_config());
335        Arc::new(registry)
336    }
337
338    fn chat_done_json(content: &str) -> serde_json::Value {
339        serde_json::json!({
340            "model": "test-model",
341            "created_at": "2024-01-01T00:00:00Z",
342            "message": { "role": "assistant", "content": content },
343            "done": true
344        })
345    }
346
347    fn build_registry(agent_names: &[&str]) -> Arc<AgentRegistry> {
348        let provider_registry = create_test_provider_registry();
349        let tools = Arc::new(Tools::from_static(Vec::<Arc<dyn Tool>>::new()));
350        let mut registry = AgentRegistry::new(provider_registry, tools);
351        for name in agent_names {
352            registry.register(name, sample_agent_config());
353        }
354        registry.register("orchestrator", sample_agent_config());
355        registry.register("router", sample_agent_config());
356        Arc::new(registry)
357    }
358
359    fn build_orchestrator(llm: ScriptedLlm, agents: &[&str], config: HashMap<String, AgentConfig>) -> OrchestratorAgent {
360        OrchestratorAgent::new(
361            Box::new(llm),
362            Arc::new(config),
363            build_registry(agents),
364        )
365    }
366
367    #[tokio::test]
368    async fn test_decompose_task_parses_valid_json() {
369        let orch = build_orchestrator(
370            ScriptedLlm::new(vec![r#"[{"agent":"sales","task":"Get Q1 revenue"},{"agent":"product","task":"Top SKUs"}]"#]),
371            &["sales", "product"],
372            create_test_agent_map(HashMap::new()),
373        );
374        let tasks = orch.decompose_task("Quarterly business review").await.expect("decompose");
375        assert_eq!(tasks.len(), 2);
376        assert_eq!(tasks[0], ("sales".to_string(), "Get Q1 revenue".to_string()));
377        assert_eq!(tasks[1], ("product".to_string(), "Top SKUs".to_string()));
378    }
379
380    #[tokio::test]
381    async fn test_decompose_task_unknown_agent_falls_back_to_product() {
382        let orch = build_orchestrator(
383            ScriptedLlm::new(vec![r#"[{"agent":"unknown-agent","task":"Do something"}]"#]),
384            &["product"],
385            create_test_agent_map(HashMap::new()),
386        );
387        let tasks = orch.decompose_task("task").await.expect("decompose");
388        assert_eq!(tasks, vec![("product".to_string(), "Do something".to_string())]);
389    }
390
391    #[tokio::test]
392    async fn test_decompose_task_invalid_json_errors() {
393        let orch = build_orchestrator(
394            ScriptedLlm::new(vec!["not-json"]),
395            &["product"],
396            create_test_agent_map(HashMap::new()),
397        );
398        let err = orch.decompose_task("task").await.unwrap_err();
399        assert!(matches!(err, AppError::LLM(_)));
400    }
401
402    #[tokio::test]
403    async fn test_decompose_task_excludes_orchestrator_and_router_from_prompt() {
404        let llm = ScriptedLlm::new(vec![r#"[]"#]);
405        let llm_clone = llm.clone();
406        let orch = build_orchestrator(llm, &["sales", "product"], create_test_agent_map(HashMap::new()));
407        orch.decompose_task("plan").await.expect("decompose");
408        let system = llm_clone.system_prompts().into_iter().next().expect("system prompt");
409        assert!(system.contains("sales"));
410        assert!(system.contains("product"));
411        assert!(!system.contains("orchestrator"));
412        assert!(!system.contains("router"));
413    }
414
415    #[tokio::test]
416    async fn test_execute_with_no_subtasks_uses_direct_generation() {
417        let orch = build_orchestrator(
418            ScriptedLlm::new(vec![r#"[]"#, "direct-answer"]),
419            &["product"],
420            create_test_agent_map(HashMap::new()),
421        );
422        let resp = orch.execute("simple question", &test_context()).await.expect("execute");
423        assert_eq!(resp.content, "direct-answer");
424    }
425
426    #[tokio::test]
427    async fn test_execute_subtasks_run_concurrently() {
428        use std::time::{Duration, Instant};
429
430        async fn sleep_ok(label: &'static str) -> std::result::Result<&'static str, &'static str> {
431            tokio::time::sleep(Duration::from_millis(80)).await;
432            Ok(label)
433        }
434
435        let start = Instant::now();
436        let results = join_subtask_results(vec![sleep_ok("a"), sleep_ok("b")])
437            .await
438            .expect("join");
439        let elapsed = start.elapsed();
440        assert_eq!(results, vec!["a", "b"]);
441        assert!(
442            elapsed < Duration::from_millis(140),
443            "expected concurrent join under 140ms, got {elapsed:?}"
444        );
445    }
446
447    #[test]
448    fn test_system_prompt_from_config() {
449        let mut agents = HashMap::new();
450        agents.insert(
451            "orchestrator".to_string(),
452            AgentConfig {
453                model: "default".to_string(),
454                system_prompt: Some("Custom orchestrator prompt".to_string()),
455                tools: vec![],
456                max_tool_iterations: 10,
457                parallel_tools: false,
458            allowed_tools: None,
459            extra: HashMap::new(),
460            compaction_enabled: None,
461            },
462        );
463        let orch = build_orchestrator(ScriptedLlm::new(vec![]), &[], create_test_agent_map(agents));
464        assert_eq!(orch.system_prompt(), "Custom orchestrator prompt");
465    }
466
467    #[test]
468    fn test_system_prompt_default_when_missing() {
469        let orch = build_orchestrator(ScriptedLlm::new(vec![]), &[], create_test_agent_map(HashMap::new()));
470        assert!(orch.system_prompt().contains("orchestrator agent"));
471    }
472
473    #[test]
474    fn test_agent_type_is_orchestrator() {
475        let orch = build_orchestrator(ScriptedLlm::new(vec![]), &[], create_test_agent_map(HashMap::new()));
476        assert_eq!(orch.agent_type(), AgentType::Orchestrator);
477    }
478
479    #[tokio::test]
480    async fn test_execute_with_subtasks_delegates_and_synthesizes() {
481        use wiremock::matchers::{method, path};
482        use wiremock::{Mock, MockServer, ResponseTemplate};
483
484        let server = MockServer::start().await;
485        Mock::given(method("POST"))
486            .and(path("/api/chat"))
487            .respond_with(
488                ResponseTemplate::new(200).set_body_json(chat_done_json("sales-agent-output")),
489            )
490            .mount(&server)
491            .await;
492
493        let llm = ScriptedLlm::new(vec![
494            r#"[{"agent":"sales","task":"Get Q1 revenue"}]"#,
495            "synthesized final answer",
496        ]);
497        let registry = build_registry_with_provider(&["sales"], &server.uri());
498        let orch = OrchestratorAgent::new(
499            Box::new(llm),
500            Arc::new(create_test_agent_map(HashMap::new())),
501            registry,
502        );
503
504        let resp = orch
505            .execute("quarterly business review", &test_context())
506            .await
507            .expect("execute with subtasks");
508        assert_eq!(resp.content, "synthesized final answer");
509    }
510
511    #[tokio::test]
512    async fn test_decompose_task_defaults_missing_json_fields() {
513        let orch = build_orchestrator(
514            ScriptedLlm::new(vec![r#"[{"task":"task-only"},{"agent":"sales"}]"#]),
515            &["sales", "product"],
516            create_test_agent_map(HashMap::new()),
517        );
518        let tasks = orch.decompose_task("plan").await.expect("decompose");
519        assert_eq!(
520            tasks,
521            vec![
522                ("product".to_string(), "task-only".to_string()),
523                ("sales".to_string(), String::new()),
524            ]
525        );
526    }
527
528    #[tokio::test]
529    async fn test_decompose_task_mixed_known_and_unknown_agents() {
530        let orch = build_orchestrator(
531            ScriptedLlm::new(vec![
532                r#"[{"agent":"sales","task":"Revenue"},{"agent":"ghost","task":"Haunt"},{"agent":"finance","task":"Budget"}]"#,
533            ]),
534            &["sales", "finance", "product"],
535            create_test_agent_map(HashMap::new()),
536        );
537        let tasks = orch.decompose_task("mixed").await.expect("decompose");
538        assert_eq!(
539            tasks,
540            vec![
541                ("sales".to_string(), "Revenue".to_string()),
542                ("product".to_string(), "Haunt".to_string()),
543                ("finance".to_string(), "Budget".to_string()),
544            ]
545        );
546    }
547
548    #[tokio::test]
549    async fn test_decompose_task_only_orchestrator_router_yields_empty_agent_list() {
550        let llm = ScriptedLlm::new(vec![r#"[]"#]);
551        let llm_clone = llm.clone();
552        let orch = build_orchestrator(llm, &[], create_test_agent_map(HashMap::new()));
553        orch.decompose_task("plan").await.expect("decompose");
554        let system = llm_clone.system_prompts().into_iter().next().expect("system prompt");
555        assert!(system.contains("Available agents: "));
556        assert!(!system.contains("orchestrator"));
557        assert!(!system.contains("router"));
558    }
559
560    #[tokio::test]
561    async fn test_execute_subtask_returns_registered_agent_content() {
562        use wiremock::matchers::{method, path};
563        use wiremock::{Mock, MockServer, ResponseTemplate};
564
565        let server = MockServer::start().await;
566        Mock::given(method("POST"))
567            .and(path("/api/chat"))
568            .respond_with(
569                ResponseTemplate::new(200).set_body_json(chat_done_json("subtask-body")),
570            )
571            .mount(&server)
572            .await;
573
574        let orch = OrchestratorAgent::new(
575            Box::new(ScriptedLlm::new(vec![])),
576            Arc::new(create_test_agent_map(HashMap::new())),
577            build_registry_with_provider(&["sales"], &server.uri()),
578        );
579        let content = orch
580            .execute_subtask("sales", "Get Q1 revenue", &test_context())
581            .await
582            .expect("subtask");
583        assert_eq!(content, "subtask-body");
584    }
585
586    #[tokio::test]
587    async fn test_execute_subtask_unknown_agent_errors() {
588        let orch = build_orchestrator(
589            ScriptedLlm::new(vec![]),
590            &["product"],
591            create_test_agent_map(HashMap::new()),
592        );
593        let err = orch
594            .execute_subtask("missing-agent", "task", &test_context())
595            .await
596            .unwrap_err();
597        assert!(!matches!(err, AppError::LLM(_)));
598    }
599
600    #[tokio::test]
601    async fn test_execute_multiple_subtasks_formats_results_for_synthesis() {
602        use wiremock::matchers::{method, path};
603        use wiremock::{Mock, MockServer, ResponseTemplate};
604
605        let server = MockServer::start().await;
606        Mock::given(method("POST"))
607            .and(path("/api/chat"))
608            .respond_with(
609                ResponseTemplate::new(200).set_body_json(chat_done_json("agent-output")),
610            )
611            .mount(&server)
612            .await;
613
614        let llm = ScriptedLlm::new(vec![
615            r#"[{"agent":"sales","task":"Revenue"},{"agent":"finance","task":"Budget"}]"#,
616            "combined answer",
617        ]);
618        let llm_clone = llm.clone();
619        let orch = OrchestratorAgent::new(
620            Box::new(llm),
621            Arc::new(create_test_agent_map(HashMap::new())),
622            build_registry_with_provider(&["sales", "finance"], &server.uri()),
623        );
624
625        let resp = orch
626            .execute("annual review", &test_context())
627            .await
628            .expect("execute");
629        assert_eq!(resp.content, "combined answer");
630        let synthesis = llm_clone.generate_prompts().into_iter().next().expect("synthesis prompt");
631        assert!(synthesis.contains("Original query: annual review"));
632        assert!(synthesis.contains("[sales] agent-output"));
633        assert!(synthesis.contains("[finance] agent-output"));
634        assert!(synthesis.contains("Subtask results:"));
635    }
636
637    #[test]
638    fn test_agent_type_orchestrator_serde_roundtrip() {
639        let agent = AgentType::Orchestrator;
640        let json = serde_json::to_string(&agent).expect("serialize");
641        assert_eq!(json, "\"orchestrator\"");
642        let parsed: AgentType = serde_json::from_str(&json).expect("deserialize");
643        assert_eq!(parsed, agent);
644        assert_eq!(format!("{parsed:?}"), "Orchestrator");
645        assert_eq!(agent.clone(), AgentType::Orchestrator);
646    }
647
648    #[test]
649    fn test_agent_context_debug_clone() {
650        let ctx = test_context();
651        let cloned = ctx.clone();
652        assert_eq!(cloned.user_id, "test-user");
653        assert_eq!(cloned.session_id, "test-session");
654        assert!(format!("{ctx:?}").contains("AgentContext"));
655    }
656
657    #[test]
658    fn test_agent_response_none_usage_and_metadata() {
659        let resp = AgentResponse {
660            content: "ok".to_string(),
661            usage: None,
662            metadata: None,
663        };
664        assert_eq!(resp.content, "ok");
665        assert!(resp.usage.is_none());
666        assert!(resp.metadata.is_none());
667    }
668
669
670}
ares_agent/orchestrator.rs

ares_agent/
orchestrator.rs
