Skip to main content

adk_core/
agent_loader.rs

1use crate::{AdkError, Result};
2use async_trait::async_trait;
3use std::collections::HashMap;
4use std::sync::Arc;
5
6/// Trait for loading agents by name.
7#[async_trait]
8pub trait AgentLoader: Send + Sync {
9    /// Load an agent by name (or app_name for compatibility).
10    async fn load_agent(&self, name: &str) -> Result<Arc<dyn crate::Agent>>;
11
12    /// List all available agent names.
13    fn list_agents(&self) -> Vec<String>;
14
15    /// Get the root (default) agent.
16    fn root_agent(&self) -> Arc<dyn crate::Agent>;
17}
18
19/// Single agent loader that returns the same agent for all names.
20pub 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
52/// Multi-agent loader that manages multiple agents by name.
53pub struct MultiAgentLoader {
54    agent_map: HashMap<String, Arc<dyn crate::Agent>>,
55    root: Arc<dyn crate::Agent>,
56}
57
58impl MultiAgentLoader {
59    /// Create a new MultiAgentLoader with the given agents.
60    /// The first agent becomes the root agent.
61    /// Returns an error if duplicate agent names are found.
62    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}