agent/tools/document_tool.py

"""
Document Generation Tool - Create Word documents
"""

import logging
import os
import tempfile
import base64
from typing import Any, Dict
from datetime import datetime
from docx import Document
from docx.shared import Inches, Pt, RGBColor
from docx.enum.text import WD_ALIGN_PARAGRAPH

logger = logging.getLogger(__name__)

DOCUMENT_TOOL_SPEC = {
    "name": "create_document",
    "description": (
        "Create a Word document (.docx) with formatted content. "
        "Use this to generate professional documents with headings, paragraphs, and styling. "
        "Supports reports, letters, proposals, and other document types."
    ),
    "parameters": {
        "type": "object",
        "properties": {
            "title": {
                "type": "string",
                "description": "The main title of the document",
            },
            "sections": {
                "type": "array",
                "description": "Array of section objects with 'heading' and 'paragraphs' (array of strings)",
                "items": {
                    "type": "object",
                    "properties": {
                        "heading": {"type": "string"},
                        "paragraphs": {"type": "array", "items": {"type": "string"}},
                    },
                    "required": ["heading", "paragraphs"],
                },
            },
            "filename": {
                "type": "string",
                "description": "Output filename for the document (e.g., 'document.docx')",
            },
        },
        "required": ["title", "sections", "filename"],
    },
}


async def create_document_handler(arguments: Dict[str, Any]) -> tuple[str, bool]:
    """Handler for creating Word documents."""
    try:
        title = arguments.get("title", "Document")
        sections = arguments.get("sections", [])
        filename = arguments.get("filename", "document.docx")
        
        if not sections:
            return "Error: No sections provided", False
        
        # Ensure filename has .docx extension
        if not filename.endswith(".docx"):
            filename += ".docx"
        
        # Create document
        doc = Document()
        
        # Set default font
        style = doc.styles['Normal']
        font = style.font
        font.name = 'Calibri'
        font.size = Pt(11)
        
        # Add title
        title_para = doc.add_paragraph()
        title_run = title_para.add_run(title)
        title_run.font.size = Pt(24)
        title_run.font.bold = True
        title_run.font.color.rgb = RGBColor(31, 78, 121)
        title_para.alignment = WD_ALIGN_PARAGRAPH.CENTER
        
        # Add spacing after title
        doc.add_paragraph()
        
        # Add sections
        for section in sections:
            heading = section.get("heading", "")
            paragraphs = section.get("paragraphs", [])
            
            # Add section heading
            if heading:
                heading_para = doc.add_paragraph()
                heading_run = heading_para.add_run(heading)
                heading_run.font.size = Pt(16)
                heading_run.font.bold = True
                heading_run.font.color.rgb = RGBColor(31, 78, 121)
                heading_para.space_after = Pt(6)
            
            # Add paragraphs
            for text in paragraphs:
                if text.strip():
                    p = doc.add_paragraph(text)
                    p.space_after = Pt(8)
            
            # Add spacing between sections
            doc.add_paragraph()
        
        # Save document
        output_dir = os.path.join(tempfile.gettempdir(), "kimi_agent_outputs")
        os.makedirs(output_dir, exist_ok=True)
        output_path = os.path.join(output_dir, filename)
        doc.save(output_path)
        
        # Also save to current working directory
        doc.save(filename)
        
        file_size = os.path.getsize(output_path)

        # Read file content for the event
        with open(output_path, "rb") as f:
            file_content = f.read()
        file_content_b64 = base64.b64encode(file_content).decode()

        # Create file data for file_generated event
        file_data = {
            "id": f"file_{datetime.now().strftime('%Y%m%d%H%M%S')}_{filename}",
            "name": filename,
            "type": "document",
            "content": file_content_b64,
            "language": None,
            "originalContent": None,
            "modifiedContent": file_content_b64,
            "size": file_size,
            "created_at": datetime.now().isoformat(),
        }

        return (
            f"✅ Document created successfully!\n"
            f"📄 Filename: {filename}\n"
            f"📑 Total sections: {len(sections)}\n"
            f"💾 File size: {file_size / 1024:.2f} KB\n"
            f"📁 Saved to: {output_path}",
            True,
            file_data  # Return file data for event
        )
        
    except Exception as e:
        logger.error(f"Document creation error: {e}")
        return f"❌ Error creating document: {str(e)}", False