186 lines
6.9 KiB
Python
186 lines
6.9 KiB
Python
"""按 mcp_servers.json 连接 MCP 服务，把所有服务的工具交给 agent
|
||
|
||
每一项是一个服务：键为服务名，transport / url / headers 原样交给 MultiServerMCPClient；
|
||
填写 module 的服务只能在本机，由应用以子进程拉起，不填表示服务由外部部署
|
||
各环境不同的值和密钥写成 ${变量名}，加载时从环境变量展开
|
||
"""
|
||
|
||
import asyncio
|
||
import json
|
||
import os
|
||
import re
|
||
import socket
|
||
import subprocess
|
||
import sys
|
||
from collections import Counter
|
||
from pathlib import Path
|
||
from urllib.parse import urlsplit
|
||
|
||
from langchain_core.tools import BaseTool
|
||
from langchain_mcp_adapters.client import MultiServerMCPClient
|
||
from langchain_mcp_adapters.interceptors import MCPToolCallRequest
|
||
from loguru import logger
|
||
from mcp.types import CallToolResult
|
||
|
||
from app.config import settings
|
||
from app.resources.component import get_catalog_cache
|
||
|
||
PROJECT_ROOT = Path(__file__).resolve().parents[2]
|
||
CONFIG_PATH = PROJECT_ROOT / "mcp_servers.json"
|
||
# 物料服务返回的目录版本须与主进程一致，否则模型查到的字段和校验依据不同
|
||
MATERIAL_SERVER = "materials"
|
||
LOCAL_HOSTS = {"localhost", "127.0.0.1"}
|
||
# 配置中引用环境变量的写法
|
||
ENV_REFERENCE = re.compile(r"\$\{(\w+)\}")
|
||
|
||
|
||
class MCPServiceError(RuntimeError):
|
||
"""MCP 配置错误、服务不可用，或物料目录版本不一致"""
|
||
|
||
|
||
def _expand_env(value):
|
||
"""只展开字符串值里的 ${NAME}；变量未设置时报错，不把字面量当成地址或密钥"""
|
||
if isinstance(value, dict):
|
||
return {key: _expand_env(item) for key, item in value.items()}
|
||
if isinstance(value, list):
|
||
return [_expand_env(item) for item in value]
|
||
if not isinstance(value, str):
|
||
return value
|
||
|
||
def lookup(match: re.Match) -> str:
|
||
if (found := os.environ.get(match[1])) is None:
|
||
raise MCPServiceError(f"MCP 服务配置引用的环境变量 {match[1]} 未设置")
|
||
return found
|
||
|
||
return ENV_REFERENCE.sub(lookup, value)
|
||
|
||
|
||
def load_mcp_config(path: Path = CONFIG_PATH) -> dict[str, dict]:
|
||
"""读取服务列表并展开环境变量；文件缺失或格式不对时报错"""
|
||
if not path.exists():
|
||
raise MCPServiceError(f"未找到 MCP 服务配置: {path}")
|
||
try:
|
||
config = _expand_env(json.loads(path.read_text(encoding="utf-8")))
|
||
except json.JSONDecodeError as error:
|
||
raise MCPServiceError(f"MCP 服务配置不是合法 JSON: {error}") from error
|
||
if not isinstance(config, dict) or not all(
|
||
isinstance(item, dict) and item.get("transport") and item.get("url")
|
||
for item in config.values()
|
||
):
|
||
raise MCPServiceError("MCP 服务配置的每一项都需要 transport 和 url")
|
||
return config
|
||
|
||
|
||
async def _limit_time(request: MCPToolCallRequest, handler):
|
||
try:
|
||
return await asyncio.wait_for(
|
||
handler(request), timeout=settings.mcp_tool_timeout
|
||
)
|
||
except TimeoutError as error:
|
||
raise MCPServiceError(
|
||
f"MCP 工具调用超时: {request.server_name}.{request.name}"
|
||
) from error
|
||
|
||
|
||
async def _check_catalog_version(request: MCPToolCallRequest, handler):
|
||
result = await handler(request)
|
||
if request.server_name != MATERIAL_SERVER:
|
||
return result
|
||
if isinstance(result, CallToolResult) and result.structuredContent:
|
||
served = result.structuredContent.get("catalog_version")
|
||
catalog = get_catalog_cache()
|
||
expected = catalog.version if catalog is not None else None
|
||
if served != expected:
|
||
raise MCPServiceError(
|
||
f"物料服务的目录版本 {served} 与主服务 {expected} 不一致，请重启服务"
|
||
)
|
||
return result
|
||
|
||
|
||
_tools: list[BaseTool] | None = None
|
||
|
||
|
||
async def get_mcp_tools() -> list[BaseTool]:
|
||
"""首次使用时拉取所有服务的工具并缓存；此前失败时重新拉取"""
|
||
global _tools
|
||
if _tools is None:
|
||
config = load_mcp_config()
|
||
client = MultiServerMCPClient(
|
||
{
|
||
name: {key: value for key, value in item.items() if key != "module"}
|
||
for name, item in config.items()
|
||
},
|
||
tool_interceptors=[_limit_time, _check_catalog_version],
|
||
# 参数不对等查询错误返回给模型自行修正；连接失败等仍直接抛出
|
||
handle_tool_errors=True,
|
||
)
|
||
try:
|
||
tools = await asyncio.wait_for(
|
||
client.get_tools(), timeout=settings.mcp_tool_timeout
|
||
)
|
||
except Exception as error:
|
||
raise MCPServiceError(f"MCP 服务不可用: {', '.join(config)}") from error
|
||
# 不加服务名前缀，skill 里写的工具名保持不变，因此不同服务的工具不能重名
|
||
if duplicated := [
|
||
name for name, count in Counter(t.name for t in tools).items() if count > 1
|
||
]:
|
||
raise MCPServiceError(f"不同 MCP 服务的工具重名: {duplicated}")
|
||
_tools = tools
|
||
return _tools
|
||
|
||
|
||
def _port_open(port: int) -> bool:
|
||
with socket.socket() as probe:
|
||
probe.settimeout(0.2)
|
||
return probe.connect_ex(("127.0.0.1", port)) == 0
|
||
|
||
|
||
async def _start(name: str, item: dict) -> subprocess.Popen | None:
|
||
"""本机地址且尚未运行时拉起服务；已运行则直接复用"""
|
||
target = urlsplit(item["url"])
|
||
if target.hostname not in LOCAL_HOSTS:
|
||
raise MCPServiceError(f"{name} 由应用拉起，只能使用本机地址")
|
||
port = target.port or 80
|
||
if _port_open(port):
|
||
return None
|
||
process = subprocess.Popen(
|
||
[sys.executable, "-X", "utf8", "-m", item["module"]],
|
||
cwd=PROJECT_ROOT,
|
||
env={**os.environ, "PORT": str(port)},
|
||
)
|
||
for _ in range(100):
|
||
if process.poll() is not None:
|
||
raise MCPServiceError(f"{name} 启动失败，请查看终端日志")
|
||
if _port_open(port):
|
||
logger.info("MCP 服务 {} 已启动: {}", name, item["url"])
|
||
return process
|
||
await asyncio.sleep(0.1)
|
||
await stop_mcp_servers([process])
|
||
raise MCPServiceError(f"{name} 启动超时")
|
||
|
||
|
||
async def start_mcp_servers() -> list[subprocess.Popen]:
|
||
"""拉起配置了 module 的服务，返回进程供退出时关闭；单个服务失败只记录日志"""
|
||
processes = []
|
||
for name, item in load_mcp_config().items():
|
||
if not item.get("module"):
|
||
continue
|
||
try:
|
||
if (process := await _start(name, item)) is not None:
|
||
processes.append(process)
|
||
except MCPServiceError as error:
|
||
logger.error("{}，用到它的请求将无法执行", error)
|
||
return processes
|
||
|
||
|
||
async def stop_mcp_servers(processes: list[subprocess.Popen]) -> None:
|
||
for process in processes:
|
||
if process.poll() is not None:
|
||
continue
|
||
process.terminate()
|
||
try:
|
||
await asyncio.to_thread(process.wait, timeout=5)
|
||
except subprocess.TimeoutExpired:
|
||
process.kill()
|
||
await asyncio.to_thread(process.wait)
|