feat: enhance MCP tool execution with unique ID generation and improved file handling
This commit is contained in:
CaptainB
parent
bb665b5196
commit
5e99770999
@ -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'
|
||||
|
||||
@ -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:
|
||||
|
||||
Loading…
Reference in New Issue
Block a user