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 {
27 Self { agent }
28 }
29}
30
31#[async_trait]
32impl AgentLoader for SingleAgentLoader {
33 async fn load_agent(&self, name: &str) -> Result<Arc<dyn crate::Agent>> {
34 if name.is_empty() || name == self.agent.name() {
35 Ok(self.agent.clone())
36 } else {
37 Err(AdkError::config(format!(
38 "Cannot load agent '{name}' - use empty string or '{}'",
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("MultiAgentLoader requires at least one agent"));
66 }
67
68 let mut agent_map = HashMap::new();
69 let root = agents[0].clone();
70
71 for agent in agents {
72 let name = agent.name().to_string();
73 if agent_map.contains_key(&name) {
74 return Err(AdkError::config(format!("Duplicate agent name: {name}")));
75 }
76 agent_map.insert(name, agent);
77 }
78
79 Ok(Self { agent_map, root })
80 }
81}
82
83#[async_trait]
84impl AgentLoader for MultiAgentLoader {
85 async fn load_agent(&self, name: &str) -> Result<Arc<dyn crate::Agent>> {
86 if name.is_empty() {
87 return Ok(self.root.clone());
88 }
89
90 self.agent_map.get(name).cloned().ok_or_else(|| {
91 AdkError::config(format!(
92 "Agent '{name}' not found. Available agents: {:?}",
93 self.list_agents()
94 ))
95 })
96 }
97
98 fn list_agents(&self) -> Vec<String> {
99 self.agent_map.keys().cloned().collect()
100 }
101
102 fn root_agent(&self) -> Arc<dyn crate::Agent> {
103 self.root.clone()
104 }
105}