gateway/modules/workflowAgentsRegistry.py

"""
Agent Registry Module.
Provides a central registry system for all available agents.
Optimized for the standardized task processing pattern.
"""

import os
import logging
import importlib
import uuid
from datetime import datetime
from typing import Dict, Any, List, Optional

logger = logging.getLogger(__name__)

class AgentBase:
    """
    Base class for all chat agents.
    Defines the standardized interface for task processing.
    """
    
    def __init__(self):
        """Initialize the base agent."""
        self.name = "base-agent"
        self.description = "Basic agent functionality"
        self.capabilities = []
        self.mydom = None
    
    def setDependencies(self, mydom=None):
        """Set external dependencies for the agent."""
        self.mydom = mydom
    
    def getAgentInfo(self) -> Dict[str, Any]:
        """
        Return standardized information about the agent's capabilities.
        
        Returns:
            Dictionary with name, description, and capabilities
        """
        return {
            "name": self.name,
            "description": self.description,
            "capabilities": self.capabilities
        }
    
    async def processTask(self, task: Dict[str, Any]) -> Dict[str, Any]:
        """
        Process a standardized task structure and return results.
        This method must be implemented by all concrete agent classes.
        
        Args:
            task: A dictionary containing:
                - taskId: Unique ID for this task
                - workflowId: ID of the parent workflow (optional)
                - prompt: The main instruction for the agent
                - inputDocuments: List of document objects to process
                - outputSpecifications: List of required output documents
                - context: Additional contextual information
        
        Returns:
            A dictionary containing:
                - feedback: Text response explaining what the agent did
                - documents: List of document objects created by the agent
        """
        # Base implementation - should be overridden by specialized agents
        logger.warning(f"Agent {self.name} is using the default implementation of processTask")
        return {
            "feedback": f"The processTask method was not implemented by agent '{self.name}'.",
            "documents": []
        }


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 = {}
        self.mydom = None
        self._loadAgents()
    
    def _loadAgents(self):
        """Load all available agents from modules."""
        logger.info("Loading agent modules...")
        
        # List of agent modules to load
        agentModules = []
        agentDir = os.path.dirname(__file__)
        
        # 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.{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 setMydom(self, mydom):
        """Set the AI service for all agents."""
        self.mydom = mydom
        self.updateAgentDependencies()
    
    def updateAgentDependencies(self):
        """Update dependencies for all registered agents."""
        for agentId, agent in self.agents.items():
            if hasattr(agent, 'setDependencies'):
                agent.setDependencies(mydom=self.mydom)
    
    def registerAgent(self, agent):
        """
        Register an agent in the registry.
        
        Args:
            agent: The agent to register
        """
        agentId = getattr(agent, 'name', "unknown_agent")
        # Initialize agent with dependencies
        if hasattr(agent, 'setDependencies'):
            agent.setDependencies(mydom=self.mydom)
        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:
            agent = self.agents[agentIdentifier]
            # Ensure the agent has the AI service
            if hasattr(agent, 'setDependencies') and self.mydom:
                agent.setDependencies(mydom=self.mydom)
            return agent
        logger.error(f"Agent with identifier '{agentIdentifier}' not found")
        return None
    
    def getAllAgents(self) -> Dict[str, Any]:
        """Return all registered agents."""
        return self.agents
    
    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()