From 5e997709990e82ed1a9ee47c06f39fc1bba94c74 Mon Sep 17 00:00:00 2001 From: CaptainB Date: Wed, 13 Aug 2025 12:30:30 +0800 Subject: [PATCH] feat: enhance MCP tool execution with unique ID generation and improved file handling --- .../ai_chat_step_node/impl/base_chat_node.py | 29 ++++----- apps/common/utils/tool_code.py | 60 +++++++++++++++++-- 2 files changed, 66 insertions(+), 23 deletions(-) diff --git a/apps/application/flow/step_node/ai_chat_step_node/impl/base_chat_node.py b/apps/application/flow/step_node/ai_chat_step_node/impl/base_chat_node.py index a082475b..2db06fe6 100644 --- a/apps/application/flow/step_node/ai_chat_step_node/impl/base_chat_node.py +++ b/apps/application/flow/step_node/ai_chat_step_node/impl/base_chat_node.py @@ -27,7 +27,9 @@ from application.flow.i_step_node import NodeResult, INode from application.flow.step_node.ai_chat_step_node.i_chat_node import IChatNode from application.flow.tools import Reasoning from common.utils.logger import maxkb_logger +from common.utils.rsa_util import rsa_long_decrypt from common.utils.tool_code import ToolExecutor +from maxkb.const import CONFIG from models_provider.models import Model from models_provider.tools import get_model_credential, get_model_instance_by_model_workspace_id from tools.models import Tool @@ -283,25 +285,14 @@ class BaseChatNode(IChatNode): self.context['execute_ids'] = [] for tool_id in tool_ids: tool = QuerySet(Tool).filter(id=tool_id).first() - executor = ToolExecutor(sandbox=True) - code = executor.generate_mcp_server_code(tool.code) - _id = uuid.uuid7() - self.context['execute_ids'].append(_id) - code_path = f'{executor.sandbox_path}/execute/{_id}.py' - with open(code_path, 'w') as f: - f.write(code) - os.system(f"chown {executor.user}:root {code_path}") + executor = ToolExecutor(CONFIG.get('SANDBOX')) + if tool.init_params is not None: + params = json.loads(rsa_long_decrypt(tool.init_params)) + else: + params = {} + _id, tool_config = executor.get_tool_mcp_config(tool.code, params) - tool_config = { - 'command': 'su', - 'args': [ - '-s', sys.executable, - '-c', f"exec(open('{code_path}', 'r').read())", - executor.user, - ], - 'cwd': executor.sandbox_path, - 'transport': 'stdio', - } + self.context['execute_ids'].append(_id) mcp_servers_config[str(tool.id)] = tool_config if len(mcp_servers_config) > 0: @@ -347,7 +338,7 @@ class BaseChatNode(IChatNode): def get_details(self, index: int, **kwargs): # 删除临时生成的MCP代码文件 if self.context.get('execute_ids'): - executor = ToolExecutor(sandbox=True) + executor = ToolExecutor(CONFIG.get('SANDBOX')) # 清理工具代码文件，延时删除，避免文件被占用 for tool_id in self.context.get('execute_ids'): code_path = f'{executor.sandbox_path}/execute/{tool_id}.py' diff --git a/apps/common/utils/tool_code.py b/apps/common/utils/tool_code.py index 1a8e50f0..31b892f6 100644 --- a/apps/common/utils/tool_code.py +++ b/apps/common/utils/tool_code.py @@ -83,7 +83,7 @@ except Exception as e: return result.get('data') raise Exception(result.get('msg')) - def _generate_mcp_server_code(self, _code): + def _generate_mcp_server_code(self, _code, params): self.validate_banned_keywords(_code) # 解析代码，提取导入语句和函数定义 @@ -100,7 +100,31 @@ except Exception as e: if isinstance(node, ast.Import) or isinstance(node, ast.ImportFrom): imports.append(ast.unparse(node)) elif isinstance(node, ast.FunctionDef): - # 为函数添加 @mcp.tool() 装饰器 + # 修改函数参数以包含 params 中的默认值 + func_name = node.name + if func_name in params: + func_params = params[func_name] + # 为函数参数设置默认值 + for i, arg in enumerate(node.args.args): + arg_name = arg.arg + if arg_name in func_params: + # 创建默认值节点 + default_value = func_params[arg_name] + if isinstance(default_value, str): + default_node = ast.Constant(value=default_value) + elif isinstance(default_value, (int, float, bool)): + default_node = ast.Constant(value=default_value) + else: + default_node = ast.Constant(value=str(default_value)) + + # 添加到defaults列表 + if not hasattr(node.args, 'defaults') or node.args.defaults is None: + node.args.defaults = [] + # 确保defaults列表长度正确 + while len(node.args.defaults) < len(node.args.args): + node.args.defaults.insert(0, None) + node.args.defaults[i] = default_node + func_code = ast.unparse(node) functions.append(f"@mcp.tool()\n{func_code}\n") else: @@ -116,9 +140,9 @@ except Exception as e: return "\n".join(code_parts) - def generate_mcp_server_code(self, code_str): + def generate_mcp_server_code(self, code_str, params): python_paths = CONFIG.get_sandbox_python_package_paths().split(',') - code = self._generate_mcp_server_code(code_str) + code = self._generate_mcp_server_code(code_str, params) return f""" import os import sys @@ -132,6 +156,34 @@ for key in list(env.keys()): exec({dedent(code)!a}) """ + def get_tool_mcp_config(self, code, params): + code = self.generate_mcp_server_code(code, params) + + _id = uuid.uuid7() + code_path = f'{self.sandbox_path}/execute/{_id}.py' + with open(code_path, 'w') as f: + f.write(code) + if self.sandbox: + os.system(f"chown {self.user}:root {code_path}") + + tool_config = { + 'command': 'su', + 'args': [ + '-s', sys.executable, + '-c', f"exec(open('{code_path}', 'r').read())", + self.user, + ], + 'cwd': self.sandbox_path, + 'transport': 'stdio', + } + else: + tool_config = { + 'command': sys.executable, + 'args': [code_path], + 'transport': 'stdio', + } + return _id, tool_config + def _exec_sandbox(self, _code, _id): exec_python_file = f'{self.sandbox_path}/execute/{_id}.py' with open(exec_python_file, 'w') as file: