155 lines
No EOL
5.4 KiB
Python
155 lines
No EOL
5.4 KiB
Python
"""
|
|
Agent Registry Module for managing and initializing agents.
|
|
"""
|
|
|
|
import os
|
|
import logging
|
|
import importlib
|
|
from typing import Dict, Any, List, Optional
|
|
from modules.workflow.agentBase import AgentBase
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
class AgentRegistry:
|
|
"""Central registry for all available agents in the system."""
|
|
|
|
_instance = None
|
|
|
|
@classmethod
|
|
def getInstance(cls):
|
|
"""Return 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 getInstance()")
|
|
|
|
self.agents: Dict[str, AgentBase] = {}
|
|
self._loadAgents()
|
|
|
|
def initialize(self, service=None):
|
|
"""Initialize or update the registry with workflow manager and service references."""
|
|
if service:
|
|
# Validate required interfaces
|
|
required_interfaces = ['base', 'msft', 'google']
|
|
missing_interfaces = []
|
|
for interface in required_interfaces:
|
|
if not hasattr(service, interface):
|
|
missing_interfaces.append(interface)
|
|
|
|
if missing_interfaces:
|
|
logger.warning(f"Service container missing required interfaces: {', '.join(missing_interfaces)}")
|
|
return False
|
|
|
|
# Initialize agents with service
|
|
for agent in self.agents.values():
|
|
if service and hasattr(agent, 'setService'):
|
|
agent.setService(service)
|
|
|
|
return True
|
|
|
|
def _loadAgents(self):
|
|
"""Load all available agents from modules."""
|
|
logger.info("Loading agent modules...")
|
|
|
|
# List of agent modules to load
|
|
agentModules = []
|
|
agentDir = os.path.join(os.path.dirname(os.path.dirname(__file__)), "agents")
|
|
|
|
# Search the directory for agent modules
|
|
for filename in os.listdir(agentDir):
|
|
if filename.startswith("agent") and filename.endswith(".py"):
|
|
agentModules.append(filename[0:-3]) # Remove .py extension
|
|
|
|
if not agentModules:
|
|
logger.warning("No agent modules found")
|
|
return
|
|
|
|
logger.info(f"{len(agentModules)} agent modules found")
|
|
|
|
# Load each agent module
|
|
for moduleName in agentModules:
|
|
try:
|
|
# Import the module
|
|
module = importlib.import_module(f"modules.agents.{moduleName}")
|
|
|
|
# Look for agent class or get_*_agent function
|
|
agentName = moduleName.split("agent")[-1]
|
|
className = f"Agent{agentName}"
|
|
getterName = f"getAgent{agentName}"
|
|
|
|
agent = None
|
|
|
|
# Try to get the agent via the get*Agent function
|
|
if hasattr(module, getterName):
|
|
getterFunc = getattr(module, getterName)
|
|
agent = getterFunc()
|
|
logger.info(f"Agent '{agent.name}' loaded via {getterName}()")
|
|
|
|
# Alternatively, try to instantiate the agent directly
|
|
elif hasattr(module, className):
|
|
agentClass = getattr(module, className)
|
|
agent = agentClass()
|
|
logger.info(f"Agent '{agent.name}' directly instantiated")
|
|
|
|
if agent:
|
|
# Register the agent
|
|
self.registerAgent(agent)
|
|
else:
|
|
logger.warning(f"No agent class or getter function found in module {moduleName}")
|
|
|
|
except ImportError as e:
|
|
logger.error(f"Module {moduleName} could not be imported: {e}")
|
|
except Exception as e:
|
|
logger.error(f"Error loading agent from module {moduleName}: {e}")
|
|
|
|
def registerAgent(self, agent):
|
|
"""
|
|
Register an agent in the registry.
|
|
|
|
Args:
|
|
agent: The agent to register
|
|
"""
|
|
agentId = getattr(agent, 'name', "unknown_agent")
|
|
self.agents[agentId] = agent
|
|
logger.debug(f"Agent '{agent.name}' registered")
|
|
|
|
def getAgent(self, agentIdentifier: str):
|
|
"""
|
|
Return an agent instance
|
|
Args:
|
|
agentIdentifier: ID or type of the desired agent
|
|
Returns:
|
|
Agent instance or None if not found
|
|
"""
|
|
if agentIdentifier in self.agents:
|
|
return self.agents[agentIdentifier]
|
|
logger.error(f"Agent with identifier '{agentIdentifier}' not found")
|
|
return None
|
|
|
|
def getAllAgents(self) -> Dict[str, AgentBase]:
|
|
"""
|
|
Get all registered agents.
|
|
|
|
Returns:
|
|
Dictionary mapping agent names to agent instances
|
|
"""
|
|
return self.agents.copy()
|
|
|
|
def getAgentInfos(self) -> List[Dict[str, Any]]:
|
|
"""Return information about all registered agents."""
|
|
agentInfos = []
|
|
seenAgents = set()
|
|
for agent in self.agents.values():
|
|
if agent not in seenAgents:
|
|
agentInfos.append(agent.getAgentInfo())
|
|
seenAgents.add(agent)
|
|
return agentInfos
|
|
|
|
|
|
# Singleton factory for the agent registry
|
|
def getAgentRegistry():
|
|
return AgentRegistry.getInstance() |