office_translator/mcp_server.py

#!/usr/bin/env python3
"""
MCP Server for Document Translation API
Model Context Protocol server for AI assistant integration
"""

import sys
import json
import asyncio
import base64
import requests
from pathlib import Path
from typing import Any, Optional

# MCP Protocol Constants
JSONRPC_VERSION = "2.0"

class MCPServer:
    """MCP Server for Document Translation"""
    
    def __init__(self):
        self.api_base = "http://localhost:8000"
        self.capabilities = {
            "tools": {}
        }
        
    def get_tools(self) -> list:
        """Return list of available tools"""
        return [
            {
                "name": "translate_document",
                "description": "Translate a document (Excel, Word, PowerPoint) to another language while preserving formatting",
                "inputSchema": {
                    "type": "object",
                    "properties": {
                        "file_path": {
                            "type": "string",
                            "description": "Path to the document file (.xlsx, .docx, .pptx)"
                        },
                        "target_language": {
                            "type": "string",
                            "description": "Target language code (e.g., 'en', 'fr', 'es', 'fa', 'de')"
                        },
                        "provider": {
                            "type": "string",
                            "enum": ["google", "google_cloud", "openai", "openrouter", "deepseek", "minimax", "zai"],
                            "description": "Translation provider (default: google)"
                        },
                        "translate_images": {
                            "type": "boolean",
                            "description": "Extract and translate text from images using vision model"
                        },
                        "system_prompt": {
                            "type": "string",
                            "description": "Custom system prompt with context, glossary, or instructions for LLM translation"
                        },
                        "output_path": {
                            "type": "string",
                            "description": "Path where to save the translated document (optional)"
                        }
                    },
                    "required": ["file_path", "target_language"]
                }
            },

            {
                "name": "get_supported_languages",
                "description": "Get list of supported language codes for translation",
                "inputSchema": {
                    "type": "object",
                    "properties": {}
                }
            },
            {
                "name": "configure_translation",
                "description": "Configure translation settings",
                "inputSchema": {
                    "type": "object",
                    "properties": {
                        "provider": {
                            "type": "string",
                            "enum": ["google", "google_cloud", "openai", "openrouter", "deepseek", "minimax", "zai"],
                            "description": "Default translation provider"
                        }
                    }
                }
            },
            {
                "name": "check_api_health",
                "description": "Check if the translation API is running and healthy",
                "inputSchema": {
                    "type": "object",
                    "properties": {}
                }
            }
        ]
    
    async def handle_tool_call(self, name: str, arguments: dict) -> dict:
        """Handle tool calls"""
        try:
            if name == "translate_document":
                return await self.translate_document(arguments)

            elif name == "get_supported_languages":
                return await self.get_supported_languages()
            elif name == "configure_translation":
                return await self.configure_translation(arguments)
            elif name == "check_api_health":
                return await self.check_api_health()
            else:
                return {"error": f"Unknown tool: {name}"}
        except Exception as e:
            return {"error": str(e)}
    
    async def translate_document(self, args: dict) -> dict:
        """Translate a document file"""
        file_path = Path(args["file_path"])
        
        if not file_path.exists():
            return {"error": f"File not found: {file_path}"}
        
        # Prepare form data
        with open(file_path, 'rb') as f:
            files = {'file': (file_path.name, f)}
            data = {
                'target_language': args['target_language'],
                'provider': args.get('provider', 'google'),
                'translate_images': str(args.get('translate_images', False)).lower(),
            }
            

            
            if args.get('system_prompt'):
                data['system_prompt'] = args['system_prompt']
            
            try:
                response = requests.post(
                    f"{self.api_base}/translate",
                    files=files,
                    data=data,
                    timeout=300  # 5 minutes timeout
                )
                
                if response.status_code == 200:
                    # Save translated file
                    output_path = args.get('output_path')
                    if not output_path:
                        output_path = file_path.parent / f"translated_{file_path.name}"
                    
                    output_path = Path(output_path)
                    with open(output_path, 'wb') as out:
                        out.write(response.content)
                    
                    return {
                        "success": True,
                        "message": f"Document translated successfully",
                        "output_path": str(output_path),
                        "source_file": str(file_path),
                        "target_language": args['target_language'],
                        "provider": args.get('provider', 'google')
                    }
                else:
                    error_detail = response.json() if response.headers.get('content-type') == 'application/json' else response.text
                    return {"error": f"Translation failed: {error_detail}"}
                    
            except requests.exceptions.ConnectionError:
                return {"error": "Cannot connect to translation API. Make sure the server is running on http://localhost:8000"}
            except requests.exceptions.Timeout:
                return {"error": "Translation request timed out"}
    

    
    async def get_supported_languages(self) -> dict:
        """Get supported language codes"""
        return {
            "languages": [
                {"code": "en", "name": "English"},
                {"code": "fa", "name": "Persian/Farsi"},
                {"code": "fr", "name": "French"},
                {"code": "es", "name": "Spanish"},
                {"code": "de", "name": "German"},
                {"code": "it", "name": "Italian"},
                {"code": "pt", "name": "Portuguese"},
                {"code": "ru", "name": "Russian"},
                {"code": "zh", "name": "Chinese"},
                {"code": "ja", "name": "Japanese"},
                {"code": "ko", "name": "Korean"},
                {"code": "ar", "name": "Arabic"},
                {"code": "nl", "name": "Dutch"},
                {"code": "pl", "name": "Polish"},
                {"code": "tr", "name": "Turkish"},
                {"code": "vi", "name": "Vietnamese"},
                {"code": "th", "name": "Thai"},
                {"code": "hi", "name": "Hindi"},
                {"code": "he", "name": "Hebrew"},
                {"code": "sv", "name": "Swedish"}
            ]
        }
    
    async def configure_translation(self, args: dict) -> dict:
        """Configure translation settings"""
        config = {}
        

        
        config['provider'] = args.get('provider', 'google')
        
        return {
            "success": True,
            "configuration": config
        }
    
    async def check_api_health(self) -> dict:
        """Check API health status"""
        try:
            response = requests.get(f"{self.api_base}/health", timeout=5)
            
            if response.status_code == 200:
                return {
                    "status": "healthy",
                    "api_url": self.api_base,
                    "details": response.json()
                }
            else:
                return {"status": "unhealthy", "error": "API returned non-200 status"}
                
        except requests.exceptions.ConnectionError:
            return {
                "status": "unavailable",
                "error": "Cannot connect to API. Start the server with: python main.py"
            }
    
    def create_response(self, id: Any, result: Any) -> dict:
        """Create JSON-RPC response"""
        return {
            "jsonrpc": JSONRPC_VERSION,
            "id": id,
            "result": result
        }
    
    def create_error(self, id: Any, code: int, message: str) -> dict:
        """Create JSON-RPC error response"""
        return {
            "jsonrpc": JSONRPC_VERSION,
            "id": id,
            "error": {
                "code": code,
                "message": message
            }
        }
    
    async def handle_message(self, message: dict) -> Optional[dict]:
        """Handle incoming JSON-RPC message"""
        msg_id = message.get("id")
        method = message.get("method")
        params = message.get("params", {})
        
        if method == "initialize":
            return self.create_response(msg_id, {
                "protocolVersion": "2024-11-05",
                "capabilities": self.capabilities,
                "serverInfo": {
                    "name": "document-translator",
                    "version": "1.0.0"
                }
            })
        
        elif method == "notifications/initialized":
            return None  # No response needed for notifications
        
        elif method == "tools/list":
            return self.create_response(msg_id, {
                "tools": self.get_tools()
            })
        
        elif method == "tools/call":
            tool_name = params.get("name")
            tool_args = params.get("arguments", {})
            result = await self.handle_tool_call(tool_name, tool_args)
            
            return self.create_response(msg_id, {
                "content": [
                    {
                        "type": "text",
                        "text": json.dumps(result, indent=2, ensure_ascii=False)
                    }
                ]
            })
        
        elif method == "ping":
            return self.create_response(msg_id, {})
        
        else:
            return self.create_error(msg_id, -32601, f"Method not found: {method}")
    
    async def run(self):
        """Run the MCP server using stdio"""
        while True:
            try:
                line = sys.stdin.readline()
                if not line:
                    break
                
                message = json.loads(line)
                response = await self.handle_message(message)
                
                if response:
                    sys.stdout.write(json.dumps(response) + "\n")
                    sys.stdout.flush()
                    
            except json.JSONDecodeError as e:
                error = self.create_error(None, -32700, f"Parse error: {e}")
                sys.stdout.write(json.dumps(error) + "\n")
                sys.stdout.flush()
            except Exception as e:
                error = self.create_error(None, -32603, f"Internal error: {e}")
                sys.stdout.write(json.dumps(error) + "\n")
                sys.stdout.flush()


def main():
    """Main entry point"""
    server = MCPServer()
    asyncio.run(server.run())


if __name__ == "__main__":
    main()