mirror of
https://github.com/FuQuan233/nonebot-plugin-llmchat.git
synced 2025-09-04 10:20:45 +00:00
85 lines
3.7 KiB
Python
85 lines
3.7 KiB
Python
from contextlib import AsyncExitStack
|
||
import asyncio
|
||
|
||
from mcp import ClientSession, StdioServerParameters
|
||
from mcp.client.sse import sse_client
|
||
from mcp.client.stdio import stdio_client
|
||
from nonebot import logger
|
||
|
||
from .config import MCPServerConfig
|
||
|
||
|
||
class MCPClient:
|
||
def __init__(self, server_config: dict[str, MCPServerConfig]):
|
||
logger.info(f"正在初始化MCPClient，共有{len(server_config)}个服务器配置")
|
||
self.server_config = server_config
|
||
self.sessions = {}
|
||
self.exit_stack = AsyncExitStack()
|
||
logger.debug("MCPClient初始化成功")
|
||
|
||
async def connect_to_servers(self):
|
||
logger.info(f"开始连接{len(self.server_config)}个MCP服务器")
|
||
for server_name, config in self.server_config.items():
|
||
logger.debug(f"正在连接服务器[{server_name}]")
|
||
if config.url:
|
||
sse_transport = await self.exit_stack.enter_async_context(sse_client(url=config.url))
|
||
read, write = sse_transport
|
||
self.sessions[server_name] = await self.exit_stack.enter_async_context(ClientSession(read, write))
|
||
await self.sessions[server_name].initialize()
|
||
elif config.command:
|
||
stdio_transport = await self.exit_stack.enter_async_context(
|
||
stdio_client(StdioServerParameters(**config.model_dump()))
|
||
)
|
||
read, write = stdio_transport
|
||
self.sessions[server_name] = await self.exit_stack.enter_async_context(ClientSession(read, write))
|
||
await self.sessions[server_name].initialize()
|
||
else:
|
||
raise ValueError("Server config must have either url or command")
|
||
|
||
logger.info(f"已成功连接到MCP服务器[{server_name}]")
|
||
|
||
async def get_available_tools(self):
|
||
logger.info(f"正在从{len(self.sessions)}个已连接的服务器获取可用工具")
|
||
available_tools = []
|
||
|
||
for server_name, session in self.sessions.items():
|
||
logger.debug(f"正在列出服务器[{server_name}]中的工具")
|
||
response = await session.list_tools()
|
||
tools = response.tools
|
||
logger.debug(f"在服务器[{server_name}]中找到{len(tools)}个工具")
|
||
|
||
available_tools.extend(
|
||
{
|
||
"type": "function",
|
||
"function": {
|
||
"name": f"{server_name}___{tool.name}",
|
||
"description": tool.description,
|
||
"parameters": tool.inputSchema,
|
||
},
|
||
}
|
||
for tool in tools
|
||
)
|
||
return available_tools
|
||
|
||
async def call_tool(self, tool_name: str, tool_args: dict):
|
||
server_name, real_tool_name = tool_name.split("___")
|
||
logger.info(f"正在服务器[{server_name}]上调用工具[{real_tool_name}]")
|
||
session = self.sessions[server_name]
|
||
try:
|
||
response = await asyncio.wait_for(
|
||
session.call_tool(real_tool_name, tool_args), timeout=10
|
||
)
|
||
except asyncio.TimeoutError:
|
||
logger.error(f"调用工具[{real_tool_name}]超时")
|
||
return f"调用工具[{real_tool_name}]超时"
|
||
logger.debug(f"工具[{real_tool_name}]调用完成，响应: {response}")
|
||
return response.content
|
||
|
||
def get_friendly_name(self, tool_name: str):
|
||
server_name, real_tool_name = tool_name.split("___")
|
||
return (self.server_config[server_name].friendly_name or server_name) + " - " + real_tool_name
|
||
|
||
async def cleanup(self):
|
||
logger.debug("正在清理MCPClient资源")
|
||
await self.exit_stack.aclose()
|
||
logger.debug("MCPClient资源清理完成")
|