Spaces:

onewayto
/

water3

Sleeping

water3 / agent /tools /document_tool.py

Upload 187 files

070daf8 verified 8 months ago

5.13 kB

"""
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

	"""
	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