mirror of
https://github.com/xinnan-tech/xiaozhi-esp32-server.git
synced 2026-07-22 07:03:53 +08:00
Merge pull request #2571 from xinnan-tech/add-mcp-callback
Add ProgressFnT LoggingFnT...
This commit is contained in:
@@ -10,10 +10,13 @@ import concurrent.futures
|
||||
from contextlib import AsyncExitStack
|
||||
from typing import Optional, List, Dict, Any
|
||||
|
||||
from mcp import ClientSession, StdioServerParameters
|
||||
from mcp import ClientSession, StdioServerParameters, Implementation
|
||||
from mcp.client.session import SamplingFnT, ElicitationFnT, ListRootsFnT, LoggingFnT, MessageHandlerFnT
|
||||
from mcp.client.stdio import stdio_client
|
||||
from mcp.client.sse import sse_client
|
||||
from mcp.client.streamable_http import streamablehttp_client
|
||||
from mcp.shared.session import ProgressFnT
|
||||
|
||||
from config.logger import setup_logging
|
||||
from core.utils.util import sanitize_tool_name
|
||||
|
||||
@@ -41,13 +44,25 @@ class ServerMCPClient:
|
||||
self.tools_dict: Dict[str, Any] = {}
|
||||
self.name_mapping: Dict[str, str] = {}
|
||||
|
||||
async def initialize(self):
|
||||
async def initialize(self, read_timeout_seconds: timedelta | None = None,
|
||||
sampling_callback: SamplingFnT | None = None,
|
||||
elicitation_callback: ElicitationFnT | None = None,
|
||||
list_roots_callback: ListRootsFnT | None = None,
|
||||
logging_callback: LoggingFnT | None = None,
|
||||
message_handler: MessageHandlerFnT | None = None,
|
||||
client_info: Implementation | None = None):
|
||||
"""初始化MCP客户端连接"""
|
||||
if self._worker_task:
|
||||
return
|
||||
|
||||
self._worker_task = asyncio.create_task(
|
||||
self._worker(), name="ServerMCPClientWorker"
|
||||
self._worker(read_timeout_seconds=read_timeout_seconds,
|
||||
sampling_callback=sampling_callback,
|
||||
elicitation_callback=elicitation_callback,
|
||||
list_roots_callback=list_roots_callback,
|
||||
logging_callback=logging_callback,
|
||||
message_handler=message_handler,
|
||||
client_info=client_info), name="ServerMCPClientWorker"
|
||||
)
|
||||
await self._ready_evt.wait()
|
||||
|
||||
@@ -97,12 +112,15 @@ class ServerMCPClient:
|
||||
for name, tool in self.tools_dict.items()
|
||||
]
|
||||
|
||||
async def call_tool(self, name: str, args: dict) -> Any:
|
||||
async def call_tool(self, name: str, arguments: dict, read_timeout_seconds: timedelta | None = None, progress_callback: ProgressFnT | None = None, *, meta: dict[str, Any] | None = None) -> Any:
|
||||
"""调用指定工具
|
||||
|
||||
Args:
|
||||
name: 工具名称
|
||||
args: 工具参数
|
||||
arguments: 工具参数
|
||||
read_timeout_seconds:
|
||||
progress_callback: 进度回调函数
|
||||
meta:
|
||||
|
||||
Returns:
|
||||
Any: 工具执行结果
|
||||
@@ -115,7 +133,7 @@ class ServerMCPClient:
|
||||
|
||||
real_name = self.name_mapping.get(name, name)
|
||||
loop = self._worker_task.get_loop()
|
||||
coro = self.session.call_tool(real_name, args)
|
||||
coro = self.session.call_tool(real_name, arguments=arguments, read_timeout_seconds=read_timeout_seconds, progress_callback=progress_callback, meta=meta)
|
||||
|
||||
if loop is asyncio.get_running_loop():
|
||||
return await coro
|
||||
@@ -144,7 +162,13 @@ class ServerMCPClient:
|
||||
# 所有检查都通过,连接正常
|
||||
return True
|
||||
|
||||
async def _worker(self):
|
||||
async def _worker(self, read_timeout_seconds: timedelta | None = None,
|
||||
sampling_callback: SamplingFnT | None = None,
|
||||
elicitation_callback: ElicitationFnT | None = None,
|
||||
list_roots_callback: ListRootsFnT | None = None,
|
||||
logging_callback: LoggingFnT | None = None,
|
||||
message_handler: MessageHandlerFnT | None = None,
|
||||
client_info: Implementation | None = None):
|
||||
"""MCP客户端工作协程"""
|
||||
async with AsyncExitStack() as stack:
|
||||
try:
|
||||
@@ -208,7 +232,13 @@ class ServerMCPClient:
|
||||
ClientSession(
|
||||
read_stream=read_stream,
|
||||
write_stream=write_stream,
|
||||
read_timeout_seconds=timedelta(seconds=15),
|
||||
read_timeout_seconds=read_timeout_seconds,
|
||||
sampling_callback=sampling_callback,
|
||||
elicitation_callback=elicitation_callback,
|
||||
list_roots_callback=list_roots_callback,
|
||||
logging_callback=logging_callback,
|
||||
message_handler=message_handler,
|
||||
client_info=client_info
|
||||
)
|
||||
)
|
||||
await self.session.initialize()
|
||||
|
||||
@@ -3,7 +3,14 @@
|
||||
import asyncio
|
||||
import os
|
||||
import json
|
||||
from datetime import timedelta
|
||||
from typing import Dict, Any, List
|
||||
|
||||
from mcp import Implementation
|
||||
from mcp.client.session import SamplingFnT, ElicitationFnT, ListRootsFnT, LoggingFnT, MessageHandlerFnT
|
||||
from mcp.shared.session import ProgressFnT
|
||||
from mcp.types import LoggingMessageNotificationParams
|
||||
|
||||
from config.config_loader import get_project_dir
|
||||
from config.logger import setup_logging
|
||||
from .mcp_client import ServerMCPClient
|
||||
@@ -56,7 +63,7 @@ class ServerMCPManager:
|
||||
# 初始化服务端MCP客户端
|
||||
logger.bind(tag=TAG).info(f"初始化服务端MCP客户端: {name}")
|
||||
client = ServerMCPClient(srv_config)
|
||||
await client.initialize()
|
||||
await client.initialize(logging_callback=self.logging_callback)
|
||||
self.clients[name] = client
|
||||
client_tools = client.get_available_tools()
|
||||
self.tools.extend(client_tools)
|
||||
@@ -109,7 +116,7 @@ class ServerMCPManager:
|
||||
# 带重试机制的工具调用
|
||||
for attempt in range(max_retries):
|
||||
try:
|
||||
return await target_client.call_tool(tool_name, arguments)
|
||||
return await target_client.call_tool(tool_name, arguments, progress_callback=self.progress_callback)
|
||||
except Exception as e:
|
||||
# 最后一次尝试失败时直接抛出异常
|
||||
if attempt == max_retries - 1:
|
||||
@@ -131,7 +138,7 @@ class ServerMCPManager:
|
||||
config = self.load_config()
|
||||
if client_name in config:
|
||||
client = ServerMCPClient(config[client_name])
|
||||
await client.initialize()
|
||||
await client.initialize(logging_callback=self.logging_callback)
|
||||
self.clients[client_name] = client
|
||||
target_client = client
|
||||
logger.bind(tag=TAG).info(
|
||||
@@ -159,3 +166,11 @@ class ServerMCPManager:
|
||||
except (asyncio.TimeoutError, Exception) as e:
|
||||
logger.bind(tag=TAG).error(f"关闭服务端MCP客户端 {name} 时出错: {e}")
|
||||
self.clients.clear()
|
||||
|
||||
# 可选回调方法
|
||||
|
||||
async def logging_callback(self, params: LoggingMessageNotificationParams):
|
||||
logger.bind(tag=TAG).info(f"[Server Log - {params.level.upper()}] {params.data}")
|
||||
|
||||
async def progress_callback(self, progress: float, total: float | None, message: str | None) -> None:
|
||||
logger.bind(tag=TAG).info(f"[Progress {progress}/{total}]: {message}")
|
||||
Reference in New Issue
Block a user