gateway/modules/agentservice_base.py

"""
Enhanced base agent class for the Agentservice.
Provides improved communication and document handling capabilities.
"""

import logging
import json
from typing import Dict, Any, List, Optional, Tuple, Union
import asyncio
from datetime import datetime
import uuid

logger = logging.getLogger(__name__)

class AgentBase:
    """
    Enhanced base agent class with improved communication capabilities.
    All specialized agents should inherit from this class.
    """
    
    def __init__(self):
        """Initialize the enhanced agent."""
        self.name = "base"
        self.capabilities = "Basic agent operations"
        self.result_format = "Text"
        # System dependencies
        self.ai_service = None
    
    def set_dependencies(self, ai_service=None, document_handler=None, lucydom_interface=None):
        self.ai_service = ai_service
    
    def get_config(self) -> Dict[str, Any]:
        """
        Get detailed information about the agent.
        
        Returns:
            Dictionary with agent information
        """
        return {
            "name": self.name,
            "capabilities": self.capabilities,
            "result_format": self.result_format,
        }
    
    def get_capabilities(self) -> List[str]:
        """
        Get a list of agent capabilities.
        
        Returns:
            List of capability strings
        """
        # Split capabilities into a list
        if isinstance(self.capabilities, str):
            return [cap.strip() for cap in self.capabilities.split(",")]
        return []
    
    def get_supported_formats(self) -> List[str]:
        """
        Get supported output formats.
        
        Returns:
            List of supported format strings
        """
        if isinstance(self.result_format, str):
            return [fmt.strip() for fmt in self.result_format.split(",")]
        return ["Text"]
    
    async def process_message(self, message: Dict[str, Any], context: Dict[str, Any] = None) -> Dict[str, Any]:
        """
        Process a message and generate a response.
        
        Args:
            message: Input message
            context: Optional context information
            
        Returns:
            Response message
        """
        # Basic implementation - should be overridden by specialized agents
        if not self.ai_service:
            logger.warning(f"Agent {self.id} has no AI service configured")
            return {
                "role": "assistant",
                "content": f"I'm {self.name}, but I'm not properly configured. Please set up the AI service.",
                "agent_id": self.id,
                "agent_type": self.type,
                "result_format": "Text"
            }
        
        # Process documents if available and set up document handler
        document_context = ""
        if self.supports_documents and self.document_handler and message.get("documents"):
            document_context = await self._process_documents(message)
        
        # Create enhanced prompt
        prompt = self._create_enhanced_prompt(message, document_context, context)
        
        # Generate response
        try:
            response_content = await self.ai_service.call_api([
                {"role": "system", "content": self._get_system_prompt()},
                {"role": "user", "content": prompt}
            ])
            
            # Process the response to extract any special instructions or status
            content, status = self._process_response(response_content)
            
            return {
                "role": "assistant",
                "content": content,
                "agent_id": self.id,
                "agent_type": self.type,
                "agent_name": self.name,
                "result_format": self.result_format,
                "status": status,
                "workflow_id": message.get("workflow_id"),
                "documents": message.get("documents", [])  # Pass through documents
            }
        except Exception as e:
            logger.error(f"Error in agent {self.id}: {str(e)}")
            return {
                "role": "assistant",
                "content": f"I encountered an error: {str(e)}",
                "agent_id": self.id,
                "agent_type": self.type,
                "result_format": "Text",
                "status": "error"
            }
    
    async def _process_documents(self, message: Dict[str, Any]) -> str:
        """
        Process documents in the message.
        
        Args:
            message: Input message with documents
            
        Returns:
            Document context as text
        """
        # Simply extract text from documents
        if not self.document_handler:
            return ""
        
        return self.document_handler.merge_document_contents(message)
    
    def _create_enhanced_prompt(self, message: Dict[str, Any], document_context: str, context: Dict[str, Any] = None) -> str:
        """
        Create an enhanced prompt with context.
        
        Args:
            message: Input message
            document_context: Document context
            context: Optional additional context
            
        Returns:
            Enhanced prompt
        """
        prompt = message.get("content", "")
        
        # Add document context if available
        if document_context:
            prompt += f"\n\n=== DOCUMENT CONTEXT ===\n{document_context}"
        
        # Add any additional context
        if context:
            # Add expected format if specified
            if "expected_format" in context:
                prompt += f"\n\nPlease format your response as: {context['expected_format']}"
            
            # Add dependency outputs if available
            if "dependency_outputs" in context:
                prompt += "\n\n=== OUTPUTS FROM PREVIOUS ACTIVITIES ===\n"
                for key, value in context["dependency_outputs"].items():
                    if isinstance(value, dict) and "content" in value:
                        prompt += f"\n--- {key} ---\n{value['content']}\n"
                    else:
                        prompt += f"\n--- {key} ---\n{str(value)}\n"
        
        return prompt
    
    def _get_system_prompt(self) -> str:
        """
        Get the system prompt for the agent.
        
        Returns:
            System prompt string
        """
        return f"""
        You are {self.name}, a specialized {self.type} agent.
        
        {self.description}
        
        Your capabilities include: {self.capabilities}
        
        You should format your responses according to: {self.result_format}
        
        Respond clearly and helpfully to the user's request. 
        When appropriate, include a status indicator at the end of your message:
        
        [STATUS: COMPLETE] - When you've fully addressed the request
        [STATUS: PARTIAL] - When you've partially addressed the request
        [STATUS: QUESTION] - When you need more information
        """
    
    def _process_response(self, response: str) -> Tuple[str, str]:
        """
        Process the response to extract status and clean content.
        
        Args:
            response: Raw response from the AI
            
        Returns:
            Tuple of (cleaned content, status)
        """
        # Default status
        status = "complete"
        
        # Check for status tags
        import re
        status_match = re.search(r'\[STATUS:\s*(COMPLETE|PARTIAL|QUESTION)\]', response, re.IGNORECASE)
        
        if status_match:
            status_value = status_match.group(1).lower()
            # Remove the status tag
            content = re.sub(r'\[STATUS:\s*(COMPLETE|PARTIAL|QUESTION)\]', '', response, flags=re.IGNORECASE).strip()
            return content, status_value
        
        return response, status

# Factory functions
def get_enhanced_base_agent() -> BaseAgent:
    """Get an instance of the enhanced base agent."""
    return BaseAgent()