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    /// Creates a new loader that always returns the given agent.
26    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
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("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}