pull/1019/merge
nathanogaga118 1 week ago committed by GitHub
commit 4452b3e726
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194

@ -96,23 +96,24 @@ class AgentRegistry:
Args: Args:
agent (Agent): The agent to add. agent (Agent): The agent to add.
Raises: Raises:
ValueError: If the agent_name already exists in the registry. ValueError: If the agent_name is invalid or already exists in the registry.
ValidationError: If the input data is invalid. ValidationError: If the input data is invalid.
""" """
name = agent.agent_name name = agent.agent_name
# ✅ Validation for agent_name
if not isinstance(name, str) or not name.strip():
logger.error("Invalid agent_name. It must be a non-empty string.")
raise ValueError("Invalid agent_name. It must be a non-empty string.")
self.agent_to_py_model(agent) self.agent_to_py_model(agent)
with self.lock: with self.lock:
if name in self.agents: if name in self.agents:
logger.error( logger.error(f"Agent with name {name} already exists.")
f"Agent with name {name} already exists." raise ValueError(f"Agent with name {name} already exists.")
)
raise ValueError(
f"Agent with name {name} already exists."
)
try: try:
self.agents[name] = agent self.agents[name] = agent
logger.info(f"Agent {name} added successfully.") logger.info(f"Agent {name} added successfully.")
@ -121,7 +122,7 @@ class AgentRegistry:
raise raise
def add_many(self, agents: List[Agent]) -> None: def add_many(self, agents: List[Agent]) -> None:
""" """
Adds multiple agents to the registry. Adds multiple agents to the registry.
Args: Args:
@ -131,11 +132,15 @@ class AgentRegistry:
ValueError: If any of the agent_names already exist in the registry. ValueError: If any of the agent_names already exist in the registry.
ValidationError: If the input data is invalid. ValidationError: If the input data is invalid.
""" """
for agent in agents:
if not isinstance(agent.agent_name, str) or not agent.agent_name.strip():
logger.error(f"Invalid agent_name in batch: {agent.agent_name!r}")
raise ValueError(
f"Invalid agent_name in batch: {agent.agent_name!r}"
)
with ThreadPoolExecutor() as executor: with ThreadPoolExecutor() as executor:
futures = { futures = {executor.submit(self.add, agent): agent for agent in agents}
executor.submit(self.add, agent): agent
for agent in agents
}
for future in as_completed(futures): for future in as_completed(futures):
try: try:
future.result() future.result()
@ -156,9 +161,7 @@ class AgentRegistry:
with self.lock: with self.lock:
try: try:
del self.agents[agent_name] del self.agents[agent_name]
logger.info( logger.info(f"Agent {agent_name} deleted successfully.")
f"Agent {agent_name} deleted successfully."
)
except KeyError as e: except KeyError as e:
logger.error(f"Error: {e}") logger.error(f"Error: {e}")
raise raise
@ -177,17 +180,11 @@ class AgentRegistry:
""" """
with self.lock: with self.lock:
if agent_name not in self.agents: if agent_name not in self.agents:
logger.error( logger.error(f"Agent with name {agent_name} does not exist.")
f"Agent with name {agent_name} does not exist." raise KeyError(f"Agent with name {agent_name} does not exist.")
)
raise KeyError(
f"Agent with name {agent_name} does not exist."
)
try: try:
self.agents[agent_name] = new_agent self.agents[agent_name] = new_agent
logger.info( logger.info(f"Agent {agent_name} updated successfully.")
f"Agent {agent_name} updated successfully."
)
except ValidationError as e: except ValidationError as e:
logger.error(f"Validation error: {e}") logger.error(f"Validation error: {e}")
raise raise
@ -208,9 +205,7 @@ class AgentRegistry:
with self.lock: with self.lock:
try: try:
agent = self.agents[agent_name] agent = self.agents[agent_name]
logger.info( logger.info(f"Agent {agent_name} retrieved successfully.")
f"Agent {agent_name} retrieved successfully."
)
return agent return agent
except KeyError as e: except KeyError as e:
logger.error(f"Error: {e}") logger.error(f"Error: {e}")
@ -233,7 +228,7 @@ class AgentRegistry:
raise e raise e
def return_all_agents(self) -> List[Agent]: def return_all_agents(self) -> List[Agent]:
""" """
Returns all agents from the registry. Returns all agents from the registry.
Returns: Returns:
@ -250,8 +245,8 @@ class AgentRegistry:
def query( def query(
self, condition: Optional[Callable[[Agent], bool]] = None self, condition: Optional[Callable[[Agent], bool]] = None
) -> List[Agent]: ) -> List[Agent]:
""" """
Queries agents based on a condition. Queries agents based on a condition.
Args: Args:
@ -267,12 +262,7 @@ class AgentRegistry:
agents = list(self.agents.values()) agents = list(self.agents.values())
logger.info("Querying all agents.") logger.info("Querying all agents.")
return agents return agents
agents = [agent for agent in self.agents.values() if condition(agent)]
agents = [
agent
for agent in self.agents.values()
if condition(agent)
]
logger.info("Querying agents with condition.") logger.info("Querying agents with condition.")
return agents return agents
except Exception as e: except Exception as e:
@ -310,20 +300,17 @@ class AgentRegistry:
return self.agents.get(agent_id) return self.agents.get(agent_id)
def agents_to_json(self) -> str: def agents_to_json(self) -> str:
""" """
Converts all agents in the registry to a JSON string. Converts all agents in the registry to a JSON string.
Returns: Returns:
str: A JSON string representation of all agents, keyed by their names. str: A JSON string representation of all agents, keyed by their names.
""" """
agents_dict = { agents_dict = {name: agent.to_dict() for name, agent in self.agents.items()}
name: agent.to_dict()
for name, agent in self.agents.items()
}
return json.dumps(agents_dict, indent=4) return json.dumps(agents_dict, indent=4)
def agent_to_py_model(self, agent: Agent): def agent_to_py_model(self, agent: Agent):
""" """
Converts an agent to a Pydantic model. Converts an agent to a Pydantic model.
Args: Args:
@ -331,9 +318,7 @@ class AgentRegistry:
""" """
agent_name = agent.agent_name agent_name = agent.agent_name
agent_description = ( agent_description = (
agent.description agent.description if agent.description else "No description provided"
if agent.description
else "No description provided"
) )
schema = AgentConfigSchema( schema = AgentConfigSchema(
@ -343,10 +328,7 @@ class AgentRegistry:
config=agent.to_dict(), config=agent.to_dict(),
) )
logger.info( logger.info(f"Agent {agent_name} converted to Pydantic model.")
f"Agent {agent_name} converted to Pydantic model."
)
self.agent_registry.agents.append(schema) self.agent_registry.agents.append(schema)

Loading…
Cancel
Save