hanzo_agents/
registry.rs

1//! Agent registry for managing and accessing specialized agents
2
3use crate::agents::*;
4use crate::tools::ToolRegistry;
5use crate::traits::{AgentError, Result, SpecializedAgent};
6use std::collections::HashMap;
7use std::sync::Arc;
8
9/// Types of specialized agents available
10#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
11pub enum AgentType {
12    /// High-level system design and architecture
13    Architect,
14    /// Technical leadership and code quality
15    Cto,
16    /// Code review and quality assurance
17    Reviewer,
18    /// Codebase exploration and documentation
19    Explorer,
20    /// Task planning and implementation strategy
21    Planner,
22    /// Research and analysis
23    Scientist,
24}
25
26impl AgentType {
27    /// Get the string name of this agent type
28    pub fn name(&self) -> &'static str {
29        match self {
30            AgentType::Architect => "architect",
31            AgentType::Cto => "cto",
32            AgentType::Reviewer => "reviewer",
33            AgentType::Explorer => "explorer",
34            AgentType::Planner => "planner",
35            AgentType::Scientist => "scientist",
36        }
37    }
38
39    /// Parse agent type from string
40    pub fn from_str(s: &str) -> Option<Self> {
41        match s.to_lowercase().as_str() {
42            "architect" => Some(AgentType::Architect),
43            "cto" => Some(AgentType::Cto),
44            "reviewer" | "review" => Some(AgentType::Reviewer),
45            "explorer" | "explore" => Some(AgentType::Explorer),
46            "planner" | "plan" => Some(AgentType::Planner),
47            "scientist" | "research" => Some(AgentType::Scientist),
48            _ => None,
49        }
50    }
51
52    /// List all available agent types
53    pub fn all() -> Vec<AgentType> {
54        vec![
55            AgentType::Architect,
56            AgentType::Cto,
57            AgentType::Reviewer,
58            AgentType::Explorer,
59            AgentType::Planner,
60            AgentType::Scientist,
61        ]
62    }
63}
64
65impl std::fmt::Display for AgentType {
66    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
67        write!(f, "{}", self.name())
68    }
69}
70
71/// Registry for managing specialized agents
72///
73/// The registry provides a central point for:
74/// - Creating and accessing agents
75/// - Managing shared tool registry
76/// - Configuring agent defaults
77pub struct AgentRegistry {
78    tool_registry: Arc<ToolRegistry>,
79    agents: HashMap<AgentType, Arc<dyn SpecializedAgent>>,
80}
81
82impl AgentRegistry {
83    /// Create a new agent registry with default tool registry
84    pub fn new() -> Self {
85        let tool_registry = Arc::new(ToolRegistry::with_defaults());
86        Self {
87            tool_registry,
88            agents: HashMap::new(),
89        }
90    }
91
92    /// Create a registry with a custom tool registry
93    pub fn with_tools(tool_registry: Arc<ToolRegistry>) -> Self {
94        Self {
95            tool_registry,
96            agents: HashMap::new(),
97        }
98    }
99
100    /// Get the shared tool registry
101    pub fn tool_registry(&self) -> Arc<ToolRegistry> {
102        self.tool_registry.clone()
103    }
104
105    /// Get or create an agent of the specified type
106    pub fn get(&mut self, agent_type: AgentType) -> Result<Arc<dyn SpecializedAgent>> {
107        if let Some(agent) = self.agents.get(&agent_type) {
108            return Ok(agent.clone());
109        }
110
111        let agent: Arc<dyn SpecializedAgent> = match agent_type {
112            AgentType::Architect => Arc::new(ArchitectAgent::new(self.tool_registry.clone())),
113            AgentType::Cto => Arc::new(CtoAgent::new(self.tool_registry.clone())),
114            AgentType::Reviewer => Arc::new(ReviewerAgent::new(self.tool_registry.clone())),
115            AgentType::Explorer => Arc::new(ExplorerAgent::new(self.tool_registry.clone())),
116            AgentType::Planner => Arc::new(PlannerAgent::new(self.tool_registry.clone())),
117            AgentType::Scientist => Arc::new(ScientistAgent::new(self.tool_registry.clone())),
118        };
119
120        self.agents.insert(agent_type, agent.clone());
121        Ok(agent)
122    }
123
124    /// Get an agent by name
125    pub fn get_by_name(&mut self, name: &str) -> Result<Arc<dyn SpecializedAgent>> {
126        let agent_type = AgentType::from_str(name)
127            .ok_or_else(|| AgentError::ConfigError(format!("Unknown agent type: {}", name)))?;
128        self.get(agent_type)
129    }
130
131    /// List all available agent types with descriptions
132    pub fn list_agents(&self) -> Vec<(AgentType, &'static str)> {
133        vec![
134            (
135                AgentType::Architect,
136                "High-level system design and architecture",
137            ),
138            (AgentType::Cto, "Technical leadership and code quality"),
139            (AgentType::Reviewer, "Code review and quality assurance"),
140            (
141                AgentType::Explorer,
142                "Codebase exploration and documentation",
143            ),
144            (
145                AgentType::Planner,
146                "Task planning and implementation strategy",
147            ),
148            (AgentType::Scientist, "Research and analysis"),
149        ]
150    }
151}
152
153impl Default for AgentRegistry {
154    fn default() -> Self {
155        Self::new()
156    }
157}
158
159#[cfg(test)]
160mod tests {
161    use super::*;
162
163    #[test]
164    fn test_agent_type_from_str() {
165        assert_eq!(AgentType::from_str("architect"), Some(AgentType::Architect));
166        assert_eq!(AgentType::from_str("CTO"), Some(AgentType::Cto));
167        assert_eq!(AgentType::from_str("review"), Some(AgentType::Reviewer));
168        assert_eq!(AgentType::from_str("unknown"), None);
169    }
170
171    #[test]
172    fn test_registry_creation() {
173        let mut registry = AgentRegistry::new();
174        let agent = registry.get(AgentType::Architect).unwrap();
175        assert_eq!(agent.name(), "architect");
176    }
177
178    #[test]
179    fn test_registry_caching() {
180        let mut registry = AgentRegistry::new();
181        let agent1 = registry.get(AgentType::Cto).unwrap();
182        let agent2 = registry.get(AgentType::Cto).unwrap();
183        // Both should be the same Arc
184        assert!(Arc::ptr_eq(&agent1, &agent2));
185    }
186}
hanzo_agents/registry.rs

hanzo_agents/
registry.rs
