feat: 接入 MCP 外部工具服务,AI 可调用外部工具获取书外信息
- 后端新增 MCP 客户端(Streamable HTTP / 老式 SSE)与流式 agent 工具循环 - 新增 /api/mcp/servers 管理接口与 mcp_servers 表 - 对话中 OpenAI 兼容模型可调用启用服务器的工具,SSE 推送 tool 事件 - 前端新增 MCP 服务管理弹窗与侧栏入口,流式回答展示工具调用状态 - 修复用户消息气泡缺少删除按钮,并级联删除其名下回答版本 Co-Authored-By: Claude Code <noreply@anthropic.com>
This commit is contained in:
+42
-5
@@ -14,6 +14,7 @@ from app.rag.engine import (
|
||||
search_segments, build_chat_prompt, parse_citations,
|
||||
estimate_tokens, match_chapters, _trim_messages,
|
||||
)
|
||||
from app.mcp.agent import agent_stream, build_mcp_context, execute_mcp_tool
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
@@ -164,10 +165,24 @@ async def send_message(conv_id: str, req: SendMessageRequest, db: AsyncSession =
|
||||
chapter_segments=chapter_segments,
|
||||
)
|
||||
|
||||
total_tokens = estimate_tokens("".join(m["content"] for m in messages))
|
||||
provider = await resolve_provider(db, model_name)
|
||||
if provider is None:
|
||||
raise HTTPException(400, f"模型 {model_name} 未配置,请先在页面「管理模型」中添加或选择可用模型")
|
||||
|
||||
# MCP 外部工具:仅当模型原生支持 function calling 时接入工具循环;
|
||||
# 无 MCP 配置 / 单服务器失败时自动跳过,行为与改造前完全一致
|
||||
mcp_ctx = await build_mcp_context(db) if provider.supports_tools else None
|
||||
if mcp_ctx is not None:
|
||||
system = messages[0]["content"]
|
||||
system += (
|
||||
"\n\n## 外部信息工具\n"
|
||||
"当问题需要书中没有的外部信息(最新动态、时事、网络资料、行业数据等)时,"
|
||||
f"可调用以下外部 MCP 工具获取:{', '.join(mcp_ctx.names)}\n"
|
||||
"- 外部信息要说明来源与时效;回答主体仍以书中内容为准,并保留 [[ref:N]] 引用标注。\n"
|
||||
)
|
||||
messages[0]["content"] = system
|
||||
|
||||
total_tokens = estimate_tokens("".join(m["content"] for m in messages))
|
||||
if total_tokens > provider.max_context * 0.8:
|
||||
# 内容感知裁剪:保留 system(含检索原文 + 整章注入),从最旧历史开始丢
|
||||
messages = _trim_messages(messages, int(provider.max_context * 0.8))
|
||||
@@ -175,9 +190,25 @@ async def send_message(conv_id: str, req: SendMessageRequest, db: AsyncSession =
|
||||
async def stream_response():
|
||||
full_content = ""
|
||||
try:
|
||||
async for token in provider.chat_stream(messages, temperature=0.7):
|
||||
full_content += token
|
||||
yield f"data: {json.dumps({'type': 'token', 'content': token}, ensure_ascii=False)}\n\n"
|
||||
if mcp_ctx is not None:
|
||||
async for kind, payload in agent_stream(
|
||||
messages,
|
||||
provider,
|
||||
mcp_ctx.ai_tools,
|
||||
execute=lambda name, args: execute_mcp_tool(mcp_ctx, name, args),
|
||||
describe=mcp_ctx.display_name,
|
||||
):
|
||||
if kind == "token":
|
||||
full_content += payload
|
||||
yield f"data: {json.dumps({'type': 'token', 'content': payload}, ensure_ascii=False)}\n\n"
|
||||
elif kind == "tool":
|
||||
yield f"data: {json.dumps({'type': 'tool', 'name': payload, 'status': 'start'}, ensure_ascii=False)}\n\n"
|
||||
elif kind == "tool_done":
|
||||
yield f"data: {json.dumps({'type': 'tool', 'name': payload, 'status': 'done'}, ensure_ascii=False)}\n\n"
|
||||
else:
|
||||
async for token in provider.chat_stream(messages, temperature=0.7):
|
||||
full_content += token
|
||||
yield f"data: {json.dumps({'type': 'token', 'content': token}, ensure_ascii=False)}\n\n"
|
||||
except Exception as e:
|
||||
logger.exception("流式生成失败(conversation=%s)", conv_id)
|
||||
yield f"data: {json.dumps({'type': 'error', 'content': str(e)}, ensure_ascii=False)}\n\n"
|
||||
@@ -266,9 +297,15 @@ async def delete_message(msg_id: str, db: AsyncSession = Depends(get_db)):
|
||||
msg = await db.get(Message, msg_id)
|
||||
if not msg:
|
||||
raise HTTPException(404, "消息不存在")
|
||||
# 级联删除:该消息名下的所有回答版本(parent 指向它),避免残留孤儿版本
|
||||
versions = (await db.execute(
|
||||
select(Message).where(Message.parent_message_id == msg_id)
|
||||
)).scalars().all()
|
||||
for v in versions:
|
||||
await db.delete(v)
|
||||
await db.delete(msg)
|
||||
await db.commit()
|
||||
return {"deleted": msg_id}
|
||||
return {"deleted": msg_id, "cascade": len(versions)}
|
||||
|
||||
|
||||
@router.delete("/conversations/{conv_id}")
|
||||
|
||||
@@ -0,0 +1,146 @@
|
||||
import json
|
||||
|
||||
from fastapi import APIRouter, Depends, HTTPException
|
||||
from pydantic import BaseModel
|
||||
from sqlalchemy import select
|
||||
from sqlalchemy.ext.asyncio import AsyncSession
|
||||
|
||||
from app.core.database import get_db
|
||||
from app.models.mcp_server import McpServer
|
||||
from app.mcp.client import McpServerConfig, McpClientError, list_mcp_tools, invalidate_mcp_cache
|
||||
|
||||
router = APIRouter(prefix="/mcp", tags=["mcp"])
|
||||
|
||||
|
||||
class McpServerIn(BaseModel):
|
||||
name: str
|
||||
transport: str = "streamable-http" # streamable-http | sse
|
||||
url: str
|
||||
headers: dict[str, str] = {}
|
||||
enabled: bool = True
|
||||
|
||||
|
||||
class McpServerUpdate(BaseModel):
|
||||
name: str | None = None
|
||||
transport: str | None = None
|
||||
url: str | None = None
|
||||
headers: dict[str, str] | None = None
|
||||
enabled: bool | None = None
|
||||
|
||||
|
||||
class McpTestRequest(BaseModel):
|
||||
"""「测试连接」用临时配置(不落库)。"""
|
||||
|
||||
name: str = "test"
|
||||
transport: str = "streamable-http"
|
||||
url: str
|
||||
headers: dict[str, str] = {}
|
||||
|
||||
|
||||
def _to_out(m: McpServer) -> dict:
|
||||
return {
|
||||
"id": m.id,
|
||||
"name": m.name,
|
||||
"transport": m.transport,
|
||||
"url": m.url,
|
||||
"headers": m.headers_dict(),
|
||||
"enabled": m.enabled,
|
||||
"created_at": m.created_at.isoformat() if m.created_at else "",
|
||||
}
|
||||
|
||||
|
||||
@router.get("/servers")
|
||||
async def list_servers(db: AsyncSession = Depends(get_db)):
|
||||
result = await db.execute(select(McpServer).order_by(McpServer.created_at))
|
||||
return [_to_out(m) for m in result.scalars().all()]
|
||||
|
||||
|
||||
@router.post("/servers")
|
||||
async def add_server(req: McpServerIn, db: AsyncSession = Depends(get_db)):
|
||||
name = (req.name or "").strip()
|
||||
url = (req.url or "").strip()
|
||||
if not name:
|
||||
raise HTTPException(400, "服务器名称不能为空")
|
||||
if not url:
|
||||
raise HTTPException(400, "服务器 URL 不能为空")
|
||||
if req.transport not in ("streamable-http", "sse"):
|
||||
raise HTTPException(400, "不支持的传输协议")
|
||||
row = McpServer(
|
||||
name=name,
|
||||
transport=req.transport,
|
||||
url=url,
|
||||
headers=json.dumps(req.headers or {}, ensure_ascii=False),
|
||||
enabled=req.enabled,
|
||||
)
|
||||
db.add(row)
|
||||
await db.commit()
|
||||
await db.refresh(row)
|
||||
return _to_out(row)
|
||||
|
||||
|
||||
@router.put("/servers/{server_id}")
|
||||
async def update_server(server_id: str, req: McpServerUpdate, db: AsyncSession = Depends(get_db)):
|
||||
row = await db.get(McpServer, server_id)
|
||||
if not row:
|
||||
raise HTTPException(404, "服务器不存在")
|
||||
if req.name is not None:
|
||||
name = req.name.strip()
|
||||
if not name:
|
||||
raise HTTPException(400, "服务器名称不能为空")
|
||||
row.name = name
|
||||
if req.transport is not None:
|
||||
if req.transport not in ("streamable-http", "sse"):
|
||||
raise HTTPException(400, "不支持的传输协议")
|
||||
row.transport = req.transport
|
||||
if req.url is not None:
|
||||
url = req.url.strip()
|
||||
if not url:
|
||||
raise HTTPException(400, "服务器 URL 不能为空")
|
||||
row.url = url
|
||||
if req.headers is not None:
|
||||
row.headers = json.dumps(req.headers, ensure_ascii=False)
|
||||
if req.enabled is not None:
|
||||
row.enabled = req.enabled
|
||||
await db.commit()
|
||||
await db.refresh(row)
|
||||
await invalidate_mcp_cache(server_id) # 配置变更 → 失效旧会话 / 工具缓存
|
||||
return _to_out(row)
|
||||
|
||||
|
||||
@router.delete("/servers/{server_id}")
|
||||
async def delete_server(server_id: str, db: AsyncSession = Depends(get_db)):
|
||||
row = await db.get(McpServer, server_id)
|
||||
if not row:
|
||||
raise HTTPException(404, "服务器不存在")
|
||||
await db.delete(row)
|
||||
await db.commit()
|
||||
await invalidate_mcp_cache(server_id)
|
||||
return {"deleted": server_id}
|
||||
|
||||
|
||||
@router.post("/servers/test")
|
||||
async def test_server(req: McpTestRequest, db: AsyncSession = Depends(get_db)):
|
||||
"""用临时配置连一次 MCP 服务,拉取工具列表验证连通性(不落库)。"""
|
||||
url = (req.url or "").strip()
|
||||
if not url:
|
||||
raise HTTPException(400, "服务器 URL 不能为空")
|
||||
if req.transport not in ("streamable-http", "sse"):
|
||||
raise HTTPException(400, "不支持的传输协议")
|
||||
cfg = McpServerConfig(
|
||||
id=f"test-{hash(url)}",
|
||||
name=(req.name or "test").strip() or "test",
|
||||
transport=req.transport,
|
||||
url=url,
|
||||
headers=req.headers or {},
|
||||
enabled=True,
|
||||
)
|
||||
try:
|
||||
tools = await list_mcp_tools(cfg, fresh=True)
|
||||
except McpClientError as e:
|
||||
return {"ok": False, "error": str(e)}
|
||||
except Exception as e:
|
||||
return {"ok": False, "error": f"{type(e).__name__}: {e}"}
|
||||
return {
|
||||
"ok": True,
|
||||
"tools": [{"name": t.name, "description": t.description or "", "input_schema": t.input_schema} for t in tools],
|
||||
}
|
||||
+39
-4
@@ -18,24 +18,30 @@ class ChatResponse:
|
||||
class LLMProvider(ABC):
|
||||
name: str = "base"
|
||||
max_context: int = 128_000
|
||||
# 是否原生支持 function calling(OpenAI 兼容协议);Anthropic 暂未接入工具循环
|
||||
supports_tools: bool = False
|
||||
|
||||
@abstractmethod
|
||||
async def chat(self, messages: list[dict], temperature: float = 0.7) -> ChatResponse:
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
async def chat_stream(self, messages: list[dict], temperature: float = 0.7) -> AsyncGenerator[str, None]:
|
||||
async def chat_stream(self, messages: list[dict], temperature: float = 0.7, tools: list | None = None) -> AsyncGenerator[str, None]:
|
||||
...
|
||||
|
||||
|
||||
class OpenAIProvider(LLMProvider):
|
||||
name = "openai"
|
||||
max_context = 128_000
|
||||
supports_tools = True
|
||||
|
||||
def __init__(self, api_key: str | None = None, base_url: str | None = None, model: str | None = None):
|
||||
self.api_key = api_key if api_key is not None else os.getenv("OPENAI_API_KEY", "")
|
||||
self.base_url = base_url if base_url is not None else os.getenv("OPENAI_BASE_URL", "https://api.openai.com/v1")
|
||||
self.model = model if model is not None else os.getenv("OPENAI_MODEL", "gpt-4o")
|
||||
# 流式工具循环的工作态:chat_stream 结束时填充 tool_calls / last_content
|
||||
self.tool_calls: list[dict] = []
|
||||
self.last_content: str = ""
|
||||
|
||||
async def chat(self, messages: list[dict], temperature: float = 0.7) -> ChatResponse:
|
||||
import httpx
|
||||
@@ -56,15 +62,27 @@ class OpenAIProvider(LLMProvider):
|
||||
completion_tokens=data.get("usage", {}).get("completion_tokens", 0),
|
||||
)
|
||||
|
||||
async def chat_stream(self, messages: list[dict], temperature: float = 0.7) -> AsyncGenerator[str, None]:
|
||||
async def chat_stream(self, messages: list[dict], temperature: float = 0.7, tools: list | None = None) -> AsyncGenerator[str, None]:
|
||||
import httpx
|
||||
payload = {
|
||||
"model": self.model,
|
||||
"messages": messages,
|
||||
"temperature": temperature,
|
||||
"stream": True,
|
||||
}
|
||||
if tools:
|
||||
payload["tools"] = tools
|
||||
# 工作态复位:上一轮(若有)的 tool_calls 不残留
|
||||
self.tool_calls = []
|
||||
self.last_content = ""
|
||||
async with httpx.AsyncClient(timeout=300) as client:
|
||||
async with client.stream(
|
||||
"POST",
|
||||
f"{self.base_url}/chat/completions",
|
||||
headers={"Authorization": f"Bearer {self.api_key}", "Content-Type": "application/json"},
|
||||
json={"model": self.model, "messages": messages, "temperature": temperature, "stream": True},
|
||||
json=payload,
|
||||
) as resp:
|
||||
tool_calls_acc: dict[int, dict] = {}
|
||||
async for line in resp.aiter_lines():
|
||||
if line.startswith("data: "):
|
||||
data_str = line[6:]
|
||||
@@ -79,9 +97,25 @@ class OpenAIProvider(LLMProvider):
|
||||
delta = choices[0].get("delta", {})
|
||||
content = delta.get("content", "")
|
||||
if content:
|
||||
self.last_content += content
|
||||
yield content
|
||||
# 流式 tool_calls:每个 index 首个 delta 带 id/name,后续增量拼 arguments
|
||||
for tc in delta.get("tool_calls") or []:
|
||||
idx = tc.get("index", 0)
|
||||
entry = tool_calls_acc.setdefault(idx, {
|
||||
"id": "", "type": "function", "function": {"name": "", "arguments": ""},
|
||||
})
|
||||
if tc.get("id"):
|
||||
entry["id"] = tc["id"]
|
||||
fn = tc.get("function") or {}
|
||||
if fn.get("name"):
|
||||
entry["function"]["name"] = fn["name"]
|
||||
if fn.get("arguments"):
|
||||
entry["function"]["arguments"] += fn["arguments"]
|
||||
except json.JSONDecodeError:
|
||||
continue
|
||||
if tool_calls_acc:
|
||||
self.tool_calls = [tool_calls_acc[i] for i in sorted(tool_calls_acc)]
|
||||
|
||||
|
||||
class AnthropicProvider(LLMProvider):
|
||||
@@ -135,8 +169,9 @@ class AnthropicProvider(LLMProvider):
|
||||
completion_tokens=data.get("usage", {}).get("output_tokens", 0),
|
||||
)
|
||||
|
||||
async def chat_stream(self, messages: list[dict], temperature: float = 0.7) -> AsyncGenerator[str, None]:
|
||||
async def chat_stream(self, messages: list[dict], temperature: float = 0.7, tools: list | None = None) -> AsyncGenerator[str, None]:
|
||||
import httpx
|
||||
# tools 参数为兼容 ABC 而存在:Anthropic 暂未接入工具循环(supports_tools=False)
|
||||
system_msg, anthropic_msgs = self._to_anthropic_messages(messages)
|
||||
body = {
|
||||
"model": self.model,
|
||||
|
||||
@@ -10,10 +10,12 @@ from app.core.database import init_db, Base
|
||||
from app.api.books import router as books_router
|
||||
from app.api.chat import router as chat_router
|
||||
from app.api.models import router as models_router
|
||||
from app.api.mcp import router as mcp_router
|
||||
|
||||
from app.models.book import Book, Chapter, Segment
|
||||
from app.models.chat import Conversation, Message
|
||||
from app.models.ai_model import AiModel
|
||||
from app.models.mcp_server import McpServer
|
||||
|
||||
|
||||
@asynccontextmanager
|
||||
@@ -35,6 +37,7 @@ app.add_middleware(
|
||||
app.include_router(books_router, prefix="/api")
|
||||
app.include_router(chat_router, prefix="/api")
|
||||
app.include_router(models_router, prefix="/api")
|
||||
app.include_router(mcp_router, prefix="/api")
|
||||
|
||||
|
||||
@app.get("/api/health")
|
||||
|
||||
@@ -0,0 +1,197 @@
|
||||
"""agent — MCP 工具与对话 agent 的桥接层
|
||||
|
||||
把启用中的 MCP 服务器暴露的工具转换成 OpenAI 兼容 function calling 格式,
|
||||
并在对话时执行「流式工具循环」:模型文本即时流式输出,出现工具调用则执行并续一轮。
|
||||
单服务器工具获取失败只跳过该服务器,绝不阻断对话。
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import logging
|
||||
import re
|
||||
from typing import AsyncGenerator, Awaitable, Callable
|
||||
|
||||
from sqlalchemy import select
|
||||
from sqlalchemy.ext.asyncio import AsyncSession
|
||||
|
||||
from app.llm.hub import LLMProvider
|
||||
from app.models.mcp_server import McpServer
|
||||
from app.mcp.client import McpClientError, McpServerConfig, McpTool, call_mcp_tool, list_mcp_tools
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# MCP 工具名统一加前缀防冲突:mcp_<server短id>_<toolName>
|
||||
MCP_PREFIX = "mcp_"
|
||||
|
||||
# OpenAI function name 限制:^[a-zA-Z0-9_-]{1,64}$
|
||||
_INVALID_NAME_CHARS = re.compile(r"[^a-zA-Z0-9_-]")
|
||||
|
||||
MAX_TOOL_ROUNDS = 3
|
||||
TOOL_RESULT_CHARS = 6000 # 单条工具结果回传模型的上限(保护上下文预算)
|
||||
|
||||
|
||||
class McpContext:
|
||||
"""一次对话请求内启用中的 MCP 服务器快照:工具列表 + 短 id → 服务器映射。"""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self.servers: list[McpServerConfig] = []
|
||||
self.ai_tools: list[dict] = []
|
||||
self.names: list[str] = [] # 完整 AI 工具名(用于 system 提示词)
|
||||
self._exec_map: dict[str, McpServerConfig] = {} # 短 id → 服务器
|
||||
|
||||
@property
|
||||
def has_tools(self) -> bool:
|
||||
return bool(self.ai_tools)
|
||||
|
||||
def resolve(self, ai_tool_name: str) -> tuple[McpServerConfig, str]:
|
||||
"""把 mcp_ 前缀的 AI 工具名还原为 (服务器, 工具名)。"""
|
||||
rest = ai_tool_name[len(MCP_PREFIX):]
|
||||
sep = rest.find("_")
|
||||
if sep <= 0:
|
||||
raise McpClientError(f"工具名格式异常:{ai_tool_name}")
|
||||
short_id, tool_name = rest[:sep], rest[sep + 1:]
|
||||
server = self._exec_map.get(short_id)
|
||||
if server is None:
|
||||
raise McpClientError(f"MCP 服务器不存在或已停用({short_id})")
|
||||
return server, tool_name
|
||||
|
||||
def display_name(self, ai_tool_name: str) -> str:
|
||||
"""前端展示用名称:<服务器名>/<工具名>。"""
|
||||
try:
|
||||
server, tool_name = self.resolve(ai_tool_name)
|
||||
except McpClientError:
|
||||
return ai_tool_name
|
||||
return f"{server.name}/{tool_name}"
|
||||
|
||||
|
||||
def _sanitize_tool_name(name: str) -> str:
|
||||
return _INVALID_NAME_CHARS.sub("_", name)
|
||||
|
||||
|
||||
def _short_server_id(server_id: str, taken: set[str]) -> str:
|
||||
"""短 id:取 uuid 前缀,必要时加长保证本次请求内唯一(避免截断冲突)。"""
|
||||
base = server_id.replace("-", "")
|
||||
n = 8
|
||||
short = base[:n]
|
||||
while short in taken:
|
||||
n += 4
|
||||
short = base[:n]
|
||||
taken.add(short)
|
||||
return short
|
||||
|
||||
|
||||
async def build_mcp_context(db: AsyncSession) -> McpContext | None:
|
||||
"""收集启用中的 MCP 服务器及其工具;无服务器或全部失败返回 None。"""
|
||||
result = await db.execute(
|
||||
select(McpServer).where(McpServer.enabled.is_(True)).order_by(McpServer.created_at)
|
||||
)
|
||||
rows = list(result.scalars().all())
|
||||
if not rows:
|
||||
return None
|
||||
|
||||
ctx = McpContext()
|
||||
taken: set[str] = set()
|
||||
for row in rows:
|
||||
cfg = McpServerConfig(
|
||||
id=row.id, name=row.name, transport=row.transport,
|
||||
url=row.url, headers=row.headers_dict(), enabled=row.enabled,
|
||||
)
|
||||
try:
|
||||
tools = await list_mcp_tools(cfg)
|
||||
except Exception as e:
|
||||
logger.warning("[mcp] 服务器「%s」工具获取失败,已跳过:%s", row.name, e)
|
||||
continue
|
||||
if not tools:
|
||||
continue
|
||||
short = _short_server_id(row.id, taken)
|
||||
ctx.servers.append(cfg)
|
||||
ctx._exec_map[short] = cfg
|
||||
for tool in tools:
|
||||
ai_name = f"{MCP_PREFIX}{short}_{_sanitize_tool_name(tool.name)}"
|
||||
if len(ai_name) > 64: # OpenAI function name 长度上限
|
||||
ai_name = ai_name[:64]
|
||||
ctx.ai_tools.append({
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": ai_name,
|
||||
"description": tool.description or f"外部 MCP 工具({row.name} / {tool.name})",
|
||||
"parameters": tool.input_schema or {"type": "object", "properties": {}},
|
||||
},
|
||||
})
|
||||
ctx.names.append(ai_name)
|
||||
if not ctx.ai_tools:
|
||||
return None
|
||||
return ctx
|
||||
|
||||
|
||||
async def execute_mcp_tool(ctx: McpContext, ai_tool_name: str, args: object) -> str:
|
||||
"""执行一个 mcp_ 前缀工具;失败返回错误串,绝不断流。"""
|
||||
try:
|
||||
server, tool_name = ctx.resolve(ai_tool_name)
|
||||
except McpClientError as e:
|
||||
return f"[mcp] {e}"
|
||||
if not server.enabled:
|
||||
return f"[mcp] 服务器已停用:{server.name}"
|
||||
try:
|
||||
return await call_mcp_tool(server, tool_name, args)
|
||||
except Exception as e:
|
||||
msg = str(e) if not isinstance(e, McpClientError) else str(e)
|
||||
return f"[mcp:{tool_name}] 调用失败:{msg}"
|
||||
|
||||
|
||||
async def agent_stream(
|
||||
messages: list[dict],
|
||||
provider: LLMProvider,
|
||||
tools: list[dict],
|
||||
execute: Callable[[str, object], Awaitable[str]],
|
||||
describe: Callable[[str], str] | None = None,
|
||||
) -> AsyncGenerator[tuple[str, str], None]:
|
||||
"""流式 agent 循环:每轮携带 tools 请求模型,出现工具调用则执行并续一轮。
|
||||
|
||||
产出 (kind, payload) 事件:
|
||||
- ("token", 文本)
|
||||
- ("tool", "<服务器>/<工具>") —— 工具调用开始
|
||||
- ("tool_done", "<服务器>/<工具>") —— 工具执行完成
|
||||
"""
|
||||
for _round in range(MAX_TOOL_ROUNDS):
|
||||
provider.tool_calls = []
|
||||
provider.last_content = ""
|
||||
async for token in provider.chat_stream(messages, temperature=0.7, tools=tools):
|
||||
yield ("token", token)
|
||||
|
||||
calls: list[dict] = getattr(provider, "tool_calls", []) or []
|
||||
if not calls:
|
||||
return
|
||||
|
||||
assistant_msg: dict = {"role": "assistant", "content": getattr(provider, "last_content", "") or None}
|
||||
for call in calls:
|
||||
name = (call.get("function") or {}).get("name") or ""
|
||||
call_id = call.get("id") or f"call_{name}_{_round}"
|
||||
assistant_msg.setdefault("tool_calls", []).append({
|
||||
"id": call_id,
|
||||
"type": "function",
|
||||
"function": {"name": name, "arguments": (call.get("function") or {}).get("arguments") or "{}"},
|
||||
})
|
||||
messages.append(assistant_msg)
|
||||
|
||||
for call in assistant_msg["tool_calls"]:
|
||||
name = call["function"]["name"]
|
||||
display = describe(name) if describe else name
|
||||
try:
|
||||
args = json.loads(call["function"]["arguments"] or "{}")
|
||||
if not isinstance(args, dict):
|
||||
args = {"value": args}
|
||||
except (json.JSONDecodeError, TypeError):
|
||||
args = {}
|
||||
yield ("tool", display)
|
||||
out = await execute(name, args)
|
||||
yield ("tool_done", display)
|
||||
messages.append({
|
||||
"role": "tool",
|
||||
"tool_call_id": call["id"],
|
||||
"content": (out or "[工具无返回内容]")[:TOOL_RESULT_CHARS],
|
||||
})
|
||||
# 达到最大轮数仍未见收尾:用已收集上下文走无工具流式生成,保证必出最终答案
|
||||
async for token in provider.chat_stream(messages, temperature=0.7):
|
||||
yield ("token", token)
|
||||
@@ -0,0 +1,523 @@
|
||||
"""client — 最小 MCP 客户端(httpx 手写协议实现,零 MCP 第三方依赖)
|
||||
|
||||
支持两种远程传输:
|
||||
- Streamable HTTP:POST JSON-RPC 到 cfg.url,会话经 `Mcp-Session-Id` 响应头下发后回带。
|
||||
- 老式 SSE:GET /sse 端点,等 `event: endpoint` 事件拿到 messages URL,再 POST JSON-RPC,
|
||||
响应经原 SSE 流按请求 id 关联回来。
|
||||
|
||||
会话 / 工具列表按服务器 id 内存缓存(TTL 5 分钟),配置变更时用 invalidate_mcp_cache 失效。
|
||||
任何失败抛 McpClientError(消息为中文、可直接透传给用户),由上层捕获转错误串,绝不阻断对话。
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import itertools
|
||||
import json
|
||||
import time
|
||||
from dataclasses import dataclass, field
|
||||
from typing import AsyncGenerator, NamedTuple
|
||||
|
||||
import httpx
|
||||
|
||||
MCP_PROTOCOL_VERSION = "2025-06-18"
|
||||
CLIENT_NAME = "bookmate"
|
||||
CLIENT_VERSION = "1.0.0"
|
||||
|
||||
REQUEST_TIMEOUT_MS = 15_000
|
||||
CACHE_TTL_S = 5 * 60
|
||||
|
||||
_rpc_seq = itertools.count(1)
|
||||
|
||||
|
||||
class McpClientError(RuntimeError):
|
||||
"""MCP 调用失败(消息可直接展示给用户)。"""
|
||||
|
||||
|
||||
@dataclass
|
||||
class McpTool:
|
||||
"""服务器暴露的一个 MCP 工具(input_schema 本就是 JSON Schema)。"""
|
||||
|
||||
name: str
|
||||
description: str = ""
|
||||
input_schema: dict | None = None
|
||||
|
||||
|
||||
@dataclass
|
||||
class McpServerConfig:
|
||||
"""一个已配置的 MCP 服务器(与 DB 行一一对应)。"""
|
||||
|
||||
id: str
|
||||
name: str
|
||||
transport: str = "streamable-http"
|
||||
url: str = ""
|
||||
headers: dict[str, str] = field(default_factory=dict)
|
||||
enabled: bool = True
|
||||
|
||||
|
||||
@dataclass
|
||||
class McpCallResult:
|
||||
text: str
|
||||
is_error: bool = False
|
||||
|
||||
|
||||
def next_rpc_id() -> int:
|
||||
return next(_rpc_seq)
|
||||
|
||||
|
||||
def _resolve_url(url: str, base: str) -> str:
|
||||
"""把(可能相对)的端点 URL 基于服务器 URL 解析为绝对 URL。"""
|
||||
from urllib.parse import urljoin
|
||||
|
||||
try:
|
||||
return urljoin(base, url)
|
||||
except Exception:
|
||||
return url
|
||||
|
||||
|
||||
def build_headers(cfg: McpServerConfig, extra: dict[str, str]) -> dict[str, str]:
|
||||
"""合并自定义请求头 + 公共头;空值头会被过滤。"""
|
||||
headers: dict[str, str] = {}
|
||||
for k, v in cfg.headers.items():
|
||||
if k and v:
|
||||
headers[k] = v
|
||||
for k, v in extra.items():
|
||||
if v:
|
||||
headers[k] = v
|
||||
headers.setdefault("Accept", "application/json, text/event-stream")
|
||||
return headers
|
||||
|
||||
|
||||
def expect_result(cfg: McpServerConfig, method: str, res: dict | None) -> dict:
|
||||
if not res:
|
||||
raise McpClientError(f"[{cfg.name}] {method} 无响应")
|
||||
if res.get("error"):
|
||||
err = res["error"]
|
||||
raise McpClientError(f"[{cfg.name}] {method} 失败:{err.get('message')}(code {err.get('code')})")
|
||||
result = res.get("result")
|
||||
if not isinstance(result, dict):
|
||||
raise McpClientError(f"[{cfg.name}] {method} 返回异常 result")
|
||||
return result
|
||||
|
||||
|
||||
def _friendly_http_error(cfg: McpServerConfig, exc: httpx.HTTPError) -> McpClientError:
|
||||
detail = getattr(exc, "request", None)
|
||||
if isinstance(exc, httpx.TimeoutException):
|
||||
return McpClientError(f"[{cfg.name}] 请求超时:{detail.url if detail else cfg.url}")
|
||||
if isinstance(exc, httpx.ConnectError):
|
||||
return McpClientError(f"[{cfg.name}] 连接失败:{detail.url if detail else cfg.url}")
|
||||
return McpClientError(f"[{cfg.name}] 网络错误:{exc}")
|
||||
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# SSE 帧解析
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
class SseEvent(NamedTuple):
|
||||
event: str
|
||||
data: str
|
||||
|
||||
|
||||
async def read_sse_events(resp: httpx.Response) -> AsyncGenerator[SseEvent, None]:
|
||||
"""把响应体流解析为 SSE 事件序列(逐行、空行分帧)。"""
|
||||
event_name = ""
|
||||
data_lines: list[str] = []
|
||||
async for line in resp.aiter_lines():
|
||||
if line == "":
|
||||
if data_lines:
|
||||
yield SseEvent(event=event_name or "message", data="\n".join(data_lines))
|
||||
event_name = ""
|
||||
data_lines = []
|
||||
elif line.startswith(":"):
|
||||
continue
|
||||
elif line.startswith("event:"):
|
||||
event_name = line[6:].strip()
|
||||
elif line.startswith("data:"):
|
||||
data_lines.append(line[5:].lstrip())
|
||||
if data_lines:
|
||||
yield SseEvent(event=event_name or "message", data="\n".join(data_lines))
|
||||
|
||||
|
||||
def parse_jsonrpc(data: str) -> dict | None:
|
||||
try:
|
||||
obj = json.loads(data)
|
||||
if isinstance(obj, dict) and "jsonrpc" in obj:
|
||||
return obj
|
||||
except (json.JSONDecodeError, TypeError):
|
||||
pass
|
||||
return None
|
||||
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Streamable HTTP 传输
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
class StreamableHttpTransport:
|
||||
"""每个请求独立 POST,会话经 sessionId 头维持。"""
|
||||
|
||||
def __init__(self, cfg: McpServerConfig):
|
||||
self.cfg = cfg
|
||||
self.session_id: str | None = None
|
||||
self.client = httpx.AsyncClient(timeout=REQUEST_TIMEOUT_MS / 1000, follow_redirects=True)
|
||||
|
||||
async def close(self) -> None:
|
||||
await self.client.aclose()
|
||||
|
||||
async def post(self, body: dict) -> dict | None:
|
||||
headers = build_headers(self.cfg, {
|
||||
"Content-Type": "application/json",
|
||||
"MCP-Protocol-Version": MCP_PROTOCOL_VERSION,
|
||||
**(self.session_id and {"Mcp-Session-Id": self.session_id} or {}),
|
||||
})
|
||||
try:
|
||||
resp = await self.client.post(self.cfg.url, headers=headers, json=body)
|
||||
except httpx.HTTPError as e:
|
||||
raise _friendly_http_error(self.cfg, e) from e
|
||||
if resp.status_code >= 400:
|
||||
raise McpClientError(f"[{self.cfg.name}] HTTP {resp.status_code} {resp.reason_phrase}")
|
||||
sid = resp.headers.get("mcp-session-id")
|
||||
if sid:
|
||||
self.session_id = sid
|
||||
ctype = (resp.headers.get("content-type") or "").lower()
|
||||
if "application/json" in ctype:
|
||||
if not resp.content: # 通知等无响应体的 200/202
|
||||
return None
|
||||
data = resp.json()
|
||||
return data if isinstance(data, dict) else None
|
||||
# SSE 响应体:读完整条响应,按请求 id 取对应消息
|
||||
async for ev in read_sse_events(resp):
|
||||
parsed = parse_jsonrpc(ev.data)
|
||||
if parsed and parsed.get("id") == body.get("id"):
|
||||
return parsed
|
||||
return None
|
||||
|
||||
async def init(self) -> None:
|
||||
res = await self.post({
|
||||
"jsonrpc": "2.0", "id": next_rpc_id(), "method": "initialize",
|
||||
"params": {
|
||||
"protocolVersion": MCP_PROTOCOL_VERSION,
|
||||
"capabilities": {},
|
||||
"clientInfo": {"name": CLIENT_NAME, "version": CLIENT_VERSION},
|
||||
},
|
||||
})
|
||||
expect_result(self.cfg, "initialize", res)
|
||||
try: # 通知(无 id,忽略错误)
|
||||
await self.post({"jsonrpc": "2.0", "method": "notifications/initialized"})
|
||||
except McpClientError:
|
||||
pass
|
||||
|
||||
async def list_tools(self) -> list[McpTool]:
|
||||
if not self.session_id:
|
||||
await self.init()
|
||||
body = {"jsonrpc": "2.0", "id": next_rpc_id(), "method": "tools/list", "params": {}}
|
||||
res = await self.post(body)
|
||||
result = expect_result(self.cfg, "tools/list", res)
|
||||
tools = [McpTool(name=t.get("name", ""), description=t.get("description") or "", input_schema=t.get("inputSchema"))
|
||||
for t in result.get("tools") or [] if t.get("name")]
|
||||
cursor = result.get("nextCursor")
|
||||
while cursor: # 分页续拉
|
||||
page = await self.post({
|
||||
"jsonrpc": "2.0", "id": next_rpc_id(), "method": "tools/list",
|
||||
"params": {"cursor": cursor},
|
||||
})
|
||||
page_result = expect_result(self.cfg, "tools/list", page)
|
||||
tools.extend(McpTool(name=t.get("name", ""), description=t.get("description") or "", input_schema=t.get("inputSchema"))
|
||||
for t in page_result.get("tools") or [] if t.get("name"))
|
||||
cursor = page_result.get("nextCursor")
|
||||
return tools
|
||||
|
||||
async def call_tool(self, name: str, args: object) -> McpCallResult:
|
||||
if not self.session_id:
|
||||
await self.init()
|
||||
res = await self.post({
|
||||
"jsonrpc": "2.0", "id": next_rpc_id(), "method": "tools/call",
|
||||
"params": {"name": name, "arguments": args if isinstance(args, dict) else {}},
|
||||
})
|
||||
result = expect_result(self.cfg, f"tools/call {name}", res)
|
||||
return McpCallResult(text=format_tool_content(result.get("content")), is_error=result.get("isError") is True)
|
||||
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# 老式 SSE 传输:GET /sse 建连,POST messages URL,响应经 SSE 流按 id 关联
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
class LegacySseTransport:
|
||||
"""后台 reader 任务持续消费 SSE 流,把 id 匹配的 JSON-RPC 响应派发给挂起的请求。"""
|
||||
|
||||
def __init__(self, cfg: McpServerConfig):
|
||||
self.cfg = cfg
|
||||
self._pending: dict[int, asyncio.Future] = {}
|
||||
self._lock = asyncio.Lock()
|
||||
self._conn_ready = asyncio.Event()
|
||||
self._client: httpx.AsyncClient | None = None
|
||||
self._reader_task: asyncio.Task | None = None
|
||||
self._endpoint_url: str | None = None
|
||||
self._session_id: str | None = None
|
||||
self._error: McpClientError | None = None
|
||||
|
||||
async def close(self) -> None:
|
||||
if self._reader_task:
|
||||
self._reader_task.cancel()
|
||||
try:
|
||||
await self._reader_task
|
||||
except (asyncio.CancelledError, Exception):
|
||||
pass
|
||||
self._reader_task = None
|
||||
if self._client:
|
||||
await self._client.aclose()
|
||||
self._client = None
|
||||
for fut in self._pending.values():
|
||||
if not fut.done():
|
||||
fut.cancel()
|
||||
self._pending.clear()
|
||||
|
||||
async def _wait_ready(self) -> None:
|
||||
try:
|
||||
await asyncio.wait_for(self._conn_ready.wait(), timeout=REQUEST_TIMEOUT_MS / 1000)
|
||||
except asyncio.TimeoutError:
|
||||
raise McpClientError(f"[{self.cfg.name}] SSE 连接超时({REQUEST_TIMEOUT_MS}ms)") from None
|
||||
if self._error:
|
||||
raise self._error
|
||||
|
||||
async def _ensure_connected(self) -> None:
|
||||
if self._conn_ready.is_set():
|
||||
return
|
||||
if self._reader_task is not None:
|
||||
await self._wait_ready()
|
||||
return
|
||||
async with self._lock:
|
||||
if self._reader_task is not None:
|
||||
await self._wait_ready()
|
||||
return
|
||||
self._client = httpx.AsyncClient(timeout=None, follow_redirects=True)
|
||||
try:
|
||||
# 必须流式取响应:SSE 连接长期不关闭,非流式 get() 会一直等到 EOF
|
||||
resp = await self._client.send(
|
||||
httpx.Request("GET", self.cfg.url, headers=build_headers(self.cfg, {})),
|
||||
stream=True,
|
||||
)
|
||||
except httpx.HTTPError as e:
|
||||
self._error = _friendly_http_error(self.cfg, e)
|
||||
self._conn_ready.set()
|
||||
raise self._error from e
|
||||
if resp.status_code >= 400:
|
||||
self._error = McpClientError(f"[{self.cfg.name}] HTTP {resp.status_code} {resp.reason_phrase}")
|
||||
self._conn_ready.set()
|
||||
raise self._error
|
||||
self._reader_task = asyncio.create_task(self._read_loop(resp))
|
||||
await self._wait_ready() # 等 endpoint 事件就绪后再继续
|
||||
|
||||
async def _read_loop(self, resp: httpx.Response) -> None:
|
||||
"""解析 SSE 流:先等 endpoint 事件,之后把 JSON-RPC 响应分发给挂起 Future。"""
|
||||
try:
|
||||
async for ev in read_sse_events(resp):
|
||||
if self._endpoint_url is None:
|
||||
if ev.event == "endpoint" and ev.data.strip():
|
||||
endpoint = ev.data.strip()
|
||||
self._endpoint_url = _resolve_url(endpoint, self.cfg.url)
|
||||
try:
|
||||
u = httpx.URL(self._endpoint_url)
|
||||
self._session_id = u.params.get("sessionId") or u.params.get("session_id")
|
||||
except Exception:
|
||||
pass
|
||||
self._conn_ready.set()
|
||||
continue
|
||||
parsed = parse_jsonrpc(ev.data)
|
||||
if parsed and isinstance(parsed.get("id"), int) and parsed["id"] in self._pending:
|
||||
fut = self._pending.pop(parsed["id"])
|
||||
if not fut.done():
|
||||
fut.set_result(parsed)
|
||||
except asyncio.CancelledError:
|
||||
pass
|
||||
except Exception as e: # 连接中断 → 所有挂起请求报错
|
||||
self._error = McpClientError(f"[{self.cfg.name}] SSE 连接中断:{e}")
|
||||
finally:
|
||||
if self._endpoint_url is None:
|
||||
if self._error is None:
|
||||
self._error = McpClientError(f"[{self.cfg.name}] SSE 端点未返回 endpoint 事件")
|
||||
self._conn_ready.set()
|
||||
for fut in self._pending.values():
|
||||
if not fut.done():
|
||||
fut.set_exception(self._error or McpClientError(f"[{self.cfg.name}] SSE 连接已关闭"))
|
||||
self._pending.clear()
|
||||
|
||||
async def post(self, body: dict) -> dict | None:
|
||||
await self._ensure_connected()
|
||||
if not self._endpoint_url:
|
||||
raise self._error or McpClientError(f"[{self.cfg.name}] SSE 端点未就绪")
|
||||
headers = build_headers(self.cfg, {
|
||||
"Content-Type": "application/json",
|
||||
"MCP-Protocol-Version": MCP_PROTOCOL_VERSION,
|
||||
**(self._session_id and {"Mcp-Session-Id": self._session_id} or {}),
|
||||
})
|
||||
|
||||
rid: int = body.get("id")
|
||||
if rid is not None:
|
||||
fut: asyncio.Future = asyncio.get_running_loop().create_future()
|
||||
self._pending[rid] = fut
|
||||
try:
|
||||
try:
|
||||
resp = await self._client.post(self._endpoint_url, headers=headers, json=body)
|
||||
except httpx.HTTPError as e:
|
||||
raise _friendly_http_error(self.cfg, e) from e
|
||||
# 部分服务器直接在 POST 响应体返回 JSON
|
||||
if resp.status_code >= 400:
|
||||
raise McpClientError(f"[{self.cfg.name}] HTTP {resp.status_code} {resp.reason_phrase}")
|
||||
if "application/json" in (resp.headers.get("content-type") or "").lower():
|
||||
data = resp.json() if resp.content else None
|
||||
if rid is not None and self._pending.get(rid) is fut:
|
||||
self._pending.pop(rid, None)
|
||||
return data if isinstance(data, dict) else None
|
||||
except Exception:
|
||||
if rid is not None and self._pending.get(rid) is fut:
|
||||
self._pending.pop(rid, None)
|
||||
raise
|
||||
# 通知(无 id)没有响应:直接返回
|
||||
if rid is None:
|
||||
return None
|
||||
try:
|
||||
return await asyncio.wait_for(asyncio.shield(fut), timeout=REQUEST_TIMEOUT_MS / 1000)
|
||||
except asyncio.TimeoutError:
|
||||
self._pending.pop(rid, None)
|
||||
raise McpClientError(f"[{self.cfg.name}] 请求超时({REQUEST_TIMEOUT_MS}ms)") from None
|
||||
|
||||
async def init(self) -> None:
|
||||
res = await self.post({
|
||||
"jsonrpc": "2.0", "id": next_rpc_id(), "method": "initialize",
|
||||
"params": {
|
||||
"protocolVersion": MCP_PROTOCOL_VERSION,
|
||||
"capabilities": {},
|
||||
"clientInfo": {"name": CLIENT_NAME, "version": CLIENT_VERSION},
|
||||
},
|
||||
})
|
||||
expect_result(self.cfg, "initialize", res)
|
||||
try:
|
||||
await self.post({"jsonrpc": "2.0", "method": "notifications/initialized"})
|
||||
except McpClientError:
|
||||
pass
|
||||
|
||||
async def list_tools(self) -> list[McpTool]:
|
||||
if not self._endpoint_url:
|
||||
await self.init()
|
||||
body = {"jsonrpc": "2.0", "id": next_rpc_id(), "method": "tools/list", "params": {}}
|
||||
res = await self.post(body)
|
||||
result = expect_result(self.cfg, "tools/list", res)
|
||||
tools = [McpTool(name=t.get("name", ""), description=t.get("description") or "", input_schema=t.get("inputSchema"))
|
||||
for t in result.get("tools") or [] if t.get("name")]
|
||||
cursor = result.get("nextCursor")
|
||||
while cursor:
|
||||
page = await self.post({
|
||||
"jsonrpc": "2.0", "id": next_rpc_id(), "method": "tools/list",
|
||||
"params": {"cursor": cursor},
|
||||
})
|
||||
page_result = expect_result(self.cfg, "tools/list", page)
|
||||
tools.extend(McpTool(name=t.get("name", ""), description=t.get("description") or "", input_schema=t.get("inputSchema"))
|
||||
for t in page_result.get("tools") or [] if t.get("name"))
|
||||
cursor = page_result.get("nextCursor")
|
||||
return tools
|
||||
|
||||
async def call_tool(self, name: str, args: object) -> McpCallResult:
|
||||
if not self._endpoint_url:
|
||||
await self.init()
|
||||
res = await self.post({
|
||||
"jsonrpc": "2.0", "id": next_rpc_id(), "method": "tools/call",
|
||||
"params": {"name": name, "arguments": args if isinstance(args, dict) else {}},
|
||||
})
|
||||
result = expect_result(self.cfg, f"tools/call {name}", res)
|
||||
return McpCallResult(text=format_tool_content(result.get("content")), is_error=result.get("isError") is True)
|
||||
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# 结果格式化
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def format_tool_content(content: object) -> str:
|
||||
if not isinstance(content, list):
|
||||
return "[工具无返回内容]"
|
||||
parts: list[str] = []
|
||||
for item in content:
|
||||
if not isinstance(item, dict):
|
||||
continue
|
||||
t = item.get("type")
|
||||
if t == "text" and item.get("text"):
|
||||
parts.append(str(item["text"]))
|
||||
elif t == "image":
|
||||
data = item.get("data") or ""
|
||||
parts.append(f"[图片 {item.get('mimeType') or '未知类型'},base64 {len(data)} 字符]")
|
||||
elif t == "audio":
|
||||
parts.append(f"[音频 {item.get('mimeType') or '未知类型'}]")
|
||||
elif t == "resource":
|
||||
res = item.get("resource") or {}
|
||||
parts.append(res.get("text") or (f"[资源 {res.get('uri')}]" if res.get("uri") else "[资源]"))
|
||||
elif item.get("text"):
|
||||
parts.append(str(item["text"]))
|
||||
return "\n".join(p for p in parts if p) or "[工具无返回内容]"
|
||||
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# 缓存与高层 API
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
_transports: dict[str, StreamableHttpTransport | LegacySseTransport] = {}
|
||||
_tool_cache: dict[str, tuple[float, list[McpTool]]] = {}
|
||||
_list_locks: dict[str, asyncio.Lock] = {}
|
||||
|
||||
|
||||
def _get_transport(cfg: McpServerConfig) -> StreamableHttpTransport | LegacySseTransport:
|
||||
transport = _transports.get(cfg.id)
|
||||
if transport is None:
|
||||
transport = LegacySseTransport(cfg) if cfg.transport == "sse" else StreamableHttpTransport(cfg)
|
||||
_transports[cfg.id] = transport
|
||||
return transport
|
||||
|
||||
|
||||
async def invalidate_mcp_cache(server_id: str | None = None) -> None:
|
||||
"""失效某个(或全部)服务器的缓存;配置变更时调用。"""
|
||||
ids = [server_id] if server_id else list(_transports.keys())
|
||||
|
||||
async def _evict(sid: str) -> None:
|
||||
transport = _transports.pop(sid, None)
|
||||
if transport is not None:
|
||||
try:
|
||||
await transport.close()
|
||||
except Exception:
|
||||
pass
|
||||
_tool_cache.pop(sid, None)
|
||||
_list_locks.pop(sid, None)
|
||||
|
||||
for sid in list(ids):
|
||||
await _evict(sid)
|
||||
|
||||
|
||||
async def list_mcp_tools(cfg: McpServerConfig, fresh: bool = False) -> list[McpTool]:
|
||||
"""获取服务器的工具列表(带缓存)。fresh 为 true 时先失效缓存重新拉取(用于「测试连接」)。"""
|
||||
if fresh:
|
||||
await invalidate_mcp_cache(cfg.id)
|
||||
cached = _tool_cache.get(cfg.id)
|
||||
if not fresh and cached and cached[0] > time.time():
|
||||
return cached[1]
|
||||
|
||||
lock = _list_locks.setdefault(cfg.id, asyncio.Lock())
|
||||
async with lock:
|
||||
cached = _tool_cache.get(cfg.id)
|
||||
if not fresh and cached and cached[0] > time.time():
|
||||
return cached[1]
|
||||
transport = _get_transport(cfg)
|
||||
try:
|
||||
tools = await transport.list_tools()
|
||||
except Exception:
|
||||
await invalidate_mcp_cache(cfg.id) # 失败时释放传输,避免坏会话复用于下一次
|
||||
raise
|
||||
_tool_cache[cfg.id] = (time.time() + CACHE_TTL_S, tools)
|
||||
return tools
|
||||
|
||||
|
||||
async def call_mcp_tool(cfg: McpServerConfig, tool_name: str, args: object) -> str:
|
||||
"""调用一个 MCP 工具;返回可读文本。失败抛 McpClientError 由上层捕获。"""
|
||||
transport = _get_transport(cfg)
|
||||
try:
|
||||
result = await transport.call_tool(tool_name, args)
|
||||
except Exception:
|
||||
await invalidate_mcp_cache(cfg.id) # 会话可能已失效 → 清缓存,允许下次重建
|
||||
raise
|
||||
return f"[MCP 工具 {tool_name} 返回错误]\n{result.text}" if result.is_error else result.text
|
||||
@@ -0,0 +1,36 @@
|
||||
import json
|
||||
import uuid
|
||||
from datetime import datetime
|
||||
from sqlalchemy import String, Text, Boolean, DateTime, func
|
||||
from sqlalchemy.orm import Mapped, mapped_column
|
||||
|
||||
from app.core.database import Base
|
||||
|
||||
|
||||
def gen_uuid() -> str:
|
||||
return str(uuid.uuid4())
|
||||
|
||||
|
||||
class McpServer(Base):
|
||||
"""用户从页面配置的 MCP 服务器(供对话 agent 调用外部工具)。
|
||||
|
||||
存放远程 MCP 服务的连接信息:传输协议、URL、附加请求头(含鉴权)。
|
||||
启用中的服务器会在对话时向模型暴露其工具(mcp_ 前缀)。
|
||||
"""
|
||||
|
||||
__tablename__ = "mcp_servers"
|
||||
|
||||
id: Mapped[str] = mapped_column(String(36), primary_key=True, default=gen_uuid)
|
||||
name: Mapped[str] = mapped_column(String(200)) # 显示名,如 firecrawl
|
||||
transport: Mapped[str] = mapped_column(String(30), default="streamable-http") # streamable-http | sse
|
||||
url: Mapped[str] = mapped_column(String(2000)) # 远程端点
|
||||
headers: Mapped[str] = mapped_column(Text, default="{}") # JSON dict,附加请求头(明文本地存储,与 ai_models 同级信任)
|
||||
enabled: Mapped[bool] = mapped_column(Boolean, default=True)
|
||||
created_at: Mapped[datetime] = mapped_column(DateTime, server_default=func.now())
|
||||
|
||||
def headers_dict(self) -> dict:
|
||||
try:
|
||||
data = json.loads(self.headers or "{}")
|
||||
return data if isinstance(data, dict) else {}
|
||||
except (json.JSONDecodeError, TypeError):
|
||||
return {}
|
||||
Reference in New Issue
Block a user