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