266 lines
9.3 KiB
Python
266 lines
9.3 KiB
Python
"""
|
|
MCP 客户端桥接
|
|
==============
|
|
连接外部 MCP Server，把它们的 tools 提供给自动回复链路调用，
|
|
让客服 AI 能查文档/挂号/业务系统等，回复更精准。
|
|
|
|
服务器列表配置在 ai_config.AI_MCP_SERVERS，例如：
|
|
|
|
[
|
|
{
|
|
"name": "filesystem",
|
|
"enabled": true,
|
|
"transport": "stdio",
|
|
"command": "npx",
|
|
"args": ["-y", "@modelcontextprotocol/server-filesystem", "D:/docs"],
|
|
"env": {},
|
|
"cwd": null
|
|
}
|
|
]
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import json
|
|
import os
|
|
import re
|
|
from contextlib import AsyncExitStack
|
|
from typing import Any, Optional
|
|
|
|
import ai_config
|
|
|
|
# OpenAI / DeepSeek function name 只允许 [a-zA-Z0-9_-]
|
|
_SAFE = re.compile(r"[^a-zA-Z0-9_-]+")
|
|
|
|
|
|
def _safe_name(server: str, tool: str) -> str:
|
|
s = _SAFE.sub("_", server.strip()) or "srv"
|
|
t = _SAFE.sub("_", tool.strip()) or "tool"
|
|
return f"{s}__{t}"
|
|
|
|
|
|
def _parse_name(namespaced: str) -> tuple[str, str]:
|
|
if "__" not in namespaced:
|
|
raise ValueError(f"invalid tool name (expect server__tool): {namespaced}")
|
|
server, tool = namespaced.split("__", 1)
|
|
return server, tool
|
|
|
|
|
|
def _enabled_servers() -> list[dict]:
|
|
raw = getattr(ai_config, "AI_MCP_SERVERS", None) or []
|
|
out = []
|
|
for i, s in enumerate(raw):
|
|
if not isinstance(s, dict):
|
|
continue
|
|
if s.get("enabled", True) is False:
|
|
continue
|
|
name = (s.get("name") or f"server{i}").strip()
|
|
transport = (s.get("transport") or "stdio").strip().lower()
|
|
if transport == "stdio" and not s.get("command"):
|
|
continue
|
|
if transport in ("sse", "http", "streamable_http") and not s.get("url"):
|
|
continue
|
|
cfg = dict(s)
|
|
cfg["name"] = name
|
|
cfg["transport"] = transport
|
|
out.append(cfg)
|
|
return out
|
|
|
|
|
|
def _tool_result_to_text(result: Any) -> str:
|
|
"""把 MCP CallToolResult 转成发给模型的纯文本。"""
|
|
try:
|
|
if result is None:
|
|
return ""
|
|
# CallToolResult: content list of TextContent / ImageContent / ...
|
|
parts = []
|
|
content = getattr(result, "content", None)
|
|
if content is None and isinstance(result, dict):
|
|
content = result.get("content")
|
|
if not content:
|
|
# 兜底序列化
|
|
return json.dumps(result, ensure_ascii=False, default=str)[:8000]
|
|
for block in content:
|
|
btype = getattr(block, "type", None) or (block.get("type") if isinstance(block, dict) else None)
|
|
if btype == "text":
|
|
text = getattr(block, "text", None) or (block.get("text") if isinstance(block, dict) else "")
|
|
parts.append(text or "")
|
|
else:
|
|
parts.append(json.dumps(block, ensure_ascii=False, default=str))
|
|
text = "\n".join(p for p in parts if p)
|
|
is_err = getattr(result, "isError", False)
|
|
if is_err:
|
|
return f"[tool error]\n{text}"[:8000]
|
|
return text[:8000]
|
|
except Exception as e:
|
|
return f"[tool result parse error] {e}"
|
|
|
|
|
|
class McpHub:
|
|
"""
|
|
一次回复周期内的 MCP 连接池。
|
|
用法：
|
|
async with McpHub() as hub:
|
|
tools = hub.openai_tools()
|
|
text = await hub.call_tool("srv__search", {"q": "..."})
|
|
"""
|
|
|
|
def __init__(self, servers: list[dict] | None = None):
|
|
self._cfgs = servers if servers is not None else _enabled_servers()
|
|
self._stack: Optional[AsyncExitStack] = None
|
|
self._sessions: dict[str, Any] = {} # server_name -> ClientSession
|
|
self._tool_meta: dict[str, dict] = {} # namespaced -> {server, tool, schema...}
|
|
self._openai_tools: list[dict] = []
|
|
|
|
@property
|
|
def server_names(self) -> list[str]:
|
|
return list(self._sessions.keys())
|
|
|
|
@property
|
|
def tool_count(self) -> int:
|
|
return len(self._openai_tools)
|
|
|
|
def openai_tools(self) -> list[dict]:
|
|
return list(self._openai_tools)
|
|
|
|
async def __aenter__(self) -> "McpHub":
|
|
self._stack = AsyncExitStack()
|
|
await self._stack.__aenter__()
|
|
for cfg in self._cfgs:
|
|
try:
|
|
await self._connect(cfg)
|
|
except Exception as e:
|
|
print(f" [MCP] 连接失败 [{cfg.get('name')}]: {e}")
|
|
return self
|
|
|
|
async def __aexit__(self, *exc):
|
|
if self._stack is not None:
|
|
await self._stack.__aexit__(*exc)
|
|
self._sessions.clear()
|
|
self._tool_meta.clear()
|
|
self._openai_tools.clear()
|
|
self._stack = None
|
|
|
|
async def _connect(self, cfg: dict):
|
|
name = cfg["name"]
|
|
transport = cfg["transport"]
|
|
|
|
if transport == "stdio":
|
|
from mcp import ClientSession, StdioServerParameters
|
|
from mcp.client.stdio import stdio_client
|
|
|
|
env = os.environ.copy()
|
|
extra = cfg.get("env") or {}
|
|
env.update({str(k): str(v) for k, v in extra.items()})
|
|
params = StdioServerParameters(
|
|
command=cfg["command"],
|
|
args=list(cfg.get("args") or []),
|
|
env=env,
|
|
cwd=cfg.get("cwd") or None,
|
|
)
|
|
read, write = await self._stack.enter_async_context(stdio_client(params))
|
|
session = await self._stack.enter_async_context(ClientSession(read, write))
|
|
elif transport in ("sse", "http", "streamable_http"):
|
|
# 优先 streamable HTTP，其次 SSE（视 SDK 版本而定）
|
|
url = cfg["url"]
|
|
headers = cfg.get("headers") or {}
|
|
session = None
|
|
try:
|
|
from mcp import ClientSession
|
|
from mcp.client.streamable_http import streamablehttp_client
|
|
read, write, _ = await self._stack.enter_async_context(
|
|
streamablehttp_client(url, headers=headers)
|
|
)
|
|
session = await self._stack.enter_async_context(ClientSession(read, write))
|
|
except Exception:
|
|
from mcp import ClientSession
|
|
from mcp.client.sse import sse_client
|
|
read, write = await self._stack.enter_async_context(
|
|
sse_client(url, headers=headers)
|
|
)
|
|
session = await self._stack.enter_async_context(ClientSession(read, write))
|
|
else:
|
|
raise ValueError(f"unsupported transport: {transport}")
|
|
|
|
await session.initialize()
|
|
self._sessions[name] = session
|
|
|
|
listed = await session.list_tools()
|
|
tools = getattr(listed, "tools", None) or []
|
|
for tool in tools:
|
|
tname = getattr(tool, "name", "") or ""
|
|
if not tname:
|
|
continue
|
|
ns = _safe_name(name, tname)
|
|
desc = getattr(tool, "description", None) or f"MCP tool {tname} from {name}"
|
|
schema = getattr(tool, "inputSchema", None) or {"type": "object", "properties": {}}
|
|
if isinstance(schema, dict):
|
|
params = schema
|
|
else:
|
|
# pydantic model
|
|
try:
|
|
params = schema if isinstance(schema, dict) else dict(schema)
|
|
except Exception:
|
|
params = {"type": "object", "properties": {}}
|
|
|
|
self._tool_meta[ns] = {
|
|
"server": name,
|
|
"tool": tname,
|
|
"description": desc,
|
|
}
|
|
self._openai_tools.append({
|
|
"type": "function",
|
|
"function": {
|
|
"name": ns,
|
|
"description": f"[{name}] {desc}",
|
|
"parameters": params,
|
|
},
|
|
})
|
|
print(f" [MCP] 已连接 [{name}]，工具 {len(tools)} 个")
|
|
|
|
async def call_tool(self, namespaced: str, arguments: dict | None = None) -> str:
|
|
server, tool = _parse_name(namespaced)
|
|
# 若命名被 safe 化，用 meta 反查真实 tool 名
|
|
meta = self._tool_meta.get(namespaced)
|
|
if meta:
|
|
server = meta["server"]
|
|
tool = meta["tool"]
|
|
session = self._sessions.get(server)
|
|
if session is None:
|
|
return f"[MCP] server not connected: {server}"
|
|
try:
|
|
result = await session.call_tool(tool, arguments=arguments or {})
|
|
return _tool_result_to_text(result)
|
|
except Exception as e:
|
|
return f"[MCP] call_tool failed: {e}"
|
|
|
|
|
|
def run_coro(coro):
|
|
"""在同步代码里跑一段 async（回复线程用）。"""
|
|
try:
|
|
loop = asyncio.get_running_loop()
|
|
except RuntimeError:
|
|
loop = None
|
|
if loop and loop.is_running():
|
|
# 极少见：已在事件循环中 → 新开线程跑
|
|
import concurrent.futures
|
|
with concurrent.futures.ThreadPoolExecutor(max_workers=1) as pool:
|
|
return pool.submit(asyncio.run, coro).result()
|
|
return asyncio.run(coro)
|
|
|
|
|
|
async def probe_servers(servers: list[dict] | None = None) -> dict:
|
|
"""探测 MCP 服务器连通性与工具列表（供 GUI 测试）。"""
|
|
async with McpHub(servers) as hub:
|
|
tools = [
|
|
{"name": n, "server": m["server"], "tool": m["tool"], "description": m["description"]}
|
|
for n, m in hub._tool_meta.items()
|
|
]
|
|
return {
|
|
"ok": True,
|
|
"servers": hub.server_names,
|
|
"tool_count": hub.tool_count,
|
|
"tools": tools,
|
|
}
|