Enhance AIClient and MCPServer to support tool registration with source tracking. Added logging for tool calls and improved error handling. Introduced methods for embedding token limit extraction and budget application in OpenAIModel. Added tests for MCP tool registration and execution.
This commit is contained in:
@@ -44,6 +44,7 @@ class MCPServer:
|
||||
self.version = version
|
||||
self.resources: Dict[str, MCPResource] = {}
|
||||
self.tools: Dict[str, Callable] = {}
|
||||
self.tool_specs: Dict[str, MCPTool] = {}
|
||||
self.prompts: Dict[str, MCPPrompt] = {}
|
||||
|
||||
async def initialize(self):
|
||||
@@ -61,6 +62,7 @@ class MCPServer:
|
||||
def register_tool(self, name: str, description: str, input_schema: Dict, handler: Callable):
|
||||
"""注册工具"""
|
||||
tool = MCPTool(name=name, description=description, input_schema=input_schema)
|
||||
self.tool_specs[name] = tool
|
||||
self.tools[name] = handler
|
||||
logger.info(f"✅ MCP工具注册: {self.name}.{name}")
|
||||
|
||||
@@ -78,10 +80,7 @@ class MCPServer:
|
||||
|
||||
async def list_tools(self) -> List[MCPTool]:
|
||||
"""列出工具"""
|
||||
return [
|
||||
MCPTool(name=name, description="", input_schema={})
|
||||
for name in self.tools.keys()
|
||||
]
|
||||
return list(self.tool_specs.values())
|
||||
|
||||
async def call_tool(self, name: str, arguments: Dict[str, Any]) -> Any:
|
||||
"""调用工具"""
|
||||
@@ -89,7 +88,16 @@ class MCPServer:
|
||||
raise ValueError(f"工具不存在: {name}")
|
||||
|
||||
handler = self.tools[name]
|
||||
return await handler(**arguments)
|
||||
logger.info(
|
||||
f"MCP工具调用开始: server={self.name}, tool={name}, args={json.dumps(arguments, ensure_ascii=False)}"
|
||||
)
|
||||
try:
|
||||
result = await handler(**arguments)
|
||||
except Exception as exc:
|
||||
logger.warning(f"MCP工具调用失败: server={self.name}, tool={name}, error={exc}")
|
||||
raise
|
||||
logger.info(f"MCP工具调用成功: server={self.name}, tool={name}")
|
||||
return result
|
||||
|
||||
async def list_prompts(self) -> List[MCPPrompt]:
|
||||
"""列出提示词"""
|
||||
@@ -216,4 +224,5 @@ class MCPManager:
|
||||
raise ValueError(f"工具名格式错误: {full_tool_name}")
|
||||
|
||||
server_name, tool_name = parts
|
||||
logger.info(f"MCP执行请求: {full_tool_name}")
|
||||
return await self.client.call_tool(server_name, tool_name, arguments)
|
||||
|
||||
Reference in New Issue
Block a user