File size: 1,449 Bytes
bc9904d
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
from typing import Dict, Any, Optional

class AgentManager:
    """Manages a collection of agents with dynamic registration and retrieval."""
    def __init__(self):
        """Initialize an empty agent registry."""
        self._agents: Dict[str, Any] = {}

    def register_agent(self, agent_id: str, agent: Any) -> None:
        """
        Register a new agent with a unique ID.
        
        Args:
            agent_id (str): Unique identifier for the agent.
            agent (Any): Agent object, expected to have specific structure (e.g., llm/history/prompt or agent/history).
        
        Raises:
            ValueError: If agent_id is already registered.
        """
        if agent_id in self._agents:
            raise ValueError(f"Agent with ID '{agent_id}' is already registered.")
        self._agents[agent_id] = agent
        print(f"[INFO] Registered agent: {agent_id}")

    def get_agent(self, agent_id: str) -> Optional[Any]:
        """
        Retrieve an agent by its ID.
        
        Args:
            agent_id (str): ID of the agent to retrieve.
        
        Returns:
            Optional[Any]: The agent if found, None otherwise.
        """
        return self._agents.get(agent_id)

    def list_agents(self) -> list[str]:
        """
        List all registered agent IDs.
        
        Returns:
            list[str]: List of registered agent IDs.
        """
        return list(self._agents.keys())