fix: Use AI-generated tool code in MCP server instead of template

2026-01-30 18:37:06 +00:00
parent 56668402c3
commit 1a99ba3ac1
2 changed files with 140 additions and 88 deletions
@@ -382,21 +382,15 @@ async def {func_name}({params_str}) -> str:
            auth = tool.get("auth", {})
            request_params = tool.get("request_params", {})
            timeout = tool.get("timeout", 30)
+            generated_code = tool.get("generated_code", "")
            
-            # 构建参数
-            params = []
-            params_doc = []
+            # 构建参数信息（用于 TOOL_LIST）
            properties = {}
            required_params = []
-            
-            # 支持两种格式：
-            # 1. {"properties": {"symbol": {...}}}
-            # 2. {"symbol": {...}} (直接参数格式)
            param_props = request_params
            if request_params and request_params.get("properties"):
                param_props = request_params["properties"]
            elif request_params and not any(k in request_params for k in ["type", "required", "description"]):
-                # 直接参数格式
                param_props = request_params
            else:
                param_props = {}
@@ -405,93 +399,48 @@ async def {func_name}({params_str}) -> str:
                for p_name, p_info in param_props.items():
                    if not isinstance(p_info, dict):
                        continue
-                    p_type = self._json_type_to_python(p_info.get("type", "string"))
                    p_desc = p_info.get("description", "")
                    is_required = p_info.get("required", False)
-                    default = p_info.get("default")
-                    
                    if is_required:
-                        params.append(f"{p_name}: {p_type}")
                        required_params.append(p_name)
-                    else:
-                        default_val = f'"{default}"' if isinstance(default, str) else (default if default is not None else "None")
-                        params.append(f"{p_name}: Optional[{p_type}] = {default_val}")
-                    
-                    params_doc.append(f"        {p_name}: {p_desc}")
                    properties[p_name] = {"type": p_info.get("type", "string"), "description": p_desc}
            
-            params_str = ", ".join(params) if params else ""
-            params_doc_str = "\n".join(params_doc) if params_doc else "        无参数"
-            
-            # 生成认证代码
-            auth_headers = self._get_auth_headers_code(auth)
-            
-            # 构建参数字典代码
-            params_dict_code = ""
-            if param_props:
-                params_dict_code = "params = {"
-                for p_name in param_props.keys():
-                    if isinstance(param_props[p_name], dict):
-                        params_dict_code += f'"{p_name}": {p_name}, '
-                params_dict_code = params_dict_code.rstrip(", ") + "}"
-            else:
-                params_dict_code = "params = {}"
-            
-            # API Key in query
-            if auth and auth.get("type") == "api_key" and auth.get("in") == "query":
-                key_name = auth.get("name", "apikey")
-                params_dict_code += f'\n        params["{key_name}"] = os.getenv("TOOL_API_KEY", "")'
-            
-            # 生成函数代码
-            func_code = f'''
-@server.tool()
-async def {func_name}({params_str}) -> str:
-    """
-    {desc}
-    
-    Args:
-{params_doc_str}
-    
-    Returns:
-        API 响应结果 (JSON 格式)
-    """
-    import httpx
-    
-    url = "{url}"
-    {auth_headers}
-    {params_dict_code}
-    
-    try:
-        async with httpx.AsyncClient(timeout={timeout}) as client:
-            response = await client.request(
-                method="{method}",
-                url=url,
-                headers=headers,
-                params={{k: v for k, v in params.items() if v is not None}}
-            )
-            
-            if response.status_code == 200:
-                return json.dumps({{
-                    "success": True,
-                    "data": response.json() if response.headers.get("content-type", "").startswith("application/json") else response.text
-                }}, ensure_ascii=False, indent=2)
-            else:
-                return json.dumps({{
-                    "success": False,
-                    "status_code": response.status_code,
-                    "error": response.text[:500]
-                }}, ensure_ascii=False)
+            # 如果有 AI 生成的代码，使用它；否则使用模板
+            if generated_code:
+                # 从 AI 生成的代码中提取函数并添加 @server.tool() 装饰器
+                import re
+                # 移除开头的 docstring 和 import 语句
+                code_lines = generated_code.split('\n')
+                func_start = -1
+                for i, line in enumerate(code_lines):
+                    if line.strip().startswith('async def '):
+                        func_start = i
+                        break
                
-    except Exception as e:
-        # 使用 AI Agent 作为后备
-        result = await get_agent().run(f"请帮我处理这个请求: {params}")
-        return json.dumps({{
-            "success": True,
-            "source": "ai_agent",
-            "result": result.output
-        }}, ensure_ascii=False, indent=2)
+                if func_start >= 0:
+                    # 提取函数代码
+                    func_code_lines = code_lines[func_start:]
+                    func_body = '\n'.join(func_code_lines)
+                    
+                    # 添加 @server.tool() 装饰器
+                    func_code = f'''
+@server.tool()
+{func_body}
 '''
-            tool_functions.append(func_code)
+                    tool_functions.append(func_code)
+                else:
+                    # 无法解析，使用原始代码
+                    logger.warning(f"无法解析 AI 生成的代码: {name}")
+                    func_code = self._generate_fallback_tool_code(
+                        func_name, desc, url, method, auth, param_props, timeout
+                    )
+                    tool_functions.append(func_code)
+            else:
+                # 使用模板生成代码
+                func_code = self._generate_fallback_tool_code(
+                    func_name, desc, url, method, auth, param_props, timeout
+                )
+                tool_functions.append(func_code)
            tool_map_entries.append(f"    '{func_name}': {func_name},")
            
            tool_list_entries.append(f'''    {{
@@ -603,6 +552,106 @@ if __name__ == '__main__':
        
        return "headers = {}"
    
+    def _generate_fallback_tool_code(
+        self,
+        func_name: str,
+        desc: str,
+        url: str,
+        method: str,
+        auth: Dict,
+        param_props: Dict,
+        timeout: int
+    ) -> str:
+        """生成后备工具代码（当没有 AI 生成代码时使用）"""
+        # 构建参数
+        params = []
+        params_doc = []
+        
+        if param_props:
+            for p_name, p_info in param_props.items():
+                if not isinstance(p_info, dict):
+                    continue
+                p_type = self._json_type_to_python(p_info.get("type", "string"))
+                p_desc = p_info.get("description", "")
+                is_required = p_info.get("required", False)
+                default = p_info.get("default")
+                
+                if is_required:
+                    params.append(f"{p_name}: {p_type}")
+                else:
+                    default_val = f'"{default}"' if isinstance(default, str) else (default if default is not None else "None")
+                    params.append(f"{p_name}: Optional[{p_type}] = {default_val}")
+                
+                params_doc.append(f"        {p_name}: {p_desc}")
+        
+        params_str = ", ".join(params) if params else ""
+        params_doc_str = "\n".join(params_doc) if params_doc else "        无参数"
+        
+        # 生成认证代码
+        auth_headers = self._get_auth_headers_code(auth)
+        
+        # 构建参数字典代码
+        params_dict_code = ""
+        if param_props:
+            params_dict_code = "params = {"
+            for p_name in param_props.keys():
+                if isinstance(param_props[p_name], dict):
+                    params_dict_code += f'"{p_name}": {p_name}, '
+            params_dict_code = params_dict_code.rstrip(", ") + "}"
+        else:
+            params_dict_code = "params = {}"
+        
+        # API Key in query
+        if auth and auth.get("type") == "api_key" and auth.get("in") == "query":
+            key_name = auth.get("name", "apikey")
+            params_dict_code += f'\n        params["{key_name}"] = os.getenv("TOOL_API_KEY", "")'
+        
+        return f'''
+@server.tool()
+async def {func_name}({params_str}) -> str:
+    """
+    {desc}
+    
+    Args:
+{params_doc_str}
+    
+    Returns:
+        API 响应结果 (JSON 格式)
+    """
+    import httpx
+    
+    api_url = "{url}"
+    {auth_headers}
+    {params_dict_code}
+    
+    try:
+        async with httpx.AsyncClient(timeout={timeout}) as client:
+            response = await client.request(
+                method="{method}",
+                url=api_url,
+                headers=headers,
+                params={{k: v for k, v in params.items() if v is not None}}
+            )
+            
+            if response.status_code == 200:
+                return json.dumps({{
+                    "success": True,
+                    "data": response.json() if response.headers.get("content-type", "").startswith("application/json") else response.text
+                }}, ensure_ascii=False, indent=2)
+            else:
+                return json.dumps({{
+                    "success": False,
+                    "status_code": response.status_code,
+                    "error": response.text[:500]
+                }}, ensure_ascii=False)
+                
+    except Exception as e:
+        return json.dumps({{
+            "success": False,
+            "error": str(e)
+        }}, ensure_ascii=False, indent=2)
+'''
+    
    def _generate_system_prompt(
        self,
        agent_name: str,
@@ -844,6 +844,8 @@ async def create_agent_with_tools(request: CreateAgentWithToolsRequest):
        tools_config = []
        for tool in tools:
            config = tool.get("config", {})
+            # 获取 AI 生成的工具代码
+            tool_code = tool_storage.get_tool_code(tool.get("tool_ref_id", ""))
            tools_config.append({
                "name": tool["name"],
                "description": tool.get("description", ""),
@@ -852,7 +854,8 @@ async def create_agent_with_tools(request: CreateAgentWithToolsRequest):
                "auth": config.get("auth"),
                "request_params": config.get("request_params"),
                "request_body": config.get("request_body"),
-                "timeout": config.get("timeout", 30)
+                "timeout": config.get("timeout", 30),
+                "generated_code": tool_code  # 传递 AI 生成的代码
            })
        
        # 生成唯一后缀，确保不同用户创建同名 Agent 不会冲突