1use std::collections::HashMap;
4use serde::{Deserialize, Serialize};
5use uuid::Uuid;
6use crate::{Result, AIError, AgentConfig};
7
8#[derive(Debug, Clone, Serialize, Deserialize)]
10pub enum AgentType {
11 Researcher,
12 Coder,
13 Analyst,
14 Coordinator,
15 Specialist(String),
16}
17
18impl std::fmt::Display for AgentType {
19 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
20 match self {
21 AgentType::Researcher => write!(f, "Researcher"),
22 AgentType::Coder => write!(f, "Coder"),
23 AgentType::Analyst => write!(f, "Analyst"),
24 AgentType::Coordinator => write!(f, "Coordinator"),
25 AgentType::Specialist(s) => write!(f, "Specialist({})", s),
26 }
27 }
28}
29
30#[derive(Debug, Clone, Serialize, Deserialize)]
32pub struct Agent {
33 pub id: String,
34 pub agent_type: AgentType,
35 pub capabilities: Vec<String>,
36 pub config: HashMap<String, String>,
37}
38
39pub struct AgentManager {
41 config: AgentConfig,
42 agents: HashMap<String, Agent>,
43}
44
45impl AgentManager {
46 pub fn new(config: AgentConfig) -> Self {
47 Self {
48 config,
49 agents: HashMap::new(),
50 }
51 }
52
53 pub async fn spawn_agent(
54 &mut self,
55 agent_type: AgentType,
56 capabilities: Vec<String>,
57 custom_config: Option<HashMap<String, String>>,
58 ) -> Result<Agent> {
59 let agent = Agent {
60 id: Uuid::new_v4().to_string(),
61 agent_type,
62 capabilities,
63 config: custom_config.unwrap_or_default(),
64 };
65
66 self.agents.insert(agent.id.clone(), agent.clone());
67 Ok(agent)
68 }
69
70 pub async fn get_agent(&self, agent_id: &str) -> Result<Agent> {
71 self.agents
72 .get(agent_id)
73 .cloned()
74 .ok_or_else(|| AIError::AgentNotFound(agent_id.to_string()))
75 }
76
77 pub async fn get_agent_count(&self) -> u64 {
78 self.agents.len() as u64
79 }
80}