gateway/modules/workflow/agentBase.py

"""
Agent Base Module.
Provides the base class for all chat agents.
Defines the standardized interface for task processing.
"""

import os
import logging
import uuid
from datetime import datetime
from typing import Dict, Any, List, Optional
from modules.shared.mimeUtils import isTextMimeType, determineContentEncoding
from modules.interfaces.lucydomModel import ChatContent

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"
        self.label = "Base Agent"
        self.description = "Base agent functionality"
        self.capabilities = []
        self.workflowManager = None
        self.service = None
    
    def setWorkflowManager(self, workflowManager):
        """
        Set the workflow manager reference and validate service container.
        
        Args:
            workflowManager: The workflow manager instance
        """
        if not workflowManager:
            logger.warning("Attempted to set null workflow manager")
            return False
            
        self.workflowManager = workflowManager
        
        # Set service reference from workflow manager if available
        if hasattr(workflowManager, 'service'):
            return self.setService(workflowManager.service)
        return False
    
    def setService(self, service):
        """
        Set the service container reference and validate required interfaces.
        
        Args:
            service: The service container with required interfaces
        """
        if not service:
            logger.warning("Attempted to set null service container")
            return False
            
        # 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
            
        self.service = service
        return True
    
    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,
            "label": self.label,
            "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
                - prompt: The main instruction for the agent
                - inputDocuments: List of document objects to process
                - outputSpecifications: List of required output documents
                - context: Additional contextual information including:
                    - workflow: The complete workflow object
                    - workflowRound: Current workflow round
                    - agentType: Type of agent
                    - timestamp: Task timestamp
                    - language: User language
        
        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"
        """
        # Validate service and workflow manager
        if not self.service:
            logger.error("Service container not initialized")
            return {
                "feedback": "Error: Service container not initialized",
                "documents": []
            }
            
        if not self.workflowManager:
            logger.error("Workflow manager not initialized")
            return {
                "feedback": "Error: Workflow manager not initialized",
                "documents": []
            }
            
        # 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:
        """
        Determine if content should be base64 encoded.
        
        Args:
            filename: Name of the file
            content: Content to check
            mimeType: Optional MIME type
            
        Returns:
            Boolean indicating if content should be base64 encoded
        """
        return determineContentEncoding(filename, content, mimeType)

    def isTextMimeType(self, mimeType: str) -> bool:
        """
        Check if MIME type is text-based.
        
        Args:
            mimeType: MIME type to check
            
        Returns:
            Boolean indicating if MIME type is text-based
        """
        return isTextMimeType(mimeType)

    def formatAgentDocumentOutput(self, label: str, content: str, contentType: str, base64Encoded: bool = False) -> ChatContent:
        """
        Format agent document output using ChatContent model.
        
        Args:
            label: Document label/filename
            content: Document content
            contentType: MIME type of content
            base64Encoded: Whether content is base64 encoded
            
        Returns:
            ChatContent object with the following attributes:
                - sequenceNr: Sequence number (defaults to 1)
                - name: Document label/filename
                - mimeType: MIME type of content
                - data: Actual content
                - metadata: Additional metadata including base64Encoded flag
        """
        return ChatContent(
            sequenceNr=1,
            name=label,
            mimeType=contentType,
            data=content,
            metadata={"base64Encoded": base64Encoded}
        )