gateway/modules/workflow/agentManager.py

"""
Agent Manager Module for managing agent operations and execution.
"""

import os
import logging
import importlib
from typing import Dict, Any, List, Optional, Tuple
from datetime import datetime, UTC
import uuid
from modules.interfaces.serviceChatModel import (
    ChatMessage, ChatDocument, UserInputRequest, ChatWorkflow, AgentResponse
)
logger = logging.getLogger(__name__)

class AgentManager:
    """Manager for agent operations and execution."""
    
    _instance = None
    
    @classmethod
    def getInstance(cls):
        """Return a singleton instance of the agent manager."""
        if cls._instance is None:
            cls._instance = cls()
        return cls._instance

    # Internal Methods

    def __init__(self):
        """Initialize the agent manager."""
        if AgentManager._instance is not None:
            raise RuntimeError("Singleton instance already exists - use getInstance()")
        
        self.service = None
        self.agents = {}  # Dictionary to store agent instances
        self._loadAgents()  # Load agents on initialization
    
    def _loadAgents(self):
        """Load all available agents from modules dynamically."""
        logger.info("Loading agent modules...")
        
        # Get the agents directory path
        agentDir = os.path.join(os.path.dirname(os.path.dirname(__file__)), "agents")
        
        # Search for agent modules
        agentModules = []
        for filename in os.listdir(agentDir):
            if filename.startswith("agent") and filename.endswith(".py"):
                agentModules.append(filename[:-3])  # Remove .py extension
        
        if not agentModules:
            logger.warning("No agent modules found in directory: %s", agentDir)
            return
        
        logger.info(f"Found {len(agentModules)} agent modules: {', '.join(agentModules)}")
        
        # Load each agent module
        for moduleName in agentModules:
            try:
                # Import the module
                module = importlib.import_module(f"modules.agents.{moduleName}")
                
                # Extract agent name from module name
                agentName = moduleName.split("agent")[-1]
                className = f"Agent{agentName}"
                getterName = f"getAgent{agentName}"
                
                agent = None
                
                # Try to get the agent via the getter function first
                if hasattr(module, getterName):
                    getterFunc = getattr(module, getterName)
                    agent = getterFunc()
                    logger.info(f"Agent '{agent.name}' loaded via {getterName}()")
                
                # If no getter, try to instantiate the agent class directly
                elif hasattr(module, className):
                    agentClass = getattr(module, className)
                    agent = agentClass()
                    logger.info(f"Agent '{agent.name}' directly instantiated from {className}")
                
                if agent:
                    # Register the agent
                    if self._registerAgent(agent):
                        logger.info(f"Successfully registered agent: {agent.name}")
                    else:
                        logger.error(f"Failed to register agent from module: {moduleName}")
                else:
                    logger.warning(f"No agent class or getter function found in module: {moduleName}")
            
            except ImportError as e:
                logger.error(f"Failed to import module {moduleName}: {str(e)}")
            except Exception as e:
                logger.error(f"Error loading agent from module {moduleName}: {str(e)}")   

    def _registerAgent(self, agent: Any):
        """Register a new agent with the manager."""
        if not hasattr(agent, 'name'):
            logger.error("Agent must have a name attribute")
            return False
        
        self.agents[agent.name] = agent
        if self.service and hasattr(agent, 'setService'):
            agent.setService(self.service)
        
        return True

    # Public Methods

    def initialize(self, service: Any):
        """Initialize the manager with service reference."""
        # Store service reference
        self.service = service
        
        # Initialize agents with service
        for agent in self.agents.values():
            if hasattr(agent, 'setService'):
                agent.setService(service)
        
        return True
       
    def getAgent(self, agentIdentifier: str) -> Optional[Any]:
        """
        Get an agent instance by its identifier.
        
        Args:
            agentIdentifier: Name or identifier of the agent
            
        Returns:
            Agent instance if found, None otherwise
        """
        agent = self.agents.get(agentIdentifier)
        if not agent:
            logger.warning(f"Agent '{agentIdentifier}' not found")
        return agent
    
    def getAllAgents(self) -> Dict[str, Any]:
        """
        Get all registered agents.
        
        Returns:
            Dictionary mapping agent names to agent instances
        """
        return self.agents.copy()
    
    def getAgentInfos(self) -> List[Dict[str, Any]]:
        """Get information about all registered agents."""
        return [
            {
                'name': agent.name,
                'description': getattr(agent, 'description', ''),
                'capabilities': getattr(agent, 'capabilities', []),
                'inputTypes': getattr(agent, 'inputTypes', []),
                'outputTypes': getattr(agent, 'outputTypes', [])
            }
            for agent in self.agents.values()
        ]
    
    async def executeAgent(self, handover: Any) -> AgentResponse:
        """
        Execute an agent with the given handover.
        
        Args:
            handover: Handover object containing agent execution context
            
        Returns:
            AgentResponse object with execution results
        """
        try:
            # Get agent instance
            agent = self.agents.get(handover.currentAgent)
            if not agent:
                raise ValueError(f"Agent {handover.currentAgent} not found")
            
            # Execute agent
            response = await agent.execute(handover)
            
            # Save output files if any
            if response.message and response.message.documents:
                self.service.document['agentOutputFilesSave'](handover, response.message.documents)
            
            return response
            
        except Exception as e:
            logger.error(f"Error executing agent {handover.currentAgent}: {str(e)}")
            
            # Create error message
            errorMessage = ChatMessage(
                id=str(uuid.uuid4()),
                workflowId=handover.workflowId,
                agentName=handover.currentAgent,
                message=f"Error executing agent: {str(e)}",
                role="system",
                status="error",
                sequenceNr=0,
                startedAt=handover.startedAt,
                finishedAt=datetime.now(UTC).isoformat(),
                success=False
            )
            
            return AgentResponse(
                success=False,
                message=errorMessage,
                error=str(e),
                performance={},
                progress=0.0
            )

# Singleton factory for the agent manager
def getAgentManager():
    return AgentManager.getInstance()