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 '{name}' - use empty string or '{}'",
38 self.agent.name()
39 )))
40 }
41 }
42
43 fn list_agents(&self) -> Vec<String> {
44 vec![self.agent.name().to_string()]
45 }
46
47 fn root_agent(&self) -> Arc<dyn crate::Agent> {
48 self.agent.clone()
49 }
50}
51
52pub struct MultiAgentLoader {
54 agent_map: HashMap<String, Arc<dyn crate::Agent>>,
55 root: Arc<dyn crate::Agent>,
56}
57
58impl MultiAgentLoader {
59 pub fn new(agents: Vec<Arc<dyn crate::Agent>>) -> Result<Self> {
63 if agents.is_empty() {
64 return Err(AdkError::config("MultiAgentLoader requires at least one agent"));
65 }
66
67 let mut agent_map = HashMap::new();
68 let root = agents[0].clone();
69
70 for agent in agents {
71 let name = agent.name().to_string();
72 if agent_map.contains_key(&name) {
73 return Err(AdkError::config(format!("Duplicate agent name: {name}")));
74 }
75 agent_map.insert(name, agent);
76 }
77
78 Ok(Self { agent_map, root })
79 }
80}
81
82#[async_trait]
83impl AgentLoader for MultiAgentLoader {
84 async fn load_agent(&self, name: &str) -> Result<Arc<dyn crate::Agent>> {
85 if name.is_empty() {
86 return Ok(self.root.clone());
87 }
88
89 self.agent_map.get(name).cloned().ok_or_else(|| {
90 AdkError::config(format!(
91 "Agent '{name}' not found. Available agents: {:?}",
92 self.list_agents()
93 ))
94 })
95 }
96
97 fn list_agents(&self) -> Vec<String> {
98 self.agent_map.keys().cloned().collect()
99 }
100
101 fn root_agent(&self) -> Arc<dyn crate::Agent> {
102 self.root.clone()
103 }
104}