feat: add MCP tool integration
This commit is contained in:
Xin Wang
5
backend/services/tools/__init__.py
Normal file
5
backend/services/tools/__init__.py
Normal file
@@ -0,0 +1,5 @@
|
||||
"""Small protocol adapters used by the shared ToolExecutor."""
|
||||
|
||||
from services.tools.mcp_client import McpClientError, McpToolClient
|
||||
|
||||
__all__ = ["McpClientError", "McpToolClient"]
|
||||
152
backend/services/tools/mcp_client.py
Normal file
152
backend/services/tools/mcp_client.py
Normal file
@@ -0,0 +1,152 @@
|
||||
"""One-shot Streamable HTTP MCP discovery and execution.
|
||||
|
||||
Connections intentionally live for one operation in the MVP. This keeps the
|
||||
async context ownership simple and correct; a session pool can be introduced
|
||||
later if measurements show MCP handshakes are a meaningful source of latency.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import hashlib
|
||||
import json
|
||||
from contextlib import asynccontextmanager
|
||||
from datetime import timedelta
|
||||
from typing import Any, AsyncIterator
|
||||
|
||||
import httpx
|
||||
from mcp import ClientSession
|
||||
from mcp.client.streamable_http import streamable_http_client
|
||||
|
||||
from models import RuntimeMcpServer
|
||||
|
||||
|
||||
MAX_MCP_RESULT_BYTES = 1_000_000
|
||||
|
||||
|
||||
class McpClientError(RuntimeError):
|
||||
"""Safe error surfaced to tool callers and the administration API."""
|
||||
|
||||
|
||||
def schema_hash(input_schema: dict[str, Any]) -> str:
|
||||
serialized = json.dumps(
|
||||
input_schema,
|
||||
ensure_ascii=False,
|
||||
sort_keys=True,
|
||||
separators=(",", ":"),
|
||||
)
|
||||
return hashlib.sha256(serialized.encode("utf-8")).hexdigest()
|
||||
|
||||
|
||||
class McpToolClient:
|
||||
"""Minimal public MCP SDK wrapper shared by discovery and runtime calls."""
|
||||
|
||||
@asynccontextmanager
|
||||
async def _session(
|
||||
self,
|
||||
server: RuntimeMcpServer,
|
||||
) -> AsyncIterator[ClientSession]:
|
||||
if server.transport != "streamable_http":
|
||||
raise McpClientError(f"暂不支持 MCP transport: {server.transport}")
|
||||
|
||||
headers = {**server.headers, **server.secret_headers}
|
||||
timeout = max(1, min(int(server.timeout_seconds), 120))
|
||||
try:
|
||||
async with httpx.AsyncClient(
|
||||
headers=headers,
|
||||
timeout=httpx.Timeout(timeout),
|
||||
follow_redirects=False,
|
||||
) as http_client:
|
||||
async with streamable_http_client(
|
||||
server.url,
|
||||
http_client=http_client,
|
||||
) as streams:
|
||||
read_stream, write_stream, _ = streams
|
||||
async with ClientSession(
|
||||
read_stream,
|
||||
write_stream,
|
||||
read_timeout_seconds=timedelta(seconds=timeout),
|
||||
) as session:
|
||||
await session.initialize()
|
||||
yield session
|
||||
except McpClientError:
|
||||
raise
|
||||
except httpx.TimeoutException as exc:
|
||||
raise McpClientError("MCP Server 连接或调用超时") from exc
|
||||
except Exception as exc: # MCP SDK exposes several transport exceptions.
|
||||
raise McpClientError(f"MCP Server 调用失败: {exc}") from exc
|
||||
|
||||
async def list_tools(self, server: RuntimeMcpServer) -> list[dict[str, Any]]:
|
||||
"""Discover all remote tools, following MCP pagination cursors."""
|
||||
discovered: list[dict[str, Any]] = []
|
||||
async with self._session(server) as session:
|
||||
cursor: str | None = None
|
||||
while True:
|
||||
result = await session.list_tools(cursor=cursor)
|
||||
for tool in result.tools:
|
||||
input_schema = dict(tool.inputSchema or {})
|
||||
discovered.append(
|
||||
{
|
||||
"name": tool.name,
|
||||
"description": tool.description or "",
|
||||
"input_schema": input_schema,
|
||||
"schema_hash": schema_hash(input_schema),
|
||||
}
|
||||
)
|
||||
cursor = result.nextCursor
|
||||
if not cursor:
|
||||
break
|
||||
return discovered
|
||||
|
||||
async def call_tool(
|
||||
self,
|
||||
server: RuntimeMcpServer,
|
||||
remote_tool_name: str,
|
||||
arguments: dict[str, Any],
|
||||
) -> dict[str, Any]:
|
||||
"""Execute one remote tool and normalize the result for local callers."""
|
||||
async with self._session(server) as session:
|
||||
result = await session.call_tool(
|
||||
remote_tool_name,
|
||||
arguments=arguments,
|
||||
read_timeout_seconds=timedelta(seconds=server.timeout_seconds),
|
||||
)
|
||||
|
||||
data = self._normalize_result(result)
|
||||
encoded = json.dumps(data, ensure_ascii=False, default=str).encode("utf-8")
|
||||
if len(encoded) > MAX_MCP_RESULT_BYTES:
|
||||
raise McpClientError("MCP 工具响应超过 1 MB 限制")
|
||||
if bool(result.isError):
|
||||
message = data.get("text") or "MCP 工具返回执行错误"
|
||||
raise McpClientError(str(message)[:2048])
|
||||
return data
|
||||
|
||||
@staticmethod
|
||||
def _normalize_result(result: Any) -> dict[str, Any]:
|
||||
"""Keep useful text/JSON while deliberately omitting binary payloads."""
|
||||
content: list[dict[str, Any]] = []
|
||||
text_parts: list[str] = []
|
||||
for item in result.content or []:
|
||||
item_type = str(getattr(item, "type", "unknown"))
|
||||
if item_type == "text":
|
||||
text = str(getattr(item, "text", ""))
|
||||
content.append({"type": "text", "text": text})
|
||||
if text:
|
||||
text_parts.append(text)
|
||||
continue
|
||||
if item_type == "resource":
|
||||
resource = getattr(item, "resource", None)
|
||||
text = getattr(resource, "text", None)
|
||||
if isinstance(text, str):
|
||||
content.append({"type": "resource", "text": text})
|
||||
text_parts.append(text)
|
||||
else:
|
||||
content.append({"type": "resource", "omitted": True})
|
||||
continue
|
||||
# Audio and image base64 data are not sent into the LLM context in MVP.
|
||||
content.append({"type": item_type, "omitted": True})
|
||||
|
||||
return {
|
||||
"text": "\n".join(text_parts),
|
||||
"structuredContent": result.structuredContent,
|
||||
"content": content,
|
||||
}
|
||||
Reference in New Issue
Block a user