sdpy-kit 0.1.0__py3-none-any.whl

kit/mcp_kit.py

@@ -0,0 +1,547 @@
+"""
+Copyright (c) 2021-2026 Clark Chang. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+
+Project: sdpy
+Author: Clark Chang
+"""
+
+# date: 2026-09-04
+"""通用 MCP 客户端工具类
+
+基于 MCP（Model Context Protocol）协议的通用客户端引擎，支持连接任意远程
+MCP 服务器（Streamable HTTP / SSE 传输）并调用其工具。
+
+架构（引擎 / 档案 / 密钥三层分离）::
+
+    业务代码
+       │  McpKit.call_tool(name, arguments)          # 按档案调用
+       │  McpKit.call_tool(name, args, server="x")   # 多服务器按名切换
+       ▼
+    McpKit（本类，通用引擎，零业务逻辑）
+       ├─ 档案: config/mcp_servers.json   # 业务项目声明连哪些 MCP 服务器
+       └─ 密钥: 环境变量 / 项目根 .env     # 档案中 token_env 指向变量名，值不入库
+
+使用方式::
+
+    # 1) 按档案调用（默认服务器，业务项目配好档案 + .env 即用）
+    result = McpKit.call_tool(
+        "search_paper_by_title", {"title": "水声通信", "size": 10}
+    )
+    print(result.data["data"])
+
+    # 2) 多服务器按名切换（server=None 时使用档案 default_server）
+    result = McpKit.call_tool("query", {...}, server="lit-service-b")
+
+    # 3) 工具发现
+    tools = McpKit.list_tools("aminer")   # [{name, description}]
+
+    # 4) 直连临时服务器（不经档案）
+    result = McpKit.call_tool_url(
+        "https://example.com/mcp", "query", {...}, token="..."
+    )
+
+    # 5) 异步环境 / 批量复用连接（FastAPI 等已有事件循环进程内必须用异步 API）
+    result = await McpKit.call_tool_async("query", {...}, server="aminer")
+    async with McpKit.session("aminer") as call:
+        r1 = await call("search_a", {...})
+        r2 = await call("search_b", {...})
+
+行为规范：
+- 业务响应 ``success=false``（如付费拦截"余额不足"）抛 McpKitError，
+  detail 携带原始响应（msg / code / log_id），便于上层诊断。
+- 档案配置 fallback_url 时，主端点连接失败自动回退（仅传输层失败回退，
+  业务失败不回退）。
+- 返回 content 非 JSON 文本时降级为原始文本放入 McpToolResult.data。
+- 同步 API 内部使用 asyncio.run，禁止在已有事件循环的进程（如 FastAPI
+  请求处理）内调用；异步场景请使用 call_tool_async / session。
+
+设计约定：
+- 无内部可变状态，@classmethod 风格，线程安全（与 MinerUKit 同哲学）。
+- 短连接为主（每次调用独立握手，实测约 1s）；高频批量请用 session() 复用。
+- 零业务逻辑：不感知任何具体 MCP 服务（如文献查询）的业务编排。
+"""
+
+import asyncio
+import json
+import os
+import time
+from collections.abc import AsyncIterator, Awaitable, Callable
+from contextlib import asynccontextmanager
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any
+
+from fastmcp import Client
+from fastmcp.client.transports import SSETransport, StreamableHttpTransport
+
+from kit.config.kit_config import get_settings
+from kit.config.log_config import get_default_logger
+
+logger = get_default_logger(__name__)
+
+# 传输类型常量
+_TRANSPORT_STREAMABLE: str = "streamable-http"
+_TRANSPORT_SSE: str = "sse"
+_VALID_TRANSPORTS: tuple[str, ...] = (_TRANSPORT_STREAMABLE, _TRANSPORT_SSE)
+
+# 文件缺失报错时附带的最小可用模板（token_env 为 token 所在环境变量名）
+_MINIMAL_SERVERS_JSON: str = (
+    '{"default_server": "aminer", "servers": {"aminer": '
+    '{"transport": "streamable-http", "url": "https://mcp.aminer.cn/mcp", '
+    '"token_env": "AMINER_TOKEN", "timeout": 60}}}'
+)
+
+
+class McpKitError(Exception):
+    """MCP 客户端异常（档案缺失 / 密钥未配置 / 连接失败 / 业务失败）
+
+    Attributes:
+        detail: 诊断详情。业务失败时为服务端原始响应 dict（如
+            ``{"success": false, "msg": "余额不足，请充值"}``）；
+            连接失败时为最后一次传输层错误字符串。
+    """
+
+    def __init__(self, message: str, *, detail: Any = None) -> None:
+        super().__init__(message)
+        self.detail = detail
+
+
+@dataclass(frozen=True, slots=True)
+class McpToolResult:
+    """单次工具调用结果
+
+    Attributes:
+        server: 服务器名（档案模式）或 URL（直连模式）
+        tool: 工具名
+        elapsed_ms: 调用耗时（毫秒）
+        data: 业务数据。JSON 响应解析后的对象（dict / list）；
+              非 JSON 文本时为原始字符串
+    """
+
+    server: str
+    tool: str
+    elapsed_ms: float
+    data: Any
+
+
+@dataclass(frozen=True, slots=True)
+class _ServerProfile:
+    """档案中单个服务器的配置（加载时已校验）"""
+
+    name: str
+    url: str
+    transport: str = _TRANSPORT_STREAMABLE
+    fallback_url: str | None = None
+    token_env: str | None = None
+    timeout: int | None = None
+    description: str = ""
+
+
+class McpKit:
+    """通用 MCP 客户端工具类（引擎与业务解耦，档案驱动多服务器）
+
+    服务器档案默认读取 ``Settings.MCP_SERVERS_CONFIG_PATH`` 指向的
+    ``config/mcp_servers.json``，字段说明见该文件模板与模块 docstring。
+    """
+
+    # ------------------------------------------------------------------
+    # 档案加载与校验
+    # ------------------------------------------------------------------
+
+    @staticmethod
+    def _load_config() -> tuple[dict[str, _ServerProfile], str | None]:
+        """加载并校验服务器档案，返回 (档案字典, 默认服务器名或 None)
+
+        Raises:
+            McpKitError: 档案不存在 / JSON 非法 / 字段校验失败
+        """
+        path = Path(get_settings().MCP_SERVERS_CONFIG_PATH)
+        if not path.is_file():
+            raise McpKitError(
+                f"MCP 服务器档案不存在: {path}。"
+                f"请在项目根创建该文件，最小模板：\n{_MINIMAL_SERVERS_JSON}\n"
+                "（token_env 为 token 所在环境变量名，token 配置在 .env），"
+                "或调整 Settings.MCP_SERVERS_CONFIG_PATH"
+            )
+        try:
+            raw: Any = json.loads(path.read_text(encoding="utf-8"))
+        except json.JSONDecodeError as e:
+            raise McpKitError(f"MCP 服务器档案 JSON 解析失败: {path}: {e}") from e
+        if not isinstance(raw, dict):
+            raise McpKitError(f"MCP 服务器档案顶层必须是 JSON 对象: {path}")
+
+        servers: Any = raw.get("servers")
+        if not isinstance(servers, dict) or not servers:
+            raise McpKitError(f"MCP 服务器档案缺少 servers 配置或为空: {path}")
+        profiles = {
+            name: McpKit._build_profile(name, cfg, path)
+            for name, cfg in servers.items()
+        }
+
+        default: Any = raw.get("default_server")
+        if default is not None:
+            if not isinstance(default, str) or default not in profiles:
+                raise McpKitError(
+                    f"default_server {default!r} 未在 servers 中定义，"
+                    f"可用: {sorted(profiles)}（{path}）"
+                )
+        return profiles, default
+
+    @classmethod
+    def _build_profile(cls, name: str, cfg: Any, path: Path) -> _ServerProfile:
+        """校验并构造单个服务器档案"""
+        if not isinstance(cfg, dict):
+            raise McpKitError(f"MCP 服务器 {name!r} 配置必须是 JSON 对象: {path}")
+        url = cfg.get("url")
+        if not isinstance(url, str) or not url.strip():
+            raise McpKitError(f"MCP 服务器 {name!r} 缺少必填字段 url: {path}")
+        transport = cfg.get("transport", _TRANSPORT_STREAMABLE)
+        if transport not in _VALID_TRANSPORTS:
+            raise McpKitError(
+                f"MCP 服务器 {name!r} transport 非法: {transport!r}，"
+                f"可选 {_VALID_TRANSPORTS}"
+            )
+        timeout = cfg.get("timeout")
+        if timeout is not None and (not isinstance(timeout, int) or timeout <= 0):
+            raise McpKitError(f"MCP 服务器 {name!r} timeout 必须为正整数（秒）: {path}")
+        token_env = cfg.get("token_env")
+        if token_env is not None and (
+            not isinstance(token_env, str) or not token_env.strip()
+        ):
+            raise McpKitError(f"MCP 服务器 {name!r} token_env 必须为非空字符串: {path}")
+        return _ServerProfile(
+            name=name,
+            url=url.strip(),
+            transport=transport,
+            fallback_url=cfg.get("fallback_url"),
+            token_env=token_env,
+            timeout=timeout,
+            description=str(cfg.get("description", "")),
+        )
+
+    @classmethod
+    def _resolve_profile(cls, server: str | None) -> _ServerProfile:
+        """按名解析服务器档案；server=None 时回退档案 default_server
+
+        Raises:
+            McpKitError: server 未定义 / 未指定且无 default_server
+        """
+        profiles, default = cls._load_config()
+        name = server if server is not None else default
+        if name is None:
+            raise McpKitError(
+                f"未指定 server 且档案未配置 default_server，可用: {sorted(profiles)}"
+            )
+        profile = profiles.get(name)
+        if profile is None:
+            raise McpKitError(
+                f"MCP 服务器 {name!r} 未在档案中定义，可用: {sorted(profiles)}"
+            )
+        return profile
+
+    # ------------------------------------------------------------------
+    # 密钥与传输构造
+    # ------------------------------------------------------------------
+
+    @staticmethod
+    def _parse_env_file(env_path: Path) -> dict[str, str]:
+        """解析 .env 文件（仅支持 KEY=VALUE 行，忽略注释与空行）"""
+        result: dict[str, str] = {}
+        if not env_path.is_file():
+            return result
+        for line in env_path.read_text(encoding="utf-8").splitlines():
+            line = line.strip()
+            if not line or line.startswith("#") or "=" not in line:
+                continue
+            key, _, value = line.partition("=")
+            result[key.strip()] = value.strip()
+        return result
+
+    @classmethod
+    def _resolve_token(cls, profile: _ServerProfile) -> str | None:
+        """解析服务器 token：环境变量优先，回退项目根 .env 文件
+
+        与 pydantic Settings 的 env_file 同源（相对工作目录的 .env），
+        解决纯 .env 配置不进入 os.environ 的场景。
+
+        Raises:
+            McpKitError: 档案声明了 token_env 但两处均未配置
+        """
+        if not profile.token_env:
+            return None
+        token = os.environ.get(profile.token_env, "").strip()
+        if not token:
+            env_file = cls._parse_env_file(Path.cwd() / ".env")
+            token = env_file.get(profile.token_env, "").strip()
+        if not token:
+            raise McpKitError(
+                f"未配置环境变量 {profile.token_env}"
+                f"（请设置环境变量或在项目根 .env 中添加 {profile.token_env}=<token>）"
+            )
+        return token
+
+    @staticmethod
+    def _build_transport(url: str, token: str | None, transport_type: str) -> Any:
+        """构造传输层（Bearer 鉴权，token 为 None 时匿名）"""
+        headers = {"Authorization": f"Bearer {token}"} if token else None
+        if transport_type == _TRANSPORT_SSE:
+            return SSETransport(url, headers=headers)
+        return StreamableHttpTransport(url, headers=headers)
+
+    @staticmethod
+    def _timeout_of(profile: _ServerProfile) -> int:
+        """档案超时优先，缺省用全局 MCP_CALL_TIMEOUT"""
+        if profile.timeout is not None:
+            return profile.timeout
+        return get_settings().MCP_CALL_TIMEOUT
+
+    # ------------------------------------------------------------------
+    # 调用内核
+    # ------------------------------------------------------------------
+
+    @staticmethod
+    def _extract_text(result: Any) -> str:
+        """提取 MCP 调用结果中的全部文本段（多段以换行拼接）"""
+        parts: list[str] = []
+        for block in result.content:
+            text = getattr(block, "text", None)
+            if text is not None:
+                parts.append(text)
+        return "\n".join(parts)
+
+    @classmethod
+    async def _invoke(
+        cls, client: Any, server_name: str, tool: str, arguments: dict[str, Any]
+    ) -> McpToolResult:
+        """在已建立的连接上调用工具：耗时统计 → 解析 → 业务判定
+
+        Raises:
+            McpKitError: 业务响应 success=false（detail 携带原始响应）
+        """
+        start = time.perf_counter()
+        result = await client.call_tool(tool, arguments)
+        elapsed = round((time.perf_counter() - start) * 1000, 1)
+        raw = cls._extract_text(result)
+        try:
+            data: Any = json.loads(raw)
+        except json.JSONDecodeError:
+            data = raw
+        if isinstance(data, dict) and data.get("success") is False:
+            msg = str(data.get("msg", ""))
+            logger.error("❌ MCP 业务失败（%s/%s）: %s", server_name, tool, msg)
+            raise McpKitError(
+                f"MCP 业务失败（{server_name}/{tool}）: {msg}", detail=data
+            )
+        logger.info("✅ MCP 调用成功（%s/%s）: %sms", server_name, tool, elapsed)
+        return McpToolResult(
+            server=server_name, tool=tool, elapsed_ms=elapsed, data=data
+        )
+
+    @classmethod
+    async def _call_with_fallback(
+        cls, profile: _ServerProfile, name: str, arguments: dict[str, Any]
+    ) -> McpToolResult:
+        """档案模式调用：主端点失败自动回退 fallback_url（业务失败不回退）"""
+        token = cls._resolve_token(profile)
+        timeout = cls._timeout_of(profile)
+        urls = [profile.url]
+        if profile.fallback_url:
+            urls.append(profile.fallback_url)
+        last_error: Exception | None = None
+        for url in urls:
+            client = Client(
+                cls._build_transport(url, token, profile.transport), timeout=timeout
+            )
+            try:
+                async with client:
+                    return await cls._invoke(client, profile.name, name, arguments)
+            except McpKitError:
+                raise  # 服务端已应答的业务失败，不做端点回退
+            except Exception as e:  # noqa: BLE001 —— 传输层异常需尝试回退端点
+                last_error = e
+                logger.warning(
+                    "⚠ MCP 端点连接失败（%s → %s）: %s", profile.name, url, e
+                )
+        raise McpKitError(
+            f"MCP 连接失败（{profile.name}/{name}），"
+            f"已尝试 {len(urls)} 个端点: {last_error}",
+            detail=str(last_error),
+        ) from last_error
+
+    # ------------------------------------------------------------------
+    # 公开 API：档案模式（同步）
+    # ------------------------------------------------------------------
+
+    @classmethod
+    def list_tools(cls, server: str | None = None) -> list[dict[str, Any]]:
+        """发现服务器的工具清单
+
+        Args:
+            server: 服务器名，None 时使用档案 default_server
+
+        Returns:
+            list[dict]: [{name, description}]
+
+        Raises:
+            McpKitError: 档案/密钥问题或连接失败
+        """
+        profile = cls._resolve_profile(server)
+        return asyncio.run(cls._list_tools_async(profile))
+
+    @classmethod
+    async def _list_tools_async(cls, profile: _ServerProfile) -> list[dict[str, Any]]:
+        """工具发现（主端点失败自动回退 fallback_url）"""
+        token = cls._resolve_token(profile)
+        timeout = cls._timeout_of(profile)
+        urls = [profile.url]
+        if profile.fallback_url:
+            urls.append(profile.fallback_url)
+        last_error: Exception | None = None
+        for url in urls:
+            client = Client(
+                cls._build_transport(url, token, profile.transport), timeout=timeout
+            )
+            try:
+                async with client:
+                    tools = await client.list_tools()
+                    logger.info(
+                        "✅ MCP 工具发现（%s）: %d 个工具", profile.name, len(tools)
+                    )
+                    return [
+                        {"name": t.name, "description": t.description or ""}
+                        for t in tools
+                    ]
+            except Exception as e:  # noqa: BLE001 —— 传输层异常需尝试回退端点
+                last_error = e
+                logger.warning(
+                    "⚠ MCP 端点连接失败（%s → %s）: %s", profile.name, url, e
+                )
+        raise McpKitError(
+            f"MCP 连接失败（{profile.name}/list_tools），"
+            f"已尝试 {len(urls)} 个端点: {last_error}",
+            detail=str(last_error),
+        ) from last_error
+
+    @classmethod
+    def call_tool(
+        cls,
+        name: str,
+        arguments: dict[str, Any],
+        *,
+        server: str | None = None,
+    ) -> McpToolResult:
+        """按档案调用工具（同步短连接，开箱即用主路径）
+
+        Args:
+            name: 工具名
+            arguments: 工具入参
+            server: 服务器名，None 时使用档案 default_server
+
+        Returns:
+            McpToolResult: 调用结果（data 为解析后的业务数据）
+
+        Raises:
+            McpKitError: 业务失败（detail 含原始响应）/ 档案 / 密钥 / 连接失败
+        """
+        profile = cls._resolve_profile(server)
+        return asyncio.run(cls._call_with_fallback(profile, name, arguments))
+
+    # ------------------------------------------------------------------
+    # 公开 API：异步变体与连接复用
+    # ------------------------------------------------------------------
+
+    @classmethod
+    async def call_tool_async(
+        cls,
+        name: str,
+        arguments: dict[str, Any],
+        *,
+        server: str | None = None,
+    ) -> McpToolResult:
+        """call_tool 的异步版本（已有事件循环的进程内使用）"""
+        profile = cls._resolve_profile(server)
+        return await cls._call_with_fallback(profile, name, arguments)
+
+    @classmethod
+    @asynccontextmanager
+    async def session(
+        cls, server: str | None = None
+    ) -> AsyncIterator[Callable[[str, dict[str, Any]], Awaitable[McpToolResult]]]:
+        """建立长连接会话，批量调用复用连接（高频场景）
+
+        Usage::
+
+            async with McpKit.session("aminer") as call:
+                r1 = await call("search_paper_by_title", {...})
+                r2 = await call("get_papers_info", {...})
+
+        注意：会话固定使用主端点，不做 fallback 回退。
+        """
+        profile = cls._resolve_profile(server)
+        token = cls._resolve_token(profile)
+        timeout = cls._timeout_of(profile)
+        async with Client(
+            cls._build_transport(profile.url, token, profile.transport), timeout=timeout
+        ) as client:
+
+            async def call(tool: str, arguments: dict[str, Any]) -> McpToolResult:
+                return await cls._invoke(client, profile.name, tool, arguments)
+
+            yield call
+
+    # ------------------------------------------------------------------
+    # 公开 API：直连模式（不经档案）
+    # ------------------------------------------------------------------
+
+    @classmethod
+    def call_tool_url(
+        cls,
+        url: str,
+        name: str,
+        arguments: dict[str, Any],
+        *,
+        token: str | None = None,
+        transport: str = _TRANSPORT_STREAMABLE,
+        timeout: int | None = None,
+    ) -> McpToolResult:
+        """直连临时服务器调用工具（不经档案，用于探索/临时端点）
+
+        Args:
+            url: MCP 服务器端点
+            name: 工具名
+            arguments: 工具入参
+            token: Bearer token，None 时匿名
+            transport: 传输类型（streamable-http / sse）
+            timeout: 超时秒数，None 时用全局 MCP_CALL_TIMEOUT
+
+        Raises:
+            McpKitError: transport 非法 / 业务失败 / 连接失败
+        """
+        if transport not in _VALID_TRANSPORTS:
+            raise McpKitError(
+                f"transport 非法: {transport!r}，可选 {_VALID_TRANSPORTS}"
+            )
+
+        async def _run() -> McpToolResult:
+            if timeout is not None:
+                effective_timeout: int = timeout
+            else:
+                effective_timeout = get_settings().MCP_CALL_TIMEOUT
+            async with Client(
+                cls._build_transport(url, token, transport), timeout=effective_timeout
+            ) as client:
+                return await cls._invoke(client, url, name, arguments)
+
+        return asyncio.run(_run())