cleanup reasoning agentsg

master
Kye Gomez 1 day ago
parent e3ef675fc1
commit 73dca28bd5

@ -1,6 +1,12 @@
from typing import List, Literal, Dict, Callable, Any, Optional, Tuple, Hashable from typing import (
List,
from functools import lru_cache Literal,
Dict,
Callable,
Any,
Tuple,
Hashable,
)
@ -44,11 +50,9 @@ class ReasoningAgentRouter:
output_type (OutputType): The format of the output (e.g., dict, list). output_type (OutputType): The format of the output (e.g., dict, list).
""" """
# Class variable to store cached agent instances # Class variable to store cached agent instances
_agent_cache: Dict[Tuple[Hashable, ...], Any] = {} _agent_cache: Dict[Tuple[Hashable, ...], Any] = {}
def __init__( def __init__(
self, self,
agent_name: str = "reasoning_agent", agent_name: str = "reasoning_agent",
@ -73,15 +77,10 @@ class ReasoningAgentRouter:
self.num_knowledge_items = num_knowledge_items self.num_knowledge_items = num_knowledge_items
self.memory_capacity = memory_capacity self.memory_capacity = memory_capacity
# Added: Initialize the factory mapping dictionary # Added: Initialize the factory mapping dictionary
self._initialize_agent_factories() self._initialize_agent_factories()
def _initialize_agent_factories(self) -> None: def _initialize_agent_factories(self) -> None:
""" """
Initialize the agent factory mapping dictionary, mapping various agent types to their respective creation functions. Initialize the agent factory mapping dictionary, mapping various agent types to their respective creation functions.
@ -91,12 +90,10 @@ class ReasoningAgentRouter:
# ReasoningDuo factory method # ReasoningDuo factory method
"reasoning-duo": self._create_reasoning_duo, "reasoning-duo": self._create_reasoning_duo,
"reasoning-agent": self._create_reasoning_duo, "reasoning-agent": self._create_reasoning_duo,
# SelfConsistencyAgent factory methods # SelfConsistencyAgent factory methods
"self-consistency": self._create_consistency_agent, "self-consistency": self._create_consistency_agent,
"consistency-agent": self._create_consistency_agent, "consistency-agent": self._create_consistency_agent,
# IREAgent factory methods # IREAgent factory methods
"ire": self._create_ire_agent, "ire": self._create_ire_agent,
"ire-agent": self._create_ire_agent, "ire-agent": self._create_ire_agent,
# Other agent type factory methods # Other agent type factory methods
@ -105,7 +102,6 @@ class ReasoningAgentRouter:
"GKPAgent": self._create_gkp_agent, "GKPAgent": self._create_gkp_agent,
} }
def _get_cache_key(self) -> Tuple[Hashable, ...]: def _get_cache_key(self) -> Tuple[Hashable, ...]:
""" """
Generate a unique key for cache lookup. Generate a unique key for cache lookup.
@ -125,10 +121,9 @@ class ReasoningAgentRouter:
self.num_samples, self.num_samples,
self.output_type, self.output_type,
self.num_knowledge_items, self.num_knowledge_items,
self.memory_capacity self.memory_capacity,
) )
def _create_reasoning_duo(self): def _create_reasoning_duo(self):
"""Create an agent instance for the ReasoningDuo type""" """Create an agent instance for the ReasoningDuo type"""
return ReasoningDuo( return ReasoningDuo(
@ -189,7 +184,6 @@ class ReasoningAgentRouter:
num_knowledge_items=self.num_knowledge_items, num_knowledge_items=self.num_knowledge_items,
) )
def select_swarm(self): def select_swarm(self):
""" """
Select and initialize the appropriate reasoning swarm based on the specified swarm type. Select and initialize the appropriate reasoning swarm based on the specified swarm type.
@ -202,25 +196,23 @@ class ReasoningAgentRouter:
# Generate cache key # Generate cache key
cache_key = self._get_cache_key() cache_key = self._get_cache_key()
# Check if an instance with the same configuration already exists in the cache # Check if an instance with the same configuration already exists in the cache
if cache_key in self.__class__._agent_cache: if cache_key in self.__class__._agent_cache:
return self.__class__._agent_cache[cache_key] return self.__class__._agent_cache[cache_key]
try: try:
# Use the factory method to create a new instance # Use the factory method to create a new instance
agent = self.agent_factories[self.swarm_type]() agent = self.agent_factories[self.swarm_type]()
# Add the newly created instance to the cache # Add the newly created instance to the cache
self.__class__._agent_cache[cache_key] = agent self.__class__._agent_cache[cache_key] = agent
return agent return agent
except KeyError: except KeyError:
# Keep the same error handling as the original code # Keep the same error handling as the original code
raise ValueError(f"Invalid swarm type: {self.swarm_type}") raise ValueError(f"Invalid swarm type: {self.swarm_type}")
def run(self, task: str, *args, **kwargs): def run(self, task: str, *args, **kwargs):
""" """
Execute the reasoning process of the selected swarm on a given task. Execute the reasoning process of the selected swarm on a given task.
@ -236,7 +228,6 @@ class ReasoningAgentRouter:
swarm = self.select_swarm() swarm = self.select_swarm()
return swarm.run(task=task) return swarm.run(task=task)
def batched_run(self, tasks: List[str], *args, **kwargs): def batched_run(self, tasks: List[str], *args, **kwargs):
""" """
Execute the reasoning process on a batch of tasks. Execute the reasoning process on a batch of tasks.
@ -254,8 +245,6 @@ class ReasoningAgentRouter:
results.append(self.run(task, *args, **kwargs)) results.append(self.run(task, *args, **kwargs))
return results return results
@classmethod @classmethod
def clear_cache(cls): def clear_cache(cls):
""" """
@ -263,5 +252,3 @@ class ReasoningAgentRouter:
Use this when you need to free memory or force the creation of new instances. Use this when you need to free memory or force the creation of new instances.
""" """
cls._agent_cache.clear() cls._agent_cache.clear()

Loading…
Cancel
Save