Add MCP server and configuration for AI assistant integration
This commit is contained in:
391
mcp_server.py
Normal file
391
mcp_server.py
Normal file
@@ -0,0 +1,391 @@
|
||||
#!/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()
|
||||
Reference in New Issue
Block a user