fix: sanitize MCP tool names for OpenAI

This commit is contained in:
whats2000
2025-06-07 13:46:35 +08:00
parent 2a212ae759
commit 4d0ddd7ff3
3 changed files with 30 additions and 13 deletions
+9 -5
View File
@@ -1,7 +1,7 @@
import json import json
import asyncio import asyncio
from concurrent.futures import Future from concurrent.futures import Future
from core.utils.util import get_vision_url from core.utils.util import get_vision_url, sanitize_tool_name
from core.utils.auth import AuthToken from core.utils.auth import AuthToken
TAG = __name__ TAG = __name__
@@ -11,7 +11,8 @@ class MCPClient:
"""MCPClient,用于管理MCP状态和工具""" """MCPClient,用于管理MCP状态和工具"""
def __init__(self): def __init__(self):
self.tools = {} # Dictionary for O(1) lookup self.tools = {} # sanitized_name -> tool_data
self.name_mapping = {}
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
@@ -30,7 +31,7 @@ class MCPClient:
result = [] result = []
for tool_name, tool_data in self.tools.items(): for tool_name, tool_data in self.tools.items():
function_def = { function_def = {
"name": tool_data["name"], "name": tool_name,
"description": tool_data["description"], "description": tool_data["description"],
"parameters": { "parameters": {
"type": tool_data["inputSchema"].get("type", "object"), "type": tool_data["inputSchema"].get("type", "object"),
@@ -53,7 +54,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[tool_data["name"]] = tool_data sanitized_name = sanitize_tool_name(tool_data["name"])
self.tools[sanitized_name] = tool_data
self.name_mapping[sanitized_name] = tool_data["name"]
self._cached_available_tools = ( self._cached_available_tools = (
None # Invalidate the cache when a tool is added None # Invalidate the cache when a tool is added
) )
@@ -333,11 +336,12 @@ async def call_mcp_tool(
raise ValueError(f"参数处理失败: {str(e)}") raise ValueError(f"参数处理失败: {str(e)}")
raise e raise e
actual_name = mcp_client.name_mapping.get(tool_name, tool_name)
payload = { payload = {
"jsonrpc": "2.0", "jsonrpc": "2.0",
"id": tool_call_id, "id": tool_call_id,
"method": "tools/call", "method": "tools/call",
"params": {"name": tool_name, "arguments": arguments}, "params": {"name": actual_name, "arguments": arguments},
} }
conn.logger.bind(tag=TAG).info( conn.logger.bind(tag=TAG).info(
+16 -8
View File
@@ -9,6 +9,7 @@ from mcp import ClientSession, StdioServerParameters
from mcp.client.stdio import stdio_client from mcp.client.stdio import stdio_client
from mcp.client.sse import sse_client from mcp.client.sse import sse_client
from config.logger import setup_logging from config.logger import setup_logging
from core.utils.util import sanitize_tool_name
TAG = __name__ TAG = __name__
@@ -23,7 +24,9 @@ class MCPClient:
self._shutdown_evt = asyncio.Event() self._shutdown_evt = asyncio.Event()
self.session: Optional[ClientSession] = None self.session: Optional[ClientSession] = None
self.tools: List = [] self.tools: List = [] # original tool objects
self.tools_dict: Dict[str, Any] = {}
self.name_mapping: Dict[str, str] = {}
async def initialize(self): async def initialize(self):
if self._worker_task: if self._worker_task:
@@ -32,7 +35,7 @@ class MCPClient:
await self._ready_evt.wait() await self._ready_evt.wait()
self.logger.bind(tag=TAG).info( self.logger.bind(tag=TAG).info(
f"Connected, tools = {[t.name for t in self.tools]}" f"Connected, tools = {[name for name in self.name_mapping.values()]}"
) )
async def cleanup(self): async def cleanup(self):
@@ -48,27 +51,28 @@ class MCPClient:
self._worker_task = None self._worker_task = None
def has_tool(self, name: str) -> bool: def has_tool(self, name: str) -> bool:
return any(t.name == name for t in self.tools) return name in self.tools_dict
def get_available_tools(self): def get_available_tools(self):
return [ return [
{ {
"type": "function", "type": "function",
"function": { "function": {
"name": t.name, "name": name,
"description": t.description, "description": tool.description,
"parameters": t.inputSchema, "parameters": tool.inputSchema,
}, },
} }
for t in self.tools for name, tool in self.tools_dict.items()
] ]
async def call_tool(self, name: str, args: dict): async def call_tool(self, name: str, args: dict):
if not self.session: if not self.session:
raise RuntimeError("MCPClient not initialized") raise RuntimeError("MCPClient not initialized")
real_name = self.name_mapping.get(name, name)
loop = self._worker_task.get_loop() loop = self._worker_task.get_loop()
coro = self.session.call_tool(name, args) coro = self.session.call_tool(real_name, args)
if loop is asyncio.get_running_loop(): if loop is asyncio.get_running_loop():
return await coro return await coro
@@ -123,6 +127,10 @@ class MCPClient:
# 获取工具 # 获取工具
self.tools = (await self.session.list_tools()).tools self.tools = (await self.session.list_tools()).tools
for t in self.tools:
sanitized = sanitize_tool_name(t.name)
self.tools_dict[sanitized] = t
self.name_mapping[sanitized] = t.name
self._ready_evt.set() self._ready_evt.set()
+5
View File
@@ -976,3 +976,8 @@ def is_valid_image_file(file_data: bytes) -> bool:
return True return True
return False return False
def sanitize_tool_name(name: str) -> str:
"""Sanitize tool names for OpenAI compatibility."""
return re.sub(r"[^a-zA-Z0-9_-]", "_", name)