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, UTC
from typing import Dict, Any, List, Optional
from modules.shared.mimeUtils import isTextMimeType, determineContentEncoding
from modules.interfaces.serviceChatModel import ChatContent, Task, AgentResponse, ChatMessage

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.service = None
    
    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 execute(self, task: Task) -> AgentResponse:
        """
        Execute a task and return the response.
        This method must be implemented by all concrete agent classes.
        
        Args:
            task: Task object containing all necessary information
            
        Returns:
            AgentResponse object with execution results
        """
        # Validate service manager
        if not self.service:
            logger.error("Service container not initialized")
            return AgentResponse(
                success=False,
                message=ChatMessage(
                    id=str(uuid.uuid4()),
                    workflowId=task.workflowId,
                    agentName=self.name,
                    message="Error: Service container not initialized",
                    role="system",
                    status="error",
                    sequenceNr=0,
                    startedAt=datetime.now(UTC).isoformat(),
                    finishedAt=datetime.now(UTC).isoformat(),
                    success=False
                ),
                performance={},
                progress=0.0
            )
        
        try:
            # Process the task using the concrete implementation
            result = await self.processTask(task)
            
            # Create response message
            message = ChatMessage(
                id=str(uuid.uuid4()),
                workflowId=task.workflowId,
                agentName=self.name,
                message=result.get("feedback", ""),
                role="assistant",
                status="completed",
                sequenceNr=0,
                startedAt=datetime.now(UTC).isoformat(),
                finishedAt=datetime.now(UTC).isoformat(),
                success=True
            )
            
            # Create response with performance metrics
            return AgentResponse(
                success=True,
                message=message,
                performance=result.get("performance", {}),
                progress=result.get("progress", 100.0)
            )
            
        except Exception as e:
            logger.error(f"Error processing task: {str(e)}", exc_info=True)
            return AgentResponse(
                success=False,
                message=ChatMessage(
                    id=str(uuid.uuid4()),
                    workflowId=task.workflowId,
                    agentName=self.name,
                    message=f"Error processing task: {str(e)}",
                    role="system",
                    status="error",
                    sequenceNr=0,
                    startedAt=datetime.now(UTC).isoformat(),
                    finishedAt=datetime.now(UTC).isoformat(),
                    success=False
                ),
                performance={},
                progress=0.0
            )
    
    async def processTask(self, task: Task) -> Dict[str, Any]:
        """
        Process a task and return the results.
        This method must be implemented by all concrete agent classes.
        
        Args:
            task: Task object containing all necessary information
            
        Returns:
            Dictionary containing:
                - feedback: Text response explaining what the agent did
                - performance: Optional performance metrics
                - progress: Task progress (0-100)
        """
        raise NotImplementedError("processTask must be implemented by concrete agent classes")
    
    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}
        )