use crate::{AdkError, Result};
use async_trait::async_trait;
use std::collections::HashMap;
use std::sync::Arc;
#[async_trait]
pub trait AgentLoader: Send + Sync {
async fn load_agent(&self, name: &str) -> Result<Arc<dyn crate::Agent>>;
fn list_agents(&self) -> Vec<String>;
fn root_agent(&self) -> Arc<dyn crate::Agent>;
}
pub struct SingleAgentLoader {
agent: Arc<dyn crate::Agent>,
}
impl SingleAgentLoader {
pub fn new(agent: Arc<dyn crate::Agent>) -> Self {
Self { agent }
}
}
#[async_trait]
impl AgentLoader for SingleAgentLoader {
async fn load_agent(&self, name: &str) -> Result<Arc<dyn crate::Agent>> {
if name.is_empty() || name == self.agent.name() {
Ok(self.agent.clone())
} else {
Err(AdkError::config(format!(
"Cannot load agent '{name}' - use empty string or '{}'",
self.agent.name()
)))
}
}
fn list_agents(&self) -> Vec<String> {
vec![self.agent.name().to_string()]
}
fn root_agent(&self) -> Arc<dyn crate::Agent> {
self.agent.clone()
}
}
pub struct MultiAgentLoader {
agent_map: HashMap<String, Arc<dyn crate::Agent>>,
root: Arc<dyn crate::Agent>,
}
impl MultiAgentLoader {
pub fn new(agents: Vec<Arc<dyn crate::Agent>>) -> Result<Self> {
if agents.is_empty() {
return Err(AdkError::config("MultiAgentLoader requires at least one agent"));
}
let mut agent_map = HashMap::new();
let root = agents[0].clone();
for agent in agents {
let name = agent.name().to_string();
if agent_map.contains_key(&name) {
return Err(AdkError::config(format!("Duplicate agent name: {name}")));
}
agent_map.insert(name, agent);
}
Ok(Self { agent_map, root })
}
}
#[async_trait]
impl AgentLoader for MultiAgentLoader {
async fn load_agent(&self, name: &str) -> Result<Arc<dyn crate::Agent>> {
if name.is_empty() {
return Ok(self.root.clone());
}
self.agent_map.get(name).cloned().ok_or_else(|| {
AdkError::config(format!(
"Agent '{name}' not found. Available agents: {:?}",
self.list_agents()
))
})
}
fn list_agents(&self) -> Vec<String> {
self.agent_map.keys().cloned().collect()
}
fn root_agent(&self) -> Arc<dyn crate::Agent> {
self.root.clone()
}
}