""" 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 from modules.mimeUtils import isTextMimeType, determineContentEncoding logger = logging.getLogger(__name__) """ Updates to the AgentBase class in workflowAgentsRegistry.py to include base64Encoded flag handling. """ 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, each containing a "base64Encoded" flag in addition to "label" and "content" """ # 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": [] } def determineBase64EncodingFlag(self, filename: str, content: Any, mimeType: str = None) -> bool: """Wrapper for the utility function""" return determineContentEncoding(filename, content, mimeType) def isTextMimeType(self, mimeType: str) -> bool: """Wrapper for the utility function""" return isTextMimeType(mimeType) def formatAgentDocumentOutput(self, label: str, content: Any, contentType: str = None) -> Dict[str, Any]: """ Helper method to properly format a document output with base64Encoded flag and metadata. Args: label: Name of the document content: Content of the document contentType: Optional content type for the document Returns: Properly formatted document dictionary """ import base64 # Determine if content should be base64 encoded should_base64_encode = self.determineBase64EncodingFlag(label, content) # Process content based on type and encoding flag formatted_content = content if should_base64_encode: if isinstance(content, bytes): # Convert binary to base64 formatted_content = base64.b64encode(content).decode('utf-8') elif isinstance(content, str): try: # Check if it's already base64 encoded base64.b64decode(content) # If we get here, it appears to be valid base64 formatted_content = content except: # Not valid base64, so encode it formatted_content = base64.b64encode(content.encode('utf-8')).decode('utf-8') # Create document with metadata doc = { "label": label, "content": formatted_content, "base64Encoded": should_base64_encode, "metadata": {} } # Add content type if provided if contentType: doc["metadata"]["contentType"] = contentType return doc 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()