修复 has_tool错误;优化检索复杂度

This commit is contained in:
myifeng
2025-05-30 11:35:12 +08:00
parent a6deb3af8b
commit 43ce1df0ed
2 changed files with 34 additions and 24 deletions
+12 -6
View File
@@ -704,17 +704,23 @@ class ConnectionHandler:
"arguments": function_arguments, "arguments": function_arguments,
} }
# 处理MCP工具调用 # 处理Server端MCP工具调用
if self.mcp_manager.is_mcp_tool(function_name): if self.mcp_manager.is_mcp_tool(function_name):
result = self._handle_mcp_tool_call(function_call_data) result = self._handle_mcp_tool_call(function_call_data)
elif hasattr(self, "mcp_client") and self.mcp_client.has_tool(function_name): elif hasattr(self, "mcp_client") and self.mcp_client.has_tool(function_name):
# 如果是MCP工具调用 # 如果是小智端MCP工具调用
self.logger.bind(tag=TAG).debug( self.logger.bind(tag=TAG).debug(
f"调用MCP工具: {function_name}, 参数: {function_arguments}" f"调用小智端MCP工具: {function_name}, 参数: {function_arguments}"
) )
result = asyncio.run_coroutine_threadsafe(call_mcp_tool(self, self.mcp_client, function_name, function_arguments), self.loop).result() try:
self.logger.bind(tag=TAG).debug(f"MCP工具调用结果: {result}") result = asyncio.run_coroutine_threadsafe(call_mcp_tool(self, self.mcp_client, function_name, function_arguments), self.loop).result()
result = ActionResponse(action=Action.REQLLM, result=result, response="") self.logger.bind(tag=TAG).debug(f"MCP工具调用结果: {result}")
result = ActionResponse(action=Action.REQLLM, result=result, response="")
except Exception as e:
self.logger.bind(tag=TAG).error(f"MCP工具调用失败: {e}")
result = ActionResponse(
action=Action.REQLLM, result="MCP工具调用失败", response=""
)
else: else:
# 处理系统函数 # 处理系统函数
result = self.func_handler.handle_llm_function_call( result = self.func_handler.handle_llm_function_call(
+22 -18
View File
@@ -7,33 +7,36 @@ TAG = __name__
class MCPClient: class MCPClient:
"""MCPClient,用于管理MCP状态和工具""" """MCPClient,用于管理MCP状态和工具"""
def __init__(self): def __init__(self):
self.tools = [] self.tools = {} # Dictionary for O(1) lookup
self.ready = False self.ready = False
self.call_results = {} # To store Futures for tool call responses self.call_results = {} # To store Futures for tool call responses
self.next_id = 1 self.next_id = 1
self.lock = asyncio.Lock() self.lock = asyncio.Lock()
self._cached_available_tools = None # Cache for get_available_tools
async def has_tool(self, name: str) -> bool: def has_tool(self, name: str) -> bool:
async with self.lock: return name in self.tools
for tool in self.tools:
if tool["name"] == name:
return True
return False
def get_available_tools(self) -> list: def get_available_tools(self) -> list:
# async with self.lock: # Check if the cache is valid
if self._cached_available_tools is not None:
return self._cached_available_tools
# If cache is not valid, regenerate the list
result = [] result = []
for tool in self.tools: for tool_name, tool_data in self.tools.items():
function_def = { function_def = {
"name": tool["name"], "name": tool_data["name"],
"description": tool["description"], "description": tool_data["description"],
"parameters": { "parameters": {
"type": tool["inputSchema"].get("type", "object"), "type": tool_data["inputSchema"].get("type", "object"),
"properties": tool["inputSchema"].get("properties", {}), "properties": tool_data["inputSchema"].get("properties", {}),
"required": tool["inputSchema"].get("required", []) "required": tool_data["inputSchema"].get("required", [])
} }
} }
result.append({"type": "function", "function": function_def}) result.append({"type": "function", "function": function_def})
self._cached_available_tools = result # Store the generated list in cache
return result return result
async def is_ready(self) -> bool: async def is_ready(self) -> bool:
@@ -46,8 +49,9 @@ class MCPClient:
async def add_tool(self, tool_data: dict): async def add_tool(self, tool_data: dict):
async with self.lock: async with self.lock:
self.tools.append(tool_data) self.tools[tool_data["name"]] = tool_data
self._cached_available_tools = None # Invalidate the cache when a tool is added
async def get_next_id(self) -> int: async def get_next_id(self) -> int:
async with self.lock: async with self.lock:
current_id = self.next_id current_id = self.next_id
@@ -232,7 +236,7 @@ async def call_mcp_tool(conn, mcp_client: MCPClient, tool_name: str, args: str =
if not await mcp_client.is_ready(): if not await mcp_client.is_ready():
raise RuntimeError("MCP客户端尚未准备就绪") raise RuntimeError("MCP客户端尚未准备就绪")
if not await mcp_client.has_tool(tool_name): if not mcp_client.has_tool(tool_name):
raise ValueError(f"工具 {tool_name} 不存在") raise ValueError(f"工具 {tool_name} 不存在")
tool_call_id = await mcp_client.get_next_id() tool_call_id = await mcp_client.get_next_id()
@@ -272,4 +276,4 @@ async def call_mcp_tool(conn, mcp_client: MCPClient, tool_name: str, args: str =
raise TimeoutError("工具调用请求超时") raise TimeoutError("工具调用请求超时")
except Exception as e: except Exception as e:
await mcp_client.cleanup_call_result(tool_call_id) await mcp_client.cleanup_call_result(tool_call_id)
raise e raise e