refactor(mcp): 移除全局MCP客户端缓存并优化工具加载逻辑
移除全局MCP客户端缓存,改为每次调用时创建新客户端实例 优化工具加载逻辑,添加服务器名称校验并改进错误处理 清理未使用的导入和冗余代码
This commit is contained in:
parent
63da8d9733
commit
165ec73ca8
@ -1,10 +1,8 @@
|
||||
import os
|
||||
import uuid
|
||||
from typing import Any, cast, Annotated
|
||||
from typing import Any, cast
|
||||
from pathlib import Path
|
||||
from datetime import datetime, timezone
|
||||
|
||||
from dataclasses import dataclass, field, fields
|
||||
from langchain_core.messages import AIMessage, ToolMessage
|
||||
from langgraph.graph import StateGraph, START, END
|
||||
from langgraph.runtime import Runtime
|
||||
@ -14,7 +12,6 @@ from langgraph.checkpoint.memory import InMemorySaver
|
||||
|
||||
from src import config as sys_config
|
||||
from src.utils import logger
|
||||
from src.agents.common.utils import get_cur_time_with_utc
|
||||
from src.agents.common.base import BaseAgent
|
||||
from src.agents.common.models import load_chat_model
|
||||
from src.agents.common.mcp import get_mcp_tools
|
||||
|
||||
@ -11,8 +11,7 @@ from langchain_mcp_adapters.client import ( # type: ignore[import-untyped]
|
||||
|
||||
from src.utils import logger
|
||||
|
||||
# Global MCP client and tools cache
|
||||
_mcp_client: MultiServerMCPClient | None = None
|
||||
# Global MCP tools cache
|
||||
_mcp_tools_cache: dict[str, list[Callable[..., Any]]] = {}
|
||||
|
||||
# MCP Server configurations
|
||||
@ -30,25 +29,22 @@ MCP_SERVERS = {
|
||||
"command": "uvx",
|
||||
"args": ["mcp-server-time"],
|
||||
"transport": "stdio",
|
||||
}
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
async def get_mcp_client(
|
||||
server_configs: dict[str, Any] | None = None,
|
||||
) -> MultiServerMCPClient | None:
|
||||
"""Get or initialize the global MCP client with given server configurations."""
|
||||
global _mcp_client
|
||||
|
||||
if _mcp_client is None:
|
||||
configs = server_configs or MCP_SERVERS
|
||||
try:
|
||||
_mcp_client = MultiServerMCPClient(configs) # pyright: ignore[reportArgumentType]
|
||||
logger.info(f"Initialized MCP client with servers: {list(configs.keys())}")
|
||||
except Exception as e:
|
||||
logger.error("Failed to initialize MCP client: %s", e)
|
||||
return None
|
||||
return _mcp_client
|
||||
"""Initializes an MCP client with the given server configurations."""
|
||||
configs = server_configs or MCP_SERVERS
|
||||
try:
|
||||
client = MultiServerMCPClient(configs) # pyright: ignore[reportArgumentType]
|
||||
logger.info(f"Initialized MCP client with servers: {list(configs.keys())}")
|
||||
return client
|
||||
except Exception as e:
|
||||
logger.error("Failed to initialize MCP client: %s", e)
|
||||
return None
|
||||
|
||||
|
||||
async def get_mcp_tools(server_name: str) -> list[Callable[..., Any]]:
|
||||
@ -60,9 +56,9 @@ async def get_mcp_tools(server_name: str) -> list[Callable[..., Any]]:
|
||||
return _mcp_tools_cache[server_name]
|
||||
|
||||
try:
|
||||
assert server_name in MCP_SERVERS, f"Server {server_name} not found in MCP_SERVERS"
|
||||
client = await get_mcp_client({server_name: MCP_SERVERS[server_name]})
|
||||
if client is None:
|
||||
_mcp_tools_cache[server_name] = []
|
||||
return []
|
||||
|
||||
# Get all tools and filter by server (if tools have server metadata)
|
||||
@ -72,9 +68,11 @@ async def get_mcp_tools(server_name: str) -> list[Callable[..., Any]]:
|
||||
_mcp_tools_cache[server_name] = tools
|
||||
logger.info(f"Loaded {len(tools)} tools from MCP server '{server_name}'")
|
||||
return tools
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to load tools from MCP server '{server_name}': %s\n{traceback.format_exc()}", e)
|
||||
_mcp_tools_cache[server_name] = []
|
||||
except AssertionError as e:
|
||||
logger.warning(f"Failed to load tools from MCP server '{server_name}': {e}")
|
||||
return []
|
||||
except Exception:
|
||||
logger.opt(exception=True).warning(f"Failed to load tools from MCP server '{server_name}'")
|
||||
return []
|
||||
|
||||
async def get_all_mcp_tools() -> list[Callable[..., Any]]:
|
||||
@ -94,7 +92,6 @@ def add_mcp_server(name: str, config: dict[str, Any]) -> None:
|
||||
|
||||
|
||||
def clear_mcp_cache() -> None:
|
||||
"""Clear the MCP client and tools cache (useful for testing)."""
|
||||
global _mcp_client, _mcp_tools_cache
|
||||
_mcp_client = None
|
||||
"""Clear the MCP tools cache (useful for testing)."""
|
||||
global _mcp_tools_cache
|
||||
_mcp_tools_cache = {}
|
||||
|
||||
Loading…
Reference in New Issue
Block a user