gateway/gwserver/modules/agentservice_agent_documentation.py

"""
Dokumentations-Agent für die Erstellung von Dokumentation, Berichten und strukturierten Inhalten.
Verwendet einen adaptiven Prozess zur Erstellung hochwertiger Dokumentation basierend auf der Komplexität des Auftrags.
Angepasst für das refaktorisierte Core-Modul und AgentCommunicationProtocol.
"""

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

from modules.agentservice_base import BaseAgent
from connectors.connector_aichat_openai import ChatService
from modules.agentservice_utils import WorkflowUtils, MessageUtils, LoggingUtils
from modules.agentservice_protocol import AgentMessage, AgentCommunicationProtocol
from modules.agentservice_filemanager import FileManager  # Import the file manager

logger = logging.getLogger(__name__)

class DocumentationAgent(BaseAgent):
    """Agent for creating documentation and structured content"""
    
    def __init__(self):
        """Initialize the documentation agent"""
        super().__init__()
        self.id = "documentation_agent"
        self.name = "Documentation Specialist"
        self.type = "documentation"
        self.description = "Creates documentation and structured content"
        self.capabilities = "report_generation,documentation,content_structuring,technical_writing,knowledge_organization"
        self.result_format = "FormattedDocument"

        # Initialize AI service
        self.ai_service = None

        # Initialize document handler
        self.document_handler = None
        
        # Document capabilities
        self.supports_documents = True
        self.document_capabilities = ["read", "reference", "create"]
        self.required_context = ["document_purpose", "target_audience"]
        
        # Initialize protocol
        self.protocol = AgentCommunicationProtocol()
        
        # Initialize utilities
        self.message_utils = MessageUtils()
        
        # Track the latest generated document
        self.last_document = {}
    
    def get_agent_info(self) -> Dict[str, Any]:
        """Get agent information for agent registry"""
        info = super().get_agent_info()
        info.update({
            "metadata": {
                "document_types": ["manual", "report", "process", "presentation", "document"],
                "formats": ["markdown", "text"]
            }
        })
        return info
    
    async def process_message(self, message: Dict[str, Any], context: Dict[str, Any] = None) -> Dict[str, Any]:
        """
        Process a message and create documentation.
        
        Args:
            message: Input message
            context: Optional context
            
        Returns:
            Response with documentation
        """
        # Extract workflow_id from context or message
        workflow_id = context.get("workflow_id") if context else message.get("workflow_id", "unknown")
        
        # Get or create logging_utils
        log_func = context.get("log_func") if context else None
        logging_utils = LoggingUtils(workflow_id, log_func)
        
        # Create response structure
        response = {
            "role": "assistant",
            "content": "",
            "agent_id": self.id,
            "agent_type": self.type,
            "agent_name": self.name,
            "result_format": self.result_format,
            "workflow_id": workflow_id,
            "documents": []
        }
        
        try:
            # Create status update using protocol
            if log_func:
                status_message = self.protocol.create_status_update_message(
                    status_description="Starting document creation",
                    sender_id=self.id,
                    status="in_progress",
                    progress=0.0,
                    context_id=workflow_id
                )
                log_func(workflow_id, status_message.content, "info", self.id, self.name)
            
            # Extract task from message
            task = message.get("content", "")
            
            # Detect document type
            document_type = self._detect_document_type(task)
            logging_utils.info(f"Creating {document_type} documentation", "execution")
            
            # Process any attached documents
            document_context = ""
            if message.get("documents"):
                logging_utils.info("Processing reference documents", "execution")
                document_context = await self._process_documents(message)
                
                # Update progress
                if log_func:
                    status_message = self.protocol.create_status_update_message(
                        status_description="Reference documents processed",
                        sender_id=self.id,
                        status="in_progress",
                        progress=0.3,
                        context_id=workflow_id
                    )
                    log_func(workflow_id, status_message.content, "info", self.id, self.name)
            
            # Enhanced prompt with document context
            enhanced_prompt = f"{task}\n\n{document_context}"
            
            # Assess complexity of the task
            is_complex = await self._assess_complexity(enhanced_prompt)
            
            # Generate title
            title = await self._generate_title(enhanced_prompt, document_type)
            logging_utils.info(f"Document title: {title}", "execution")
            
            # Update progress
            if log_func:
                status_message = self.protocol.create_status_update_message(
                    status_description=f"Generating {document_type}: {title}",
                    sender_id=self.id,
                    status="in_progress",
                    progress=0.5,
                    context_id=workflow_id
                )
                log_func(workflow_id, status_message.content, "info", self.id, self.name)
            
            # Generate content based on complexity
            if is_complex:
                # For complex documents, use the AI service with enhanced prompt
                content = await self._generate_complex_document(enhanced_prompt, document_type, title)
                logging_utils.info("Complex document generated", "execution")
            else:
                # For simple documents, use direct generation
                content = await self._generate_simple_document(enhanced_prompt, document_type, title)
                logging_utils.info("Simple document generated", "execution")
            
            # Final progress update
            if log_func:
                status_message = self.protocol.create_status_update_message(
                    status_description="Document creation completed",
                    sender_id=self.id,
                    status="completed",
                    progress=1.0,
                    context_id=workflow_id
                )
                log_func(workflow_id, status_message.content, "info", self.id, self.name)
            
            # Create a document artifact if document handler is available
            if self.document_handler:
                doc_id = f"doc_{uuid.uuid4()}"
                document = {
                    "id": doc_id,
                    "source": {
                        "type": "generated",
                        "id": doc_id,
                        "name": title,
                        "content_type": "text/markdown",
                        "size": len(content)
                    },
                    "contents": [
                        {
                            "type": "text",
                            "text": content,
                            "is_extracted": True
                        }
                    ]
                }
                
                # Add document to response
                response["documents"].append(document)
                
                # Store the latest document
                self.last_document = document
                
                # Update response content to reference the document
                response["content"] = f"I've created a document titled '{title}' that contains the requested information. The document is attached to this message."
                
                # If protocol message is required, send it
                if context and context.get("require_protocol_message"):
                    result_message = self.send_document_result(
                        document_title=title,
                        document_content=content,
                        sender_id=self.id,
                        receiver_id=context.get("receiver_id", "workflow"),
                        context_id=workflow_id
                    )
                    # Just log the message creation
                    logging_utils.info(f"Created protocol result message: {result_message.id}", "execution")
            else:
                # If no document handler, just put content in response
                response["content"] = content
            
            return response
            
        except Exception as e:
            error_msg = f"Error in documentation agent: {str(e)}"
            logging_utils.error(error_msg, "error")
            
            # Create error response using protocol
            error_message = self.protocol.create_error_message(
                error_description=error_msg,
                sender_id=self.id,
                error_type="documentation",
                error_details={"traceback": traceback.format_exc()},
                context_id=workflow_id
            )
            
            # Set error in response
            response["content"] = f"## Error creating documentation\n\n{error_msg}\n\n```\n{traceback.format_exc()}\n```"
            response["status"] = "error"
            
            return response
    
    async def _assess_complexity(self, task: str) -> bool:
        """
        Assess task complexity to determine document structure.
        
        Args:
            task: The task description
            
        Returns:
            True if complex document needed, False otherwise
        """
        if not self.ai_service:
            # Default to complex if no AI service
            return True
        
        prompt = f"""
        Analyze this task and determine if it requires a complex or simple document structure:
        
        {task}
        
        Respond with only "COMPLEX" or "SIMPLE".
        """
        
        try:
            response = await self.ai_service.call_api([
                {"role": "system", "content": "You determine document complexity requirements."},
                {"role": "user", "content": prompt}
            ])
            
            return "COMPLEX" in response.upper()
        except Exception:
            # Default to complex on error
            return True
    
    async def _generate_title(self, task: str, document_type: str) -> str:
        """
        Generate a title for the document.
        
        Args:
            task: The task description
            document_type: Type of document
            
        Returns:
            Generated title
        """
        if not self.ai_service:
            # Default title if no AI service
            return f"{document_type.capitalize()} Document"
        
        prompt = f"""
        Create a concise, professional title for this {document_type}:
        
        {task}
        
        Respond with ONLY the title, nothing else.
        """
        
        try:
            title = await self.ai_service.call_api([
                {"role": "system", "content": "You create document titles."},
                {"role": "user", "content": prompt}
            ])
            
            # Clean up the title
            return title.strip('"\'#*- \n\t')
        except Exception:
            # Default title on error
            return f"{document_type.capitalize()} Document"
    
    async def _generate_complex_document(self, task: str, document_type: str, title: str) -> str:
        """
        Generate a complex document with structure.
        
        Args:
            task: The task description
            document_type: Type of document
            title: Document title
            
        Returns:
            Generated document content
        """
        if not self.ai_service:
            return f"# {title}\n\nUnable to generate complex document: AI service not available."
        
        prompt = f"""
        Create a comprehensive, well-structured {document_type} titled "{title}" based on:
        
        {task}
        
        The document should include:
        1. A clear introduction with purpose and scope
        2. Logically organized sections with headings
        3. Detailed content with examples and evidence
        4. A conclusion with key takeaways
        5. Appropriate formatting using Markdown
        
        Format the document in Markdown with proper headings, lists, and emphasis.
        """
        
        try:
            content = await self.ai_service.call_api([
                {"role": "system", "content": "You create comprehensive, well-structured documentation."},
                {"role": "user", "content": prompt}
            ])
            
            # Ensure title is at the top
            if not content.strip().startswith("# "):
                content = f"# {title}\n\n{content}"
            
            return content
        except Exception as e:
            return f"# {title}\n\nError generating document: {str(e)}"
    
    async def _generate_simple_document(self, task: str, document_type: str, title: str) -> str:
        """
        Generate a simple document without complex structure.
        
        Args:
            task: The task description
            document_type: Type of document
            title: Document title
            
        Returns:
            Generated document content
        """
        if not self.ai_service:
            return f"# {title}\n\nUnable to generate document: AI service not available."
        
        prompt = f"""
        Create a concise, focused {document_type} titled "{title}" based on:
        
        {task}
        
        The document should be clear, precise, and to the point without complex chapter structure.
        Format using Markdown with appropriate headings and formatting.
        """
        
        try:
            content = await self.ai_service.call_api([
                {"role": "system", "content": "You create concise, focused documentation."},
                {"role": "user", "content": prompt}
            ])
            
            # Ensure title is at the top
            if not content.strip().startswith("# "):
                content = f"# {title}\n\n{content}"
            
            return content
        except Exception as e:
            return f"# {title}\n\nError generating document: {str(e)}"
    
    def _detect_document_type(self, message: str) -> str:
        """
        Detect document type from the message.
        
        Args:
            message: User message
            
        Returns:
            Detected document type
        """
        message = message.lower()
        
        if any(term in message for term in ["manual", "guide", "instruction", "tutorial"]):
            return "manual"
        elif any(term in message for term in ["report", "analysis", "assessment", "review"]):
            return "report"
        elif any(term in message for term in ["process", "workflow", "procedure", "steps"]):
            return "process"
        elif any(term in message for term in ["presentation", "slides", "deck"]):
            return "presentation"
        else:
            return "document"
    
    def send_document_result(self, document_title: str, document_content: str, 
                           sender_id: str, receiver_id: str, context_id: str = None) -> AgentMessage:
        """Send a document result using the protocol"""
        metadata = {
            "document_type": self._detect_document_type(document_content),
            "title": document_title,
            "created_at": datetime.now().isoformat()
        }
        
        return self.protocol.create_result_message(
            result_content=document_content,
            sender_id=sender_id,
            receiver_id=receiver_id,
            task_id=f"doc_{uuid.uuid4()}",
            output_data=metadata,
            result_format=self.result_format,
            context_id=context_id
        )
    
    def send_error_message(self, error_description: str, sender_id: str, receiver_id: str = None, 
                         context_id: str = None) -> AgentMessage:
        """Send an error message using the protocol"""
        return self.protocol.create_error_message(
            error_description=error_description,
            sender_id=sender_id,
            receiver_id=receiver_id,
            error_type="documentation_error",
            error_details={"timestamp": datetime.now().isoformat()},
            context_id=context_id
        )

# Singleton instance
_documentation_agent = None

def get_documentation_agent():
    """Returns a singleton instance of the documentation agent"""
    global _documentation_agent
    if _documentation_agent is None:
        _documentation_agent = DocumentationAgent()
    return _documentation_agent