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 '{}' - 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
53/// Multi-agent loader that manages multiple agents by name.
54pub struct MultiAgentLoader {
55    agent_map: HashMap<String, Arc<dyn crate::Agent>>,
56    root: Arc<dyn crate::Agent>,
57}
58
59impl MultiAgentLoader {
60    /// Create a new MultiAgentLoader with the given agents.
61    /// The first agent becomes the root agent.
62    /// Returns an error if duplicate agent names are found.
63    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}