292 lines
No EOL
11 KiB
Python
292 lines
No EOL
11 KiB
Python
"""
|
|
Updated registry for all available agents in the system.
|
|
Provides centralized agent registration and access with improved error handling.
|
|
"""
|
|
|
|
import os
|
|
import logging
|
|
import importlib
|
|
from typing import Dict, Any, List, Optional
|
|
|
|
# Import direct base agent module
|
|
from modules.agentservice_base import BaseAgent
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
class AgentRegistry:
|
|
"""Registry for all available agents in the system"""
|
|
|
|
_instance = None
|
|
|
|
@classmethod
|
|
def get_instance(cls):
|
|
"""Get a singleton instance of the Agent Registry"""
|
|
if cls._instance is None:
|
|
cls._instance = cls()
|
|
return cls._instance
|
|
|
|
def __init__(self):
|
|
"""Initialize the Agent Registry"""
|
|
if AgentRegistry._instance is not None:
|
|
raise RuntimeError("Singleton instance already exists - use get_instance()")
|
|
self.agents = {}
|
|
self.ai_service = None
|
|
self.document_handler = None
|
|
self.lucydom_interface = None
|
|
self._load_agents()
|
|
|
|
def _load_agents(self):
|
|
"""Load all available agents"""
|
|
# List of all agent modules to load
|
|
logger.info("Automatically loading agent modules...")
|
|
agent_modules = []
|
|
for filename in os.listdir(os.path.dirname(__file__)):
|
|
if filename.startswith("agentservice_agent_") and filename.endswith(".py"):
|
|
agent_modules.append(filename[:-3]) # Remove .py extension
|
|
if not agent_modules:
|
|
logger.warning("No agent modules found")
|
|
return
|
|
logger.info(f"Found {len(agent_modules)} agent modules")
|
|
|
|
for module_name in agent_modules:
|
|
try:
|
|
# Import the module
|
|
module = importlib.import_module(f"modules.{module_name}")
|
|
|
|
# Look for the agent class or a get_*_agent function
|
|
agent_type = module_name.split('_')[-1]
|
|
class_name = f"{agent_type.capitalize()}Agent"
|
|
getter_name = f"get_{agent_type}_agent"
|
|
|
|
agent = None
|
|
|
|
# Try to get the agent via the get_*_agent function
|
|
if hasattr(module, getter_name):
|
|
getter_func = getattr(module, getter_name)
|
|
agent = getter_func()
|
|
logger.info(f"Agent '{agent.name}' (Type: {agent.type}) loaded via {getter_name}()")
|
|
|
|
# Alternatively, try to instantiate the agent directly
|
|
elif hasattr(module, class_name):
|
|
agent_class = getattr(module, class_name)
|
|
agent = agent_class()
|
|
logger.info(f"Agent '{agent.name}' (Type: {agent.type}) directly instantiated")
|
|
|
|
if agent:
|
|
# Register the agent
|
|
self.register_agent(agent)
|
|
else:
|
|
logger.warning(f"No agent class or getter function found in module {module_name}")
|
|
|
|
except ImportError as e:
|
|
logger.error(f"Module {module_name} could not be imported: {e}")
|
|
except Exception as e:
|
|
logger.error(f"Error loading agent from module {module_name}: {e}")
|
|
|
|
def set_dependencies(self, ai_service=None, document_handler=None, lucydom_interface=None):
|
|
"""
|
|
Set system dependencies for all agents.
|
|
|
|
Args:
|
|
ai_service: AI service for text generation
|
|
document_handler: Document handler for document operations
|
|
lucydom_interface: LucyDOM interface for database access
|
|
"""
|
|
self.ai_service = ai_service
|
|
self.document_handler = document_handler
|
|
self.lucydom_interface = lucydom_interface
|
|
# Update all registered agents
|
|
self.update_agent_dependencies()
|
|
|
|
|
|
def update_agent_dependencies(self):
|
|
"""Update dependencies for all registered agents"""
|
|
for agent_id, agent in self.agents.items():
|
|
if hasattr(agent, 'set_dependencies'):
|
|
agent.set_dependencies(
|
|
ai_service=self.ai_service,
|
|
document_handler=self.document_handler,
|
|
lucydom_interface=self.lucydom_interface
|
|
)
|
|
|
|
def register_agent(self, agent: 'BaseAgent'):
|
|
"""
|
|
Register an agent in the registry.
|
|
|
|
Args:
|
|
agent: The agent to register
|
|
"""
|
|
agent_type = agent.type
|
|
agent_id = getattr(agent, 'id', agent_type)
|
|
|
|
# Initialize enhanced agents with dependencies
|
|
if hasattr(agent, 'set_dependencies'):
|
|
agent.set_dependencies(
|
|
ai_service=self.ai_service,
|
|
document_handler=self.document_handler,
|
|
lucydom_interface=self.lucydom_interface
|
|
)
|
|
|
|
self.agents[agent_type] = agent
|
|
# Also register by ID if it's different from type
|
|
if agent_id != agent_type:
|
|
self.agents[agent_id] = agent
|
|
|
|
logger.debug(f"Agent '{agent.name}' (Type: {agent_type}, ID: {agent_id}) registered")
|
|
|
|
def get_agent(self, agent_identifier: str) -> Optional[BaseAgent]:
|
|
"""
|
|
Get an agent instance by ID or type.
|
|
|
|
Args:
|
|
agent_identifier: ID or type of the desired agent
|
|
|
|
Returns:
|
|
Agent instance or None if not found
|
|
"""
|
|
# Try to find directly by type
|
|
if agent_identifier in self.agents:
|
|
return self.agents[agent_identifier]
|
|
|
|
# If not found, try different name variants
|
|
variants = [
|
|
agent_identifier,
|
|
agent_identifier.replace('_agent', ''),
|
|
f"{agent_identifier}_agent"
|
|
]
|
|
|
|
for variant in variants:
|
|
if variant in self.agents:
|
|
return self.agents[variant]
|
|
|
|
logger.warning(f"Agent with identifier '{agent_identifier}' not found")
|
|
return None
|
|
|
|
def get_all_agents(self) -> Dict[str, BaseAgent]:
|
|
"""Get all registered agents."""
|
|
return self.agents
|
|
|
|
def get_agent_infos(self) -> List[Dict[str, Any]]:
|
|
"""Get information about all registered agents."""
|
|
agent_infos = []
|
|
# Only once per agent instance (since we register both by type and ID)
|
|
seen_agents = set()
|
|
for agent in self.agents.values():
|
|
if agent not in seen_agents:
|
|
agent_infos.append(agent.get_agent_info())
|
|
seen_agents.add(agent)
|
|
return agent_infos
|
|
|
|
def get_agent_by_format(self, required_format: str) -> Optional[BaseAgent]:
|
|
"""
|
|
Find an agent that can produce the required output format.
|
|
|
|
Args:
|
|
required_format: The required output format
|
|
|
|
Returns:
|
|
Agent that can produce the required format, or None if not found
|
|
"""
|
|
# Create mapping of result format -> agent for faster lookup
|
|
format_to_agent = {}
|
|
seen_agents = set()
|
|
|
|
for agent in self.agents.values():
|
|
if agent not in seen_agents:
|
|
# Get the agent's result format
|
|
agent_format = getattr(agent, 'result_format', None)
|
|
if agent_format:
|
|
format_to_agent[agent_format.lower()] = agent
|
|
seen_agents.add(agent)
|
|
|
|
# Try to find an exact match
|
|
if required_format.lower() in format_to_agent:
|
|
return format_to_agent[required_format.lower()]
|
|
|
|
# If no exact match, try to find a partial match
|
|
for fmt, agent in format_to_agent.items():
|
|
if required_format.lower() in fmt or fmt in required_format.lower():
|
|
return agent
|
|
|
|
# No match found
|
|
return None
|
|
|
|
def initialize_agents_for_workflow(self) -> Dict[str, Dict[str, Any]]:
|
|
"""Initialize agents for a workflow."""
|
|
initialized_agents = {}
|
|
seen_agents = set()
|
|
for agent in self.agents.values():
|
|
if agent not in seen_agents:
|
|
agent_info = agent.get_agent_info()
|
|
agent_id = agent_info["id"]
|
|
initialized_agents[agent_id] = agent_info
|
|
seen_agents.add(agent)
|
|
return initialized_agents
|
|
|
|
def get_agent_capabilities(self) -> Dict[str, List[str]]:
|
|
"""
|
|
Get a mapping of capabilities to agents.
|
|
Useful for finding the right agent for a specific task.
|
|
|
|
Returns:
|
|
Dict mapping capability keywords to agent IDs
|
|
"""
|
|
capabilities_map = {}
|
|
seen_agents = set()
|
|
|
|
for agent in self.agents.values():
|
|
if agent not in seen_agents:
|
|
# Get agent info
|
|
agent_id = getattr(agent, 'id', agent.type)
|
|
|
|
# Extract capabilities - check for get_capabilities method first
|
|
if hasattr(agent, 'get_capabilities') and callable(getattr(agent, 'get_capabilities')):
|
|
capabilities = agent.get_capabilities()
|
|
else:
|
|
# Fall back to string parsing
|
|
capabilities_str = getattr(agent, 'capabilities', "")
|
|
capabilities = [kw.strip().lower() for kw in capabilities_str.split(',') if kw.strip()]
|
|
|
|
# Add each capability to the mapping
|
|
for capability in capabilities:
|
|
if capability not in capabilities_map:
|
|
capabilities_map[capability] = []
|
|
if agent_id not in capabilities_map[capability]:
|
|
capabilities_map[capability].append(agent_id)
|
|
|
|
seen_agents.add(agent)
|
|
|
|
return capabilities_map
|
|
|
|
def get_agent_by_capability(self, capability: str) -> Optional['BaseAgent']:
|
|
"""
|
|
Find an agent with a specific capability.
|
|
|
|
Args:
|
|
capability: The required capability
|
|
|
|
Returns:
|
|
Agent with the required capability, or None if not found
|
|
"""
|
|
# Create mapping of capabilities for faster lookup
|
|
capability_map = self.get_agent_capabilities()
|
|
|
|
# Look for the capability (case-insensitive)
|
|
capability = capability.lower()
|
|
matching_agents = []
|
|
|
|
# Direct match
|
|
if capability in capability_map:
|
|
matching_agents = capability_map[capability]
|
|
else:
|
|
# Partial matches
|
|
for cap, agents in capability_map.items():
|
|
if capability in cap or cap in capability:
|
|
matching_agents.extend(agents)
|
|
|
|
# Return the first matching agent
|
|
if matching_agents:
|
|
agent_id = matching_agents[0]
|
|
return self.get_agent(agent_id)
|
|
|
|
return None |