69 lines
2.8 KiB
Python
69 lines
2.8 KiB
Python
import asyncio
|
|
import uuid
|
|
import json
|
|
from typing import Any, Dict, Union, List
|
|
from starlette.responses import Response
|
|
from tools.utils import ToolError
|
|
|
|
class MCPServer:
|
|
def __init__(self):
|
|
self.sessions: Dict[str, asyncio.Queue] = {}
|
|
self.running = True
|
|
|
|
def create_session(self) -> str:
|
|
sid = str(uuid.uuid4())
|
|
self.sessions[sid] = asyncio.Queue()
|
|
return sid
|
|
|
|
async def send_to_session(self, sid: str, message: Dict[str, Any]):
|
|
if sid in self.sessions:
|
|
await self.sessions[sid].put(json.dumps(message))
|
|
|
|
async def handle_request(self, method: str, params: Dict[str, Any], req_id: Any, tool_registry: list) -> Union[Dict[str, Any], Response]:
|
|
if method == "initialize":
|
|
return {
|
|
"protocolVersion": "2024-11-05",
|
|
"capabilities": {"tools": {}},
|
|
"serverInfo": {"name": "MooseCP", "version": "0.4.1"}
|
|
}
|
|
|
|
if method == "tools/list":
|
|
return {"tools": [t.to_mcp_dict() for t in tool_registry]}
|
|
|
|
if method == "tools/call":
|
|
name = params.get("name")
|
|
args = params.get("arguments", {})
|
|
|
|
tool = next((t for t in tool_registry if t.name == name), None)
|
|
if not tool:
|
|
return Response(
|
|
content=json.dumps({"jsonrpc": "2.0", "id": req_id, "error": {"code": -32601, "message": f"Tool {name} not found"}}),
|
|
media_type="application/json"
|
|
)
|
|
|
|
try:
|
|
result_data = await tool.handler(args)
|
|
except ToolError as e:
|
|
# Convert tool errors into a text response so the LLM can understand and potentially fix the input
|
|
result_data = {"text": str(e)}
|
|
except Exception as e:
|
|
# Catch-all for unexpected crashes to prevent server death
|
|
result_data = {"text": f"Unexpected internal error: {str(e)}"}
|
|
|
|
# Handle results that are already lists of content (for multimodal/multi-part responses)
|
|
if isinstance(result_data, list):
|
|
content = result_data
|
|
elif isinstance(result_data, dict) and "text" in result_data:
|
|
content = [{"type": "text", "text": result_data["text"]}]
|
|
elif isinstance(result_data, dict) and result_data.get("type") == "image":
|
|
content = [result_data]
|
|
else:
|
|
content = [result_data]
|
|
|
|
return {"content": content}
|
|
|
|
return Response(
|
|
content=json.dumps({"jsonrpc": "2.0", "id": req_id, "error": {"code": -32601, "message": f"Method {method} not found"}}),
|
|
media_type="application/json"
|
|
)
|