gateway/modules/agentDocumentation.py

"""
Documentation agent for creating documentation, reports, and structured content.
Reimagined with an output-first, AI-driven approach with multi-step document generation.
"""

import logging
import json
from typing import Dict, Any, List

from modules.workflowAgentsRegistry import AgentBase

logger = logging.getLogger(__name__)

class AgentDocumentation(AgentBase):
    """AI-driven agent for creating documentation and structured content using multi-step generation"""
    
    def __init__(self):
        """Initialize the documentation agent"""
        super().__init__()
        self.name = "documentation"
        self.description = "Creates structured documentation, reports, and content using AI with multi-step generation"
        self.capabilities = [
            "report_generation",
            "documentation",
            "content_structuring",
            "technical_writing",
            "knowledge_organization"
        ]

    def setDependencies(self, mydom=None):
        """Set external dependencies for the agent."""
        self.mydom = mydom
        
    async def processTask(self, task: Dict[str, Any]) -> Dict[str, Any]:
        """
        Process a task by focusing on required outputs and using AI to generate them.
        
        Args:
            task: Task dictionary with prompt, inputDocuments, outputSpecifications
            
        Returns:
            Dictionary with feedback and documents
        """
        try:
            # Extract task information
            prompt = task.get("prompt", "")
            inputDocuments = task.get("inputDocuments", [])
            outputSpecs = task.get("outputSpecifications", [])
            
            # Check AI service
            if not self.mydom:
                return {
                    "feedback": "The Documentation agent requires an AI service to function.",
                    "documents": []
                }
            
            # Extract context from input documents - focusing only on dataExtracted
            documentContext = self._extractDocumentContext(inputDocuments)
            
            # Create task analysis to understand the requirements
            documentationPlan = await self._analyzeTask(prompt, documentContext, outputSpecs)
            
            # Generate all required output documents
            documents = []
            
            # If no output specs provided, create default document
            if not outputSpecs:
                defaultFormat = documentationPlan.get("recommendedFormat", "markdown")
                defaultTitle = documentationPlan.get("title", "Documentation")
                safeTitle = self._sanitizeFilename(defaultTitle)
                
                outputSpecs = [
                    {"label": f"{safeTitle}.{defaultFormat}", "description": "Comprehensive documentation"}
                ]
            
            # Process each output specification
            for spec in outputSpecs:
                outputLabel = spec.get("label", "")
                outputDescription = spec.get("description", "")
                
                # Generate the document using multi-step approach
                document = await self._createDocumentMultiStep(
                    prompt, 
                    documentContext, 
                    outputLabel,
                    outputDescription,
                    documentationPlan
                )
                
                documents.append(document)
            
            # Generate feedback
            feedback = documentationPlan.get("feedback", f"Created {len(documents)} documents based on your requirements.")
            
            return {
                "feedback": feedback,
                "documents": documents
            }
            
        except Exception as e:
            logger.error(f"Error in documentation generation: {str(e)}", exc_info=True)
            return {
                "feedback": f"Error during documentation generation: {str(e)}",
                "documents": []
            }
    
    def _extractDocumentContext(self, documents: List[Dict[str, Any]]) -> str:
        """
        Extract context from input documents, focusing on dataExtracted.
        
        Args:
            documents: List of document objects
            
        Returns:
            Extracted context as text
        """
        contextParts = []
        
        for doc in documents:
            docName = doc.get("name", "unnamed")
            if doc.get("ext"):
                docName = f"{docName}.{doc.get('ext')}"
            
            contextParts.append(f"\n\n--- {docName} ---\n")
            
            # Process contents for dataExtracted
            for content in doc.get("contents", []):
                if content.get("dataExtracted"):
                    contextParts.append(content.get("dataExtracted", ""))
        
        return "\n".join(contextParts)
    
    def _sanitizeFilename(self, filename: str) -> str:
        """
        Sanitize a filename by removing invalid characters.
        
        Args:
            filename: Filename to sanitize
            
        Returns:
            Sanitized filename
        """
        # Replace invalid characters with underscores
        invalidChars = r'<>:"/\|?*'
        for char in invalidChars:
            filename = filename.replace(char, '_')
        
        # Trim filename if too long
        if len(filename) > 100:
            filename = filename[:97] + "..."
            
        return filename
    
    async def _analyzeTask(self, prompt: str, context: str, outputSpecs: List) -> Dict:
        """
        Use AI to analyze the task and create a documentation plan.
        
        Args:
            prompt: The task prompt
            context: Document context
            outputSpecs: Output specifications
            
        Returns:
            Documentation plan dictionary
        """
        analysisPrompt = f"""
        Analyze this documentation task and create a detailed plan.
        
        TASK: {prompt}
        
        DOCUMENT CONTEXT SAMPLE:
        {context[:1000]}... (truncated)
        
        OUTPUT REQUIREMENTS:
        {json.dumps(outputSpecs, indent=2)}
        
        Create a detailed documentation plan in JSON format with the following structure:
        {{
            "title": "Document Title",
            "documentType": "report|manual|guide|whitepaper|etc",
            "audience": "technical|general|executive|etc",
            "detailedStructure": [
                {{
                    "title": "Chapter/Section Title",
                    "keyPoints": ["point1", "point2", ...],
                    "subsections": ["subsection1", "subsection2", ...],
                    "importance": "high|medium|low",
                    "estimatedLength": "short|medium|long"
                }},
                ... more sections ...
            ],
            "keyTopics": ["topic1", "topic2", ...],
            "tone": "formal|conversational|instructional|etc",
            "recommendedFormat": "markdown|html|text|etc",
            "formattingRequirements": ["requirement1", "requirement2", ...],
            "executiveSummary": "Brief description of what the document will cover",
            "feedback": "Brief message explaining the documentation approach"
        }}
        
        Only return valid JSON. No preamble or explanations.
        """
        
        try:
            response = await self.mydom.callAi([
                {"role": "system", "content": "You are a documentation expert. Respond with valid JSON only."},
                {"role": "user", "content": analysisPrompt}
            ])
            
            # Extract JSON from response
            jsonStart = response.find('{')
            jsonEnd = response.rfind('}') + 1
            
            if jsonStart >= 0 and jsonEnd > jsonStart:
                plan = json.loads(response[jsonStart:jsonEnd])
                return plan
            else:
                # Fallback if JSON not found
                return {
                    "title": "Documentation",
                    "documentType": "report",
                    "audience": "general",
                    "detailedStructure": [
                        {
                            "title": "Introduction",
                            "keyPoints": ["Purpose", "Scope"],
                            "subsections": [],
                            "importance": "high",
                            "estimatedLength": "short"
                        },
                        {
                            "title": "Main Content",
                            "keyPoints": ["Core Information"],
                            "subsections": ["Key Findings", "Analysis"],
                            "importance": "high",
                            "estimatedLength": "long"
                        },
                        {
                            "title": "Conclusion",
                            "keyPoints": ["Summary", "Next Steps"],
                            "subsections": [],
                            "importance": "medium",
                            "estimatedLength": "short"
                        }
                    ],
                    "keyTopics": ["General Information"],
                    "tone": "formal",
                    "recommendedFormat": "markdown",
                    "formattingRequirements": ["Clear headings", "Professional formatting"],
                    "executiveSummary": "A comprehensive documentation covering the requested topics.",
                    "feedback": "Created documentation based on your requirements."
                }
                
        except Exception as e:
            logger.warning(f"Error creating documentation plan: {str(e)}")
            return {
                "title": "Documentation",
                "documentType": "report",
                "audience": "general",
                "detailedStructure": [
                    {
                        "title": "Introduction",
                        "keyPoints": ["Purpose", "Scope"],
                        "subsections": [],
                        "importance": "high",
                        "estimatedLength": "short"
                    },
                    {
                        "title": "Main Content",
                        "keyPoints": ["Core Information"],
                        "subsections": ["Key Findings", "Analysis"],
                        "importance": "high",
                        "estimatedLength": "long"
                    },
                    {
                        "title": "Conclusion",
                        "keyPoints": ["Summary", "Next Steps"],
                        "subsections": [],
                        "importance": "medium",
                        "estimatedLength": "short"
                    }
                ],
                "keyTopics": ["General Information"],
                "tone": "formal",
                "recommendedFormat": "markdown",
                "formattingRequirements": ["Clear headings", "Professional formatting"],
                "executiveSummary": "A comprehensive documentation covering the requested topics.",
                "feedback": "Created documentation based on your requirements."
            }
    
    async def _createDocumentMultiStep(self, prompt: str, context: str, outputLabel: str, 
                                     outputDescription: str, documentationPlan: Dict) -> Dict:
        """
        Create a document using a multi-step approach with separate AI calls for each section.
        
        Args:
            prompt: Original task prompt
            context: Document context
            outputLabel: Output filename
            outputDescription: Description of desired output
            documentationPlan: Documentation plan from AI
            
        Returns:
            Document object
        """
        # Determine format from filename
        formatType = outputLabel.split('.')[-1].lower() if '.' in outputLabel else "md"
        
        # Map format to contentType
        contentTypeMap = {
            "md": "text/markdown",
            "markdown": "text/markdown",
            "html": "text/html",
            "txt": "text/plain",
            "text": "text/plain",
            "json": "application/json",
            "csv": "text/csv"
        }
        
        contentType = contentTypeMap.get(formatType, "text/plain")
        
        # Get document information
        title = documentationPlan.get("title", "Documentation")
        documentType = documentationPlan.get("documentType", "document")
        audience = documentationPlan.get("audience", "general")
        tone = documentationPlan.get("tone", "formal")
        keyTopics = documentationPlan.get("keyTopics", [])
        formattingRequirements = documentationPlan.get("formattingRequirements", [])
        
        # Get the detailed structure
        detailedStructure = documentationPlan.get("detailedStructure", [])
        if not detailedStructure:
            # Fallback structure if none provided
            detailedStructure = [
                {
                    "title": "Introduction",
                    "keyPoints": ["Purpose", "Scope"],
                    "importance": "high"
                },
                {
                    "title": "Main Content",
                    "keyPoints": ["Core Information"],
                    "importance": "high"
                },
                {
                    "title": "Conclusion",
                    "keyPoints": ["Summary", "Next Steps"],
                    "importance": "medium"
                }
            ]
        
        try:
            # Step 1: Generate document introduction
            introPrompt = f"""
            Create the introduction for a {documentType} titled "{title}".
            
            DOCUMENT OVERVIEW:
            - Type: {documentType}
            - Audience: {audience}
            - Tone: {tone}
            - Key Topics: {', '.join(keyTopics)}
            - Format: {formatType}
            
            TASK CONTEXT: {prompt}
            
            This introduction should:
            1. Clearly state the purpose and scope of the document
            2. Provide context and background information
            3. Outline what the reader will find in the document
            4. Set the appropriate tone for the {audience} audience
            
            The introduction should be professional and engaging, formatted according to {formatType} standards.
            """
            
            introduction = await self.mydom.callAi([
                {"role": "system", "content": f"You are a documentation expert creating an introduction in {formatType} format."},
                {"role": "user", "content": introPrompt}
            ], produceUserAnswer = True)
            
            # Step 2: Generate executive summary (if applicable)
            if documentType in ["report", "whitepaper", "case study"]:
                summaryPrompt = f"""
                Create an executive summary for a {documentType} titled "{title}".
                
                DOCUMENT OVERVIEW:
                - Type: {documentType}
                - Audience: {audience}
                - Key Topics: {', '.join(keyTopics)}
                
                TASK CONTEXT: {prompt}
                
                This executive summary should:
                1. Provide a concise overview of the entire document
                2. Highlight key findings, recommendations, or conclusions
                3. Be suitable for executives or busy readers who may only read this section
                4. Be professionally formatted according to {formatType} standards
                
                Keep the summary focused and impactful, approximately 200-300 words.
                """
                
                executiveSummary = await self.mydom.callAi([
                    {"role": "system", "content": f"You are a documentation expert creating an executive summary in {formatType} format."},
                    {"role": "user", "content": summaryPrompt}
                ], produceUserAnswer = True)
            else:
                executiveSummary = ""
            
            # Step 3: Generate each section
            sections = []
            
            for section in detailedStructure:
                sectionTitle = section.get("title", "Section")
                keyPoints = section.get("keyPoints", [])
                subsections = section.get("subsections", [])
                importance = section.get("importance", "medium")
                
                # Adjust depth based on importance
                detailLevel = "high" if importance == "high" else "medium"
                
                sectionPrompt = f"""
                Create the "{sectionTitle}" section for a {documentType} titled "{title}".
                
                SECTION DETAILS:
                - Title: {sectionTitle}
                - Key Points to Cover: {', '.join(keyPoints)}
                - Subsections: {', '.join(subsections)}
                - Detail Level: {detailLevel}
                
                DOCUMENT CONTEXT:
                - Type: {documentType}
                - Audience: {audience}
                - Tone: {tone}
                - Format: {formatType}
                
                TASK CONTEXT: {prompt}
                
                AVAILABLE INFORMATION:
                {context[:500]}... (truncated)
                
                This section should:
                1. Be comprehensive and well-structured
                2. Cover all the key points listed
                3. Include the specified subsections with appropriate headings
                4. Maintain a {tone} tone suitable for the {audience} audience
                5. Be properly formatted according to {formatType} standards
                6. Include specific examples, data, or evidence where appropriate
                
                Be thorough in your coverage of this section, providing substantive content.
                """
                
                sectionContent = await self.mydom.callAi([
                    {"role": "system", "content": f"You are a documentation expert creating detailed content for the {sectionTitle} section."},
                    {"role": "user", "content": sectionPrompt}
                ], produceUserAnswer = True)
                
                sections.append(sectionContent)
            
            # Step 4: Generate conclusion
            conclusionPrompt = f"""
            Create the conclusion for a {documentType} titled "{title}".
            
            DOCUMENT OVERVIEW:
            - Type: {documentType}
            - Audience: {audience}
            - Key Topics: {', '.join(keyTopics)}
            
            TASK CONTEXT: {prompt}
            
            This conclusion should:
            1. Summarize the key points covered in the document
            2. Provide closure to the topics discussed
            3. Include any relevant recommendations or next steps
            4. Leave the reader with a clear understanding of the document's significance
            
            The conclusion should be professional and impactful, formatted according to {formatType} standards.
            """
            
            conclusion = await self.mydom.callAi([
                {"role": "system", "content": f"You are a documentation expert creating a conclusion in {formatType} format."},
                {"role": "user", "content": conclusionPrompt}
            ], produceUserAnswer = True)
            
            # Step 5: Assemble the complete document
            if formatType in ["md", "markdown"]:
                # Markdown format
                documentContent = f"# {title}\n\n"
                
                if executiveSummary:
                    documentContent += f"## Executive Summary\n\n{executiveSummary}\n\n"
                
                documentContent += f"{introduction}\n\n"
                
                for i, sectionContent in enumerate(sections):
                    # Ensure section starts with heading if not already
                    sectionTitle = detailedStructure[i].get("title", f"Section {i+1}")
                    if not sectionContent.strip().startswith("#"):
                        documentContent += f"## {sectionTitle}\n\n"
                    documentContent += f"{sectionContent}\n\n"
                
                documentContent += f"## Conclusion\n\n{conclusion}\n"
                
            elif formatType == "html":
                # HTML format
                documentContent = f"<html>\n<head>\n<title>{title}</title>\n</head>\n<body>\n"
                documentContent += f"<h1>{title}</h1>\n\n"
                
                if executiveSummary:
                    documentContent += f"<h2>Executive Summary</h2>\n<div>{executiveSummary}</div>\n\n"
                
                documentContent += f"<div>{introduction}</div>\n\n"
                
                for i, sectionContent in enumerate(sections):
                    sectionTitle = detailedStructure[i].get("title", f"Section {i+1}")
                    documentContent += f"<h2>{sectionTitle}</h2>\n<div>{sectionContent}</div>\n\n"
                
                documentContent += f"<h2>Conclusion</h2>\n<div>{conclusion}</div>\n"
                documentContent += "</body>\n</html>"
                
            else:
                # Plain text format
                documentContent = f"{title}\n{'=' * len(title)}\n\n"
                
                if executiveSummary:
                    documentContent += f"EXECUTIVE SUMMARY\n{'-' * 17}\n\n{executiveSummary}\n\n"
                
                documentContent += f"{introduction}\n\n"
                
                for i, sectionContent in enumerate(sections):
                    sectionTitle = detailedStructure[i].get("title", f"Section {i+1}")
                    documentContent += f"{sectionTitle}\n{'-' * len(sectionTitle)}\n\n{sectionContent}\n\n"
                
                documentContent += f"CONCLUSION\n{'-' * 10}\n\n{conclusion}\n"
            
            # Create document object
            return {
                "label": outputLabel,
                "content": documentContent,
                "metadata": {
                    "contentType": contentType
                }
            }
            
        except Exception as e:
            logger.error(f"Error creating document: {str(e)}", exc_info=True)
            
            # Create a simple error document
            if formatType in ["md", "markdown"]:
                content = f"# Error in Documentation\n\nThere was an error generating the documentation: {str(e)}"
            elif formatType == "html":
                content = f"<html><body><h1>Error in Documentation</h1><p>There was an error generating the documentation: {str(e)}</p></body></html>"
            else:
                content = f"Error in Documentation\n\nThere was an error generating the documentation: {str(e)}"
            
            return {
                "label": outputLabel,
                "content": content,
                "metadata": {
                    "contentType": contentType
                }
            }


# Factory function for the Documentation agent
def getAgentDocumentation():
    """Returns an instance of the Documentation agent."""
    return AgentDocumentation()