zhuchaowe/maxkb
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: enhance MCP tool execution with unique ID generation and improved file handling

Browse Source
This commit is contained in:
CaptainB 2025-08-13 11:47:59 +08:00
parent f0a3391897
commit bb665b5196
1 changed files with 17 additions and 9 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

26
apps/application/flow/step_node/ai_chat_step_node/impl/base_chat_node.py
View File

@ -16,6 +16,7 @@ import traceback
from functools import reduce from functools import reduce
from typing import List, Dict from typing import List, Dict
import uuid_utils.compat as uuid
from django.db.models import QuerySet from django.db.models import QuerySet
from langchain.schema import HumanMessage, SystemMessage from langchain.schema import HumanMessage, SystemMessage
from langchain_core.messages import BaseMessage, AIMessage, AIMessageChunk, ToolMessage from langchain_core.messages import BaseMessage, AIMessage, AIMessageChunk, ToolMessage
@ -27,7 +28,6 @@ from application.flow.step_node.ai_chat_step_node.i_chat_node import IChatNode
from application.flow.tools import Reasoning from application.flow.tools import Reasoning
from common.utils.logger import maxkb_logger from common.utils.logger import maxkb_logger
from common.utils.tool_code import ToolExecutor from common.utils.tool_code import ToolExecutor
from maxkb.conf import PROJECT_DIR
from models_provider.models import Model from models_provider.models import Model
from models_provider.tools import get_model_credential, get_model_instance_by_model_workspace_id from models_provider.tools import get_model_credential, get_model_instance_by_model_workspace_id
from tools.models import Tool from tools.models import Tool
@ -280,18 +280,26 @@ class BaseChatNode(IChatNode):
if tool_enable: if tool_enable:
if tool_ids and len(tool_ids) > 0: # 如果有工具ID，则将其转换为MCP if tool_ids and len(tool_ids) > 0: # 如果有工具ID，则将其转换为MCP
self.context['tool_ids'] = tool_ids self.context['tool_ids'] = tool_ids
self.context['execute_ids'] = []
for tool_id in tool_ids: for tool_id in tool_ids:
tool = QuerySet(Tool).filter(id=tool_id).first() tool = QuerySet(Tool).filter(id=tool_id).first()
executor = ToolExecutor() executor = ToolExecutor(sandbox=True)
code = executor.generate_mcp_server_code(tool.code) code = executor.generate_mcp_server_code(tool.code)
code_path = f'{executor.sandbox_path}/execute/{tool_id}.py' _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: with open(code_path, 'w') as f:
f.write(code) f.write(code)
os.system(f"chown sandbox:root {code_path}") os.system(f"chown {executor.user}:root {code_path}")
tool_config = { tool_config = {
'command': sys.executable, 'command': 'su',
'args': [code_path], 'args': [
'-s', sys.executable,
'-c', f"exec(open('{code_path}', 'r').read())",
executor.user,
],
'cwd': executor.sandbox_path,
'transport': 'stdio', 'transport': 'stdio',
} }
mcp_servers_config[str(tool.id)] = tool_config mcp_servers_config[str(tool.id)] = tool_config
@ -338,10 +346,10 @@ class BaseChatNode(IChatNode):
def get_details(self, index: int, **kwargs): def get_details(self, index: int, **kwargs):
# 删除临时生成的MCP代码文件 # 删除临时生成的MCP代码文件
if self.context.get('tool_ids'): if self.context.get('execute_ids'):
executor = ToolExecutor() executor = ToolExecutor(sandbox=True)
# 清理工具代码文件，延时删除，避免文件被占用 # 清理工具代码文件，延时删除，避免文件被占用
for tool_id in self.context.get('tool_ids'): for tool_id in self.context.get('execute_ids'):
code_path = f'{executor.sandbox_path}/execute/{tool_id}.py' code_path = f'{executor.sandbox_path}/execute/{tool_id}.py'
if os.path.exists(code_path): if os.path.exists(code_path):
os.remove(code_path) os.remove(code_path)