325 lines
8.9 KiB
Python
325 lines
8.9 KiB
Python
"""
|
|
Agent registration mechanism.
|
|
Used to manage multiple Agent instances and supports selecting different Agents by model name.
|
|
"""
|
|
|
|
from typing import Dict, Optional, Type, List, Any
|
|
from .BaseAgent import BaseAgent
|
|
from config.config import get_config
|
|
import threading
|
|
|
|
|
|
class AgentRegistry:
|
|
"""Agent registry, managing multiple Agent instances."""
|
|
|
|
def __init__(self):
|
|
self._agents: Dict[str, BaseAgent] = {}
|
|
self._agent_classes: Dict[str, Type[BaseAgent]] = {}
|
|
self._lock = threading.Lock()
|
|
|
|
def register_agent_class(self, model_name: str, agent_class: Type[BaseAgent]):
|
|
"""
|
|
Register an Agent class.
|
|
|
|
Args:
|
|
model_name: Model name, used for the model parameter in the API
|
|
agent_class: Agent class, inheriting from BaseAgent
|
|
"""
|
|
with self._lock:
|
|
self._agent_classes[model_name] = agent_class
|
|
|
|
def _create_agent_instance(
|
|
self,
|
|
model_name: str,
|
|
name: Optional[str] = None,
|
|
description: Optional[str] = None,
|
|
context_file: Optional[str] = None,
|
|
**kwargs
|
|
) -> BaseAgent:
|
|
"""
|
|
Internal method: create an Agent instance.
|
|
|
|
Args:
|
|
model_name: Model name
|
|
name: Agent name
|
|
description: Agent description
|
|
context_file: Context file path
|
|
**kwargs: Other parameters
|
|
|
|
Returns:
|
|
Agent instance
|
|
"""
|
|
if model_name not in self._agent_classes:
|
|
raise ValueError(f"Unknown model: {model_name}")
|
|
|
|
agent_class = self._agent_classes[model_name]
|
|
|
|
# Get configuration.
|
|
config = get_config()
|
|
|
|
# Use default values or passed-in parameters.
|
|
agent_name = name or f"{model_name}-agent"
|
|
agent_description = description or f"Agent instance for {model_name}"
|
|
|
|
# Create the Agent instance.
|
|
agent = agent_class(
|
|
name=agent_name,
|
|
description=agent_description,
|
|
llm_interface=config.BASIC_INTERFACE,
|
|
context_file=context_file,
|
|
model_name=model_name, # Pass model_name to the Agent.
|
|
**kwargs
|
|
)
|
|
|
|
return agent
|
|
|
|
def get_or_create_agent(
|
|
self,
|
|
model_name: str,
|
|
**kwargs
|
|
) -> BaseAgent:
|
|
"""
|
|
Get or create an Agent instance (singleton pattern).
|
|
Ensure each model_name corresponds to only one Agent instance.
|
|
|
|
Args:
|
|
model_name: Model name
|
|
**kwargs: Creation parameters
|
|
|
|
Returns:
|
|
Agent instance
|
|
"""
|
|
with self._lock:
|
|
if model_name not in self._agents:
|
|
self._agents[model_name] = self._create_agent_instance(model_name, **kwargs)
|
|
return self._agents[model_name]
|
|
|
|
def create_agent(
|
|
self,
|
|
model_name: str,
|
|
force_new: bool = False,
|
|
**kwargs
|
|
) -> BaseAgent:
|
|
"""
|
|
Create an Agent instance.
|
|
|
|
Args:
|
|
model_name: Model name
|
|
force_new: Whether to force creation of a new instance, replacing the existing instance
|
|
**kwargs: Other parameters
|
|
|
|
Returns:
|
|
Agent instance
|
|
"""
|
|
with self._lock:
|
|
if force_new or model_name not in self._agents:
|
|
self._agents[model_name] = self._create_agent_instance(model_name, **kwargs)
|
|
return self._agents[model_name]
|
|
|
|
def get_agent(self, model_name: str) -> Optional[BaseAgent]:
|
|
"""
|
|
Get an already-created Agent instance.
|
|
|
|
Args:
|
|
model_name: Model name
|
|
|
|
Returns:
|
|
Agent instance or None
|
|
"""
|
|
return self._agents.get(model_name)
|
|
|
|
def list_models(self) -> List[str]:
|
|
"""
|
|
List all registered model names.
|
|
|
|
Returns:
|
|
List of model names
|
|
"""
|
|
return list(self._agent_classes.keys())
|
|
|
|
def list_agents(self) -> List[str]:
|
|
"""
|
|
List all created Agent instances.
|
|
|
|
Returns:
|
|
List of model names for Agent instances
|
|
"""
|
|
return list(self._agents.keys())
|
|
|
|
def clear_agents(self):
|
|
"""Clear all Agent instances."""
|
|
with self._lock:
|
|
self._agents.clear()
|
|
|
|
def remove_agent(self, model_name: str) -> bool:
|
|
"""
|
|
Remove an Agent instance.
|
|
|
|
Args:
|
|
model_name: Model name
|
|
|
|
Returns:
|
|
Whether removal succeeded
|
|
"""
|
|
with self._lock:
|
|
if model_name in self._agents:
|
|
del self._agents[model_name]
|
|
return True
|
|
return False
|
|
|
|
def get_agent_info(self, model_name: str) -> Optional[Dict[str, Any]]:
|
|
"""
|
|
Get Agent information.
|
|
|
|
Args:
|
|
model_name: Model name
|
|
|
|
Returns:
|
|
Agent information dictionary or None
|
|
"""
|
|
agent = self.get_agent(model_name)
|
|
if agent:
|
|
return {
|
|
"model_name": model_name,
|
|
"name": agent.name,
|
|
"description": agent.description,
|
|
"agent_class": agent.__class__.__name__,
|
|
"toolkit_size": len(agent.toolkit),
|
|
"session_info": agent.get_session_info(),
|
|
"is_singleton": True # Mark this as a singleton instance.
|
|
}
|
|
return None
|
|
|
|
def get_all_agents_info(self) -> Dict[str, Dict[str, Any]]:
|
|
"""
|
|
Get information for all Agents.
|
|
|
|
Returns:
|
|
Dictionary of all Agent information
|
|
"""
|
|
result = {}
|
|
with self._lock:
|
|
for model_name in self._agents:
|
|
info = self.get_agent_info(model_name)
|
|
if info:
|
|
result[model_name] = info
|
|
return result
|
|
|
|
def is_agent_active(self, model_name: str) -> bool:
|
|
"""
|
|
Check whether the Agent with the specified model_name has been created.
|
|
|
|
Args:
|
|
model_name: Model name
|
|
|
|
Returns:
|
|
Whether it has been created
|
|
"""
|
|
return model_name in self._agents
|
|
|
|
def get_agent_stats(self) -> Dict[str, Any]:
|
|
"""
|
|
Get registry statistics.
|
|
|
|
Returns:
|
|
Statistics dictionary
|
|
"""
|
|
with self._lock:
|
|
return {
|
|
"registered_models": len(self._agent_classes),
|
|
"active_agents": len(self._agents),
|
|
"registered_model_list": list(self._agent_classes.keys()),
|
|
"active_agent_list": list(self._agents.keys())
|
|
}
|
|
|
|
|
|
# Global Agent registry instance.
|
|
_global_registry = AgentRegistry()
|
|
|
|
|
|
def get_agent_registry() -> AgentRegistry:
|
|
"""Get the global Agent registry."""
|
|
return _global_registry
|
|
|
|
|
|
def register_agent(model_name: str, agent_class: Type[BaseAgent]):
|
|
"""
|
|
Convenience function for registering an Agent class.
|
|
|
|
Args:
|
|
model_name: Model name
|
|
agent_class: Agent class
|
|
"""
|
|
_global_registry.register_agent_class(model_name, agent_class)
|
|
|
|
|
|
def get_agent(model_name: str, **kwargs) -> BaseAgent:
|
|
"""
|
|
Convenience function for getting or creating an Agent instance (singleton pattern).
|
|
|
|
Args:
|
|
model_name: Model name
|
|
**kwargs: Creation parameters, used only during first creation
|
|
|
|
Returns:
|
|
Agent instance
|
|
"""
|
|
return _global_registry.get_or_create_agent(model_name, **kwargs)
|
|
|
|
|
|
def get_existing_agent(model_name: str) -> Optional[BaseAgent]:
|
|
"""
|
|
Convenience function for getting an existing Agent instance.
|
|
|
|
Args:
|
|
model_name: Model name
|
|
|
|
Returns:
|
|
Agent instance or None
|
|
"""
|
|
return _global_registry.get_agent(model_name)
|
|
|
|
|
|
def create_new_agent(model_name: str, **kwargs) -> BaseAgent:
|
|
"""
|
|
Convenience function for forcing creation of a new Agent instance.
|
|
|
|
Args:
|
|
model_name: Model name
|
|
**kwargs: Creation parameters
|
|
|
|
Returns:
|
|
Agent instance
|
|
"""
|
|
return _global_registry.create_agent(model_name, force_new=True, **kwargs)
|
|
|
|
|
|
def list_available_models() -> List[str]:
|
|
"""
|
|
Convenience function for listing all available models.
|
|
|
|
Returns:
|
|
List of model names
|
|
"""
|
|
return _global_registry.list_models()
|
|
|
|
|
|
def get_registry_stats() -> Dict[str, Any]:
|
|
"""
|
|
Convenience function for getting registry statistics.
|
|
|
|
Returns:
|
|
Statistics dictionary
|
|
"""
|
|
return _global_registry.get_agent_stats()
|
|
|
|
|
|
def clear_all_agents():
|
|
"""
|
|
Convenience function for clearing all Agent instances.
|
|
"""
|
|
_global_registry.clear_agents()
|
|
# Also clear BaseAgent's class-level instance cache.
|
|
from .BaseAgent import BaseAgent
|
|
BaseAgent.clear_instances()
|