Merge pull request #898 from Wxysnx/feature/refactor/replace-conditionals-with-factory-pattern

replace-conditionals-with-factory-pattern

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