diff --git a/swarms/agents/reasoning_agents.py b/swarms/agents/reasoning_agents.py index e1419ee0..77d2201d 100644 --- a/swarms/agents/reasoning_agents.py +++ b/swarms/agents/reasoning_agents.py @@ -1,4 +1,4 @@ -from typing import List, Literal +from typing import List, Literal, Dict, Callable, Any from swarms.agents.consistency_agent import SelfConsistencyAgent from swarms.agents.flexion_agent import ReflexionAgent @@ -22,7 +22,6 @@ agent_types = Literal[ "AgentJudge", ] - class ReasoningAgentRouter: """ A Reasoning Agent that can answer questions and assist with various tasks using different reasoning strategies. @@ -61,14 +60,104 @@ class ReasoningAgentRouter: self.output_type = output_type self.num_knowledge_items = num_knowledge_items self.memory_capacity = memory_capacity + + # Added: Initialize the factory mapping dictionary + self._initialize_agent_factories() + + # Added: Factory method initialization function + def _initialize_agent_factories(self) -> None: + """ + Initialize the agent factory mapping dictionary, mapping various agent types to their respective creation functions. + This method replaces the original if-elif chain, making the code easier to maintain and extend. + """ + self.agent_factories: Dict[str, Callable[[], Any]] = { + # ReasoningDuo factory methods + "reasoning-duo": self._create_reasoning_duo, + "reasoning-agent": self._create_reasoning_duo, + + # SelfConsistencyAgent factory methods + "self-consistency": self._create_consistency_agent, + "consistency-agent": self._create_consistency_agent, + + # IREAgent factory methods + "ire": self._create_ire_agent, + "ire-agent": self._create_ire_agent, + + # Other agent type factory methods + "AgentJudge": self._create_agent_judge, + "ReflexionAgent": self._create_reflexion_agent, + "GKPAgent": self._create_gkp_agent + } + + # Added: Concrete factory methods for various agent types + def _create_reasoning_duo(self): + """Creates an agent instance for ReasoningDuo type""" + return ReasoningDuo( + agent_name=self.agent_name, + agent_description=self.description, + model_name=[self.model_name, self.model_name], + system_prompt=self.system_prompt, + output_type=self.output_type, + ) + + def _create_consistency_agent(self): + """Creates an agent instance for SelfConsistencyAgent type""" + return SelfConsistencyAgent( + agent_name=self.agent_name, + description=self.description, + model_name=self.model_name, + system_prompt=self.system_prompt, + max_loops=self.max_loops, + num_samples=self.num_samples, + output_type=self.output_type, + ) + + def _create_ire_agent(self): + """Creates an agent instance for IREAgent type""" + return IREAgent( + agent_name=self.agent_name, + description=self.description, + model_name=self.model_name, + system_prompt=self.system_prompt, + max_loops=self.max_loops, + max_iterations=self.num_samples, + output_type=self.output_type, + ) + + def _create_agent_judge(self): + """Creates an agent instance for AgentJudge type""" + return AgentJudge( + agent_name=self.agent_name, + model_name=self.model_name, + system_prompt=self.system_prompt, + max_loops=self.max_loops, + ) + + def _create_reflexion_agent(self): + """Creates an agent instance for ReflexionAgent type""" + return ReflexionAgent( + agent_name=self.agent_name, + system_prompt=self.system_prompt, + model_name=self.model_name, + max_loops=self.max_loops, + ) + + def _create_gkp_agent(self): + """Creates an agent instance for GKPAgent type""" + return GKPAgent( + agent_name=self.agent_name, + model_name=self.model_name, + num_knowledge_items=self.num_knowledge_items, + ) def select_swarm(self): """ Selects and initializes the appropriate reasoning swarm based on the specified swarm type. - Returns: An instance of the selected reasoning swarm. """ + # Commented out original if-elif chain implementation + """ if ( self.swarm_type == "reasoning-duo" or self.swarm_type == "reasoning-agent" @@ -132,6 +221,15 @@ class ReasoningAgentRouter: ) else: raise ValueError(f"Invalid swarm type: {self.swarm_type}") + """ + + # Added: Implementation using factory pattern and dictionary mapping + try: + # Get the corresponding creation function from the factory dictionary and call it + return self.agent_factories[self.swarm_type]() + except KeyError: + # Maintain the same error handling as the original code + raise ValueError(f"Invalid swarm type: {self.swarm_type}") def run(self, task: str, *args, **kwargs): """ @@ -159,4 +257,4 @@ class ReasoningAgentRouter: results = [] for task in tasks: results.append(self.run(task, *args, **kwargs)) - return results + return results \ No newline at end of file