From ef69df99e25dca0de21a51ae9b907cbb39c35685 Mon Sep 17 00:00:00 2001 From: nathanogaga118 <106948598+nathanogaga118@users.noreply.github.com> Date: Tue, 19 Aug 2025 17:56:29 -0700 Subject: [PATCH] Update agent_registry.py --- swarms/structs/agent_registry.py | 101 ++++++++++++++++++++++++++++++- 1 file changed, 100 insertions(+), 1 deletion(-) diff --git a/swarms/structs/agent_registry.py b/swarms/structs/agent_registry.py index d01490dd..88fec3de 100644 --- a/swarms/structs/agent_registry.py +++ b/swarms/structs/agent_registry.py @@ -41,6 +41,15 @@ class AgentRegistrySchema(BaseModel): class AgentRegistry: """ A class for managing a registry of agents. + + Attributes: + name (str): The name of the registry. + description (str): A description of the registry. + return_json (bool): Indicates whether to return data in JSON format. + auto_save (bool): Indicates whether to automatically save changes to the registry. + agents (Dict[str, Agent]): A dictionary of agents in the registry, keyed by agent name. + lock (Lock): A lock for thread-safe operations on the registry. + agent_registry (AgentRegistrySchema): The schema for the agent registry. """ def __init__( @@ -53,6 +62,16 @@ class AgentRegistry: *args, **kwargs, ): + """ + Initializes the AgentRegistry. + + Args: + name (str, optional): The name of the registry. Defaults to "Agent Registry". + description (str, optional): A description of the registry. Defaults to "A registry for managing agents.". + agents (Optional[List[Agent]], optional): A list of agents to initially add to the registry. Defaults to None. + return_json (bool, optional): Indicates whether to return data in JSON format. Defaults to True. + auto_save (bool, optional): Indicates whether to automatically save changes to the registry. Defaults to False. + """ self.name = name self.description = description self.return_json = return_json @@ -122,6 +141,15 @@ class AgentRegistry: raise def delete(self, agent_name: str) -> None: + """ + Deletes an agent from the registry. + + Args: + agent_name (str): The name of the agent to delete. + + Raises: + KeyError: If the agent_name does not exist in the registry. + """ with self.lock: try: del self.agents[agent_name] @@ -131,6 +159,17 @@ class AgentRegistry: raise def update_agent(self, agent_name: str, new_agent: Agent) -> None: + """ + Updates an existing agent in the registry. + + Args: + agent_name (str): The name of the agent to update. + new_agent (Agent): The new agent to replace the existing one. + + Raises: + KeyError: If the agent_name does not exist in the registry. + ValidationError: If the input data is invalid. + """ with self.lock: if agent_name not in self.agents: logger.error(f"Agent with name {agent_name} does not exist.") @@ -143,6 +182,18 @@ class AgentRegistry: raise def get(self, agent_name: str) -> Agent: + """ + Retrieves an agent from the registry. + + Args: + agent_name (str): The name of the agent to retrieve. + + Returns: + Agent: The agent associated with the given agent_name. + + Raises: + KeyError: If the agent_name does not exist in the registry. + """ with self.lock: try: agent = self.agents[agent_name] @@ -153,6 +204,12 @@ class AgentRegistry: raise def list_agents(self) -> List[str]: + """ + Lists all agent names in the registry. + + Returns: + List[str]: A list of all agent names. + """ try: with self.lock: agent_names = list(self.agents.keys()) @@ -163,6 +220,12 @@ class AgentRegistry: raise e def return_all_agents(self) -> List[Agent]: + """ + Returns all agents from the registry. + + Returns: + List[Agent]: A list of all agents. + """ try: with self.lock: agents = list(self.agents.values()) @@ -172,7 +235,19 @@ class AgentRegistry: logger.error(f"Error: {e}") raise e - def query(self, condition: Optional[Callable[[Agent], bool]] = None) -> List[Agent]: + def query( + self, condition: Optional[Callable[[Agent], bool]] = None + ) -> List[Agent]: + """ + Queries agents based on a condition. + + Args: + condition (Optional[Callable[[Agent], bool]]): A function that takes an agent and returns a boolean indicating + whether the agent meets the condition. + + Returns: + List[Agent]: A list of agents that meet the condition. + """ try: with self.lock: if condition is None: @@ -187,6 +262,15 @@ class AgentRegistry: raise e def find_agent_by_name(self, agent_name: str) -> Optional[Agent]: + """ + Find an agent by its name. + + Args: + agent_name (str): The name of the agent to find. + + Returns: + Agent: The agent with the given name. + """ try: with ThreadPoolExecutor() as executor: futures = { @@ -202,13 +286,28 @@ class AgentRegistry: raise e def find_agent_by_id(self, agent_id: str) -> Optional[Agent]: + """ + Find an agent by its ID. + """ return self.agents.get(agent_id) def agents_to_json(self) -> str: + """ + Converts all agents in the registry to a JSON string. + + Returns: + str: A JSON string representation of all agents, keyed by their names. + """ agents_dict = {name: agent.to_dict() for name, agent in self.agents.items()} return json.dumps(agents_dict, indent=4) def agent_to_py_model(self, agent: Agent): + """ + Converts an agent to a Pydantic model. + + Args: + agent (Agent): The agent to convert. + """ agent_name = agent.agent_name agent_description = ( agent.description if agent.description else "No description provided"