119 lines
4.2 KiB
Python
119 lines
4.2 KiB
Python
"""物料字段 MCP 客户端：把字段查询工具交给模型使用
|
|
|
|
服务只监听本机，由应用启动时拉起；不可用时直接报错，不退回其他查询方式
|
|
"""
|
|
|
|
import asyncio
|
|
import os
|
|
import socket
|
|
import subprocess
|
|
import sys
|
|
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
|
|
|
|
SERVER_NAME = "materials"
|
|
PROJECT_ROOT = Path(__file__).resolve().parents[2]
|
|
|
|
|
|
class MaterialFieldServiceError(RuntimeError):
|
|
"""物料字段服务不可用，或与主进程的物料目录版本不一致"""
|
|
|
|
|
|
async def _check_version(request: MCPToolCallRequest, handler):
|
|
"""限制单次查询耗时；返回的目录版本须与主进程一致，否则模型查到的字段和校验依据不同"""
|
|
try:
|
|
result = await asyncio.wait_for(
|
|
handler(request), timeout=settings.mcp_tool_timeout
|
|
)
|
|
except TimeoutError as error:
|
|
raise MaterialFieldServiceError(f"物料字段查询超时: {request.name}") from error
|
|
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 MaterialFieldServiceError(
|
|
f"物料字段服务的目录版本 {served} 与主服务 {expected} 不一致，请重启服务"
|
|
)
|
|
return result
|
|
|
|
|
|
_client = MultiServerMCPClient(
|
|
{
|
|
SERVER_NAME: {
|
|
"transport": "streamable_http",
|
|
"url": settings.mcp_material_url,
|
|
}
|
|
},
|
|
tool_interceptors=[_check_version],
|
|
# 未知组件类型、分组等查询错误返回给模型自行修正；连接失败等仍直接抛出
|
|
handle_tool_errors=True,
|
|
)
|
|
_tools: list[BaseTool] | None = None
|
|
|
|
|
|
async def get_material_tools() -> list[BaseTool]:
|
|
"""字段查询工具；首次使用或此前连接失败时向服务拉取工具清单"""
|
|
global _tools
|
|
if _tools is None:
|
|
try:
|
|
_tools = await asyncio.wait_for(
|
|
_client.get_tools(server_name=SERVER_NAME),
|
|
timeout=settings.mcp_tool_timeout,
|
|
)
|
|
except Exception as error:
|
|
raise MaterialFieldServiceError(
|
|
f"物料字段服务不可用: {settings.mcp_material_url}"
|
|
) from error
|
|
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_material_server() -> subprocess.Popen | None:
|
|
"""本机地址且尚未运行时拉起服务，返回进程供退出时关闭；已运行则直接复用"""
|
|
target = urlsplit(settings.mcp_material_url)
|
|
if target.hostname not in {"localhost", "127.0.0.1"}:
|
|
raise MaterialFieldServiceError("物料字段服务只能部署在本机，不对外暴露")
|
|
port = target.port or 80
|
|
if _port_open(port):
|
|
return None
|
|
process = subprocess.Popen(
|
|
[sys.executable, "-X", "utf8", "-m", "mcp_servers.material_server"],
|
|
cwd=PROJECT_ROOT,
|
|
env={**os.environ, "PORT": str(port)},
|
|
)
|
|
for _ in range(100):
|
|
if process.poll() is not None:
|
|
raise MaterialFieldServiceError("物料字段服务启动失败，请查看终端日志")
|
|
if _port_open(port):
|
|
logger.info("物料字段服务已启动: {}", settings.mcp_material_url)
|
|
return process
|
|
await asyncio.sleep(0.1)
|
|
await stop_material_server(process)
|
|
raise MaterialFieldServiceError("物料字段服务启动超时")
|
|
|
|
|
|
async def stop_material_server(process: subprocess.Popen | None) -> None:
|
|
if process is None or process.poll() is not None:
|
|
return
|
|
process.terminate()
|
|
try:
|
|
await asyncio.to_thread(process.wait, timeout=5)
|
|
except subprocess.TimeoutExpired:
|
|
process.kill()
|
|
await asyncio.to_thread(process.wait)
|