更新

wechat_rpa/mcp_bridge.py

@@ -0,0 +1,265 @@
+"""
+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,
+        }