office_translator/mcp_server.py

392 lines
15 KiB
Python

#!/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", "ollama", "deepl", "libre"],
"description": "Translation provider (default: google)"
},
"ollama_model": {
"type": "string",
"description": "Ollama model to use (e.g., 'llama3.2', 'gemma3:12b')"
},
"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": "list_ollama_models",
"description": "List available Ollama models for translation",
"inputSchema": {
"type": "object",
"properties": {
"base_url": {
"type": "string",
"description": "Ollama server URL (default: http://localhost:11434)"
}
}
}
},
{
"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", "ollama", "deepl", "libre"],
"description": "Default translation provider"
},
"ollama_url": {
"type": "string",
"description": "Ollama server URL"
},
"ollama_model": {
"type": "string",
"description": "Default Ollama model"
}
}
}
},
{
"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 == "list_ollama_models":
return await self.list_ollama_models(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('ollama_model'):
data['ollama_model'] = args['ollama_model']
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 list_ollama_models(self, args: dict) -> dict:
"""List available Ollama models"""
base_url = args.get('base_url', 'http://localhost:11434')
try:
response = requests.get(
f"{self.api_base}/ollama/models",
params={'base_url': base_url},
timeout=10
)
if response.status_code == 200:
data = response.json()
return {
"models": data.get('models', []),
"count": data.get('count', 0),
"ollama_url": base_url
}
else:
return {"error": "Failed to list models", "models": []}
except requests.exceptions.ConnectionError:
return {"error": "Cannot connect to API server", "models": []}
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 = {}
if args.get('ollama_url') and args.get('ollama_model'):
try:
response = requests.post(
f"{self.api_base}/ollama/configure",
data={
'base_url': args['ollama_url'],
'model': args['ollama_model']
},
timeout=10
)
if response.status_code == 200:
config['ollama'] = response.json()
except Exception as e:
config['ollama_error'] = str(e)
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()