gateway/modules/chat_registry.py

"""
Chat 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.ai_service = None
    
    def set_dependencies(self, ai_service=None):
        """Set external dependencies for the agent."""
        self.ai_service = ai_service
    
    def get_agent_info(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 process_task(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:
                - task_id: Unique ID for this task
                - workflow_id: ID of the parent workflow (optional)
                - prompt: The main instruction for the agent
                - input_documents: List of document objects to process
                - output_specifications: 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 process_task")
        return {
            "feedback": f"The process_task method was not implemented by agent '{self.name}'.",
            "documents": []
        }


class AgentRegistry:
    """Central registry for all available agents in the system."""
    
    _instance = None
    
    @classmethod
    def get_instance(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 get_instance()")
        
        self.agents = {}
        self.ai_service = None
        self._load_agents()
    
    def _load_agents(self):
        """Load all available agents from modules."""
        logger.info("Loading agent modules...")
        
        # List of agent modules to load
        agent_modules = []
        agent_dir = os.path.dirname(__file__)
        
        # Search the directory for agent modules
        for filename in os.listdir(agent_dir):
            if filename.startswith("chat_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"{len(agent_modules)} agent modules found")
        
        # Load each agent module
        for module_name in agent_modules:
            try:
                # Import the module
                module = importlib.import_module(f"modules.{module_name}")
                
                # Look for agent class or get_*_agent function
                agent_name = module_name.split('_')[-1]
                class_name = f"Agent{agent_name.capitalize()}"
                getter_name = f"get_{agent_name}_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}' 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}' 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_ai_service(self, ai_service):
        """Set the AI service for all agents."""
        self.ai_service = ai_service
        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)
    
    def register_agent(self, agent):
        """
        Register an agent in the registry.
        
        Args:
            agent: The agent to register
        """
        agent_id = getattr(agent, 'name', "unknown_agent")
        # Initialize agent with dependencies
        if hasattr(agent, 'set_dependencies'):
            agent.set_dependencies(ai_service=self.ai_service)
        self.agents[agent_id] = agent
        logger.debug(f"Agent '{agent.name}' registered")
    
    def get_agent(self, agent_identifier: str):
        """
        Return an agent instance
        Args:
            agent_identifier: ID or type of the desired agent
        Returns:
            Agent instance or None if not found
        """
        if agent_identifier in self.agents:
            return self.agents[agent_identifier]
        logger.error(f"Agent with identifier '{agent_identifier}' not found")
        return None
    
    def get_all_agents(self) -> Dict[str, Any]:
        """Return all registered agents."""
        return self.agents
    
    def get_agent_infos(self) -> List[Dict[str, Any]]:
        """Return information about all registered agents."""
        agent_infos = []
        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


# Singleton factory for the agent registry
def get_agent_registry():
    return AgentRegistry.get_instance()