74 lines
2.1 KiB
Python
74 lines
2.1 KiB
Python
from __future__ import annotations
|
||
|
||
from dataclasses import dataclass
|
||
from typing import Any
|
||
|
||
from autogen_ext.tools.mcp import StdioServerParams, mcp_server_tools
|
||
|
||
from . import config
|
||
|
||
|
||
@dataclass(frozen=True)
|
||
class MCPServerConfig:
|
||
name: str
|
||
command: str
|
||
args: list[str]
|
||
|
||
|
||
def _configured_servers() -> list[MCPServerConfig]:
|
||
if not config.ENABLE_MCP_TOOLS:
|
||
return []
|
||
|
||
servers: list[MCPServerConfig] = []
|
||
if config.MCP_WEATHER_SERVER_COMMAND:
|
||
servers.append(
|
||
MCPServerConfig(
|
||
name="weather",
|
||
command=config.MCP_WEATHER_SERVER_COMMAND,
|
||
args=config.MCP_WEATHER_SERVER_ARGS,
|
||
)
|
||
)
|
||
if config.MCP_WEBSEARCH_SERVER_COMMAND:
|
||
servers.append(
|
||
MCPServerConfig(
|
||
name="websearch",
|
||
command=config.MCP_WEBSEARCH_SERVER_COMMAND,
|
||
args=config.MCP_WEBSEARCH_SERVER_ARGS,
|
||
)
|
||
)
|
||
return servers
|
||
|
||
|
||
async def load_mcp_tools() -> list[Any]:
|
||
configured_servers = _configured_servers()
|
||
if not configured_servers:
|
||
print("ℹ️ MCP 工具未配置，跳过加载。")
|
||
return []
|
||
|
||
loaded_tools: list[Any] = []
|
||
tool_names: set[str] = set()
|
||
|
||
for server in configured_servers:
|
||
params = StdioServerParams(
|
||
command=server.command,
|
||
args=server.args,
|
||
read_timeout_seconds=config.MCP_SERVER_READ_TIMEOUT_SECONDS,
|
||
)
|
||
try:
|
||
server_tools = await mcp_server_tools(params)
|
||
for tool in server_tools:
|
||
if tool.name in tool_names:
|
||
print(f"⚠️ MCP 工具重名，已跳过: {tool.name}")
|
||
continue
|
||
loaded_tools.append(tool)
|
||
tool_names.add(tool.name)
|
||
print(f"✅ MCP 服务已加载: {server.name} ({len(server_tools)} tools)")
|
||
except Exception as exc:
|
||
print(f"⚠️ MCP 服务加载失败: {server.name}, error={exc}")
|
||
|
||
if loaded_tools:
|
||
print(f"✅ MCP 工具总数: {len(loaded_tools)}")
|
||
else:
|
||
print("ℹ️ 未加载到任何 MCP 工具。")
|
||
return loaded_tools
|