1- from typing import List , Literal , Dict , Callable , Any , Optional , Tuple , Hashable
2-
3- from functools import lru_cache
1+ from typing import (
2+ List ,
3+ Literal ,
4+ Dict ,
5+ Callable ,
6+ Any ,
7+ Tuple ,
8+ Hashable ,
9+ )
410
511
612
@@ -44,11 +50,9 @@ class ReasoningAgentRouter:
4450 output_type (OutputType): The format of the output (e.g., dict, list).
4551 """
4652
47-
4853 # Class variable to store cached agent instances
4954 _agent_cache : Dict [Tuple [Hashable , ...], Any ] = {}
5055
51-
5256 def __init__ (
5357 self ,
5458 agent_name : str = "reasoning_agent" ,
@@ -73,15 +77,10 @@ def __init__(
7377 self .num_knowledge_items = num_knowledge_items
7478 self .memory_capacity = memory_capacity
7579
76-
77-
78-
79-
8080 # Added: Initialize the factory mapping dictionary
8181
8282 self ._initialize_agent_factories ()
8383
84-
8584 def _initialize_agent_factories (self ) -> None :
8685 """
8786 Initialize the agent factory mapping dictionary, mapping various agent types to their respective creation functions.
@@ -91,12 +90,10 @@ def _initialize_agent_factories(self) -> None:
9190 # ReasoningDuo factory method
9291 "reasoning-duo" : self ._create_reasoning_duo ,
9392 "reasoning-agent" : self ._create_reasoning_duo ,
94-
9593 # SelfConsistencyAgent factory methods
9694 "self-consistency" : self ._create_consistency_agent ,
9795 "consistency-agent" : self ._create_consistency_agent ,
9896 # IREAgent factory methods
99-
10097 "ire" : self ._create_ire_agent ,
10198 "ire-agent" : self ._create_ire_agent ,
10299 # Other agent type factory methods
@@ -105,7 +102,6 @@ def _initialize_agent_factories(self) -> None:
105102 "GKPAgent" : self ._create_gkp_agent ,
106103 }
107104
108-
109105 def _get_cache_key (self ) -> Tuple [Hashable , ...]:
110106 """
111107 Generate a unique key for cache lookup.
@@ -125,10 +121,9 @@ def _get_cache_key(self) -> Tuple[Hashable, ...]:
125121 self .num_samples ,
126122 self .output_type ,
127123 self .num_knowledge_items ,
128- self .memory_capacity
124+ self .memory_capacity ,
129125 )
130126
131-
132127 def _create_reasoning_duo (self ):
133128 """Create an agent instance for the ReasoningDuo type"""
134129 return ReasoningDuo (
@@ -189,7 +184,6 @@ def _create_gkp_agent(self):
189184 num_knowledge_items = self .num_knowledge_items ,
190185 )
191186
192-
193187 def select_swarm (self ):
194188 """
195189 Select and initialize the appropriate reasoning swarm based on the specified swarm type.
@@ -202,25 +196,23 @@ def select_swarm(self):
202196
203197 # Generate cache key
204198 cache_key = self ._get_cache_key ()
205-
199+
206200 # Check if an instance with the same configuration already exists in the cache
207201 if cache_key in self .__class__ ._agent_cache :
208202 return self .__class__ ._agent_cache [cache_key ]
209-
210203
211204 try :
212205 # Use the factory method to create a new instance
213206 agent = self .agent_factories [self .swarm_type ]()
214-
207+
215208 # Add the newly created instance to the cache
216209 self .__class__ ._agent_cache [cache_key ] = agent
217-
210+
218211 return agent
219212 except KeyError :
220213 # Keep the same error handling as the original code
221214 raise ValueError (f"Invalid swarm type: { self .swarm_type } " )
222215
223-
224216 def run (self , task : str , * args , ** kwargs ):
225217 """
226218 Execute the reasoning process of the selected swarm on a given task.
@@ -236,7 +228,6 @@ def run(self, task: str, *args, **kwargs):
236228 swarm = self .select_swarm ()
237229 return swarm .run (task = task )
238230
239-
240231 def batched_run (self , tasks : List [str ], * args , ** kwargs ):
241232 """
242233 Execute the reasoning process on a batch of tasks.
@@ -254,14 +245,10 @@ def batched_run(self, tasks: List[str], *args, **kwargs):
254245 results .append (self .run (task , * args , ** kwargs ))
255246 return results
256247
257-
258-
259248 @classmethod
260249 def clear_cache (cls ):
261250 """
262251 Clear the agent instance cache.
263252 Use this when you need to free memory or force the creation of new instances.
264253 """
265254 cls ._agent_cache .clear ()
266-
267-
0 commit comments