1use crate::{AdkError, Result};
2use async_trait::async_trait;
3use std::collections::HashMap;
4use std::sync::Arc;
5
6#[async_trait]
8pub trait AgentLoader: Send + Sync {
9 async fn load_agent(&self, name: &str) -> Result<Arc<dyn crate::Agent>>;
11
12 fn list_agents(&self) -> Vec<String>;
14
15 fn root_agent(&self) -> Arc<dyn crate::Agent>;
17}
18
19pub struct SingleAgentLoader {
21 agent: Arc<dyn crate::Agent>,
22}
23
24impl SingleAgentLoader {
25 pub fn new(agent: Arc<dyn crate::Agent>) -> Self {
26 Self { agent }
27 }
28}
29
30#[async_trait]
31impl AgentLoader for SingleAgentLoader {
32 async fn load_agent(&self, name: &str) -> Result<Arc<dyn crate::Agent>> {
33 if name.is_empty() || name == self.agent.name() {
34 Ok(self.agent.clone())
35 } else {
36 Err(AdkError::Config(format!(
37 "Cannot load agent '{}' - use empty string or '{}'",
38 name,
39 self.agent.name()
40 )))
41 }
42 }
43
44 fn list_agents(&self) -> Vec<String> {
45 vec![self.agent.name().to_string()]
46 }
47
48 fn root_agent(&self) -> Arc<dyn crate::Agent> {
49 self.agent.clone()
50 }
51}
52
53pub struct MultiAgentLoader {
55 agent_map: HashMap<String, Arc<dyn crate::Agent>>,
56 root: Arc<dyn crate::Agent>,
57}
58
59impl MultiAgentLoader {
60 pub fn new(agents: Vec<Arc<dyn crate::Agent>>) -> Result<Self> {
64 if agents.is_empty() {
65 return Err(AdkError::Config(
66 "MultiAgentLoader requires at least one agent".to_string(),
67 ));
68 }
69
70 let mut agent_map = HashMap::new();
71 let root = agents[0].clone();
72
73 for agent in agents {
74 let name = agent.name().to_string();
75 if agent_map.contains_key(&name) {
76 return Err(AdkError::Config(format!("Duplicate agent name: {}", name)));
77 }
78 agent_map.insert(name, agent);
79 }
80
81 Ok(Self { agent_map, root })
82 }
83}
84
85#[async_trait]
86impl AgentLoader for MultiAgentLoader {
87 async fn load_agent(&self, name: &str) -> Result<Arc<dyn crate::Agent>> {
88 if name.is_empty() {
89 return Ok(self.root.clone());
90 }
91
92 self.agent_map.get(name).cloned().ok_or_else(|| {
93 AdkError::Config(format!(
94 "Agent '{}' not found. Available agents: {:?}",
95 name,
96 self.list_agents()
97 ))
98 })
99 }
100
101 fn list_agents(&self) -> Vec<String> {
102 self.agent_map.keys().cloned().collect()
103 }
104
105 fn root_agent(&self) -> Arc<dyn crate::Agent> {
106 self.root.clone()
107 }
108}