15308555518/ MaxKB
1
0
Fork 0
mirror of https://github.com/1Panel-dev/MaxKB.git synced 2025-12-26 01:33:05 +00:00
Code Issues Actions 2 Packages Projects Releases Wiki Activity

refactor: remove init parameter from class ToolExecutor.

This commit is contained in:
liqiang-fit2cloud 2025-12-04 15:40:50 +08:00
parent b034737b35
commit f8ada9a110
7 changed files with 33 additions and 37 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
apps/application/chat_pipeline/step/chat_step/impl/base_chat_step.py
View File

@ -253,7 +253,7 @@ class BaseChatStep(IChatStep):
tool = QuerySet(Tool).filter(id=tool_id).first()
if tool is None or tool.is_active is False:
continue
executor = ToolExecutor(CONFIG.get('SANDBOX'))
executor = ToolExecutor()
if tool.init_params is not None:
params = json.loads(rsa_long_decrypt(tool.init_params))
else:

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

@ -237,7 +237,7 @@ class BaseChatNode(IChatNode):
tool = QuerySet(Tool).filter(id=tool_id).first()
if not tool.is_active:
continue
executor = ToolExecutor(CONFIG.get('SANDBOX'))
executor = ToolExecutor()
if tool.init_params is not None:
params = json.loads(rsa_long_decrypt(tool.init_params))
else:

2
apps/application/flow/step_node/tool_lib_node/impl/base_tool_lib_node.py
View File

@ -22,7 +22,7 @@ from common.utils.tool_code import ToolExecutor
from maxkb.const import CONFIG
from tools.models import Tool
function_executor = ToolExecutor(CONFIG.get('SANDBOX'))
function_executor = ToolExecutor()
def write_context(step_variable: Dict, global_variable: Dict, node, workflow):

2
apps/application/flow/step_node/tool_node/impl/base_tool_node.py
View File

@ -17,7 +17,7 @@ from application.flow.step_node.tool_node.i_tool_node import IToolNode
from common.utils.tool_code import ToolExecutor
from maxkb.const import CONFIG
function_executor = ToolExecutor(CONFIG.get('SANDBOX'))
function_executor = ToolExecutor()
def write_context(step_variable: Dict, global_variable: Dict, node, workflow):

58
apps/common/utils/tool_code.py
View File

@ -19,24 +19,21 @@ from maxkb.const import BASE_DIR, CONFIG
from maxkb.const import PROJECT_DIR
from textwrap import dedent
_enable_sandbox = bool(CONFIG.get('SANDBOX', 0))
_run_user = 'sandbox' if _enable_sandbox else getpass.getuser()
_sandbox_path = CONFIG.get("SANDBOX_HOME", '/opt/maxkb-app/sandbox') if _enable_sandbox else os.path.join(PROJECT_DIR, 'data', 'sandbox')
_process_limit_timeout_seconds = int(CONFIG.get("SANDBOX_PYTHON_PROCESS_LIMIT_TIMEOUT_SECONDS", '3600'))
_process_limit_cpu_cores = min(max(int(CONFIG.get("SANDBOX_PYTHON_PROCESS_LIMIT_CPU_CORES", '1')), 1), len(os.sched_getaffinity(0))) if sys.platform.startswith("linux") else os.cpu_count() # 只支持linux，window和mac不支持
_process_limit_mem_mb = int(CONFIG.get("SANDBOX_PYTHON_PROCESS_LIMIT_MEM_MB", '256'))
class ToolExecutor:
enable_sandbox = bool(CONFIG.get('SANDBOX', 0))
sandbox_path = CONFIG.get("SANDBOX_HOME", '/opt/maxkb-app/sandbox') if enable_sandbox else os.path.join(PROJECT_DIR, 'data', 'sandbox')
process_timeout_seconds = int(CONFIG.get("SANDBOX_PYTHON_PROCESS_TIMEOUT_SECONDS", '3600'))
process_limit_mem_mb = int(CONFIG.get("SANDBOX_PYTHON_PROCESS_LIMIT_MEM_MB", '256'))
process_limit_cpu_cores = min(max(int(CONFIG.get("SANDBOX_PYTHON_PROCESS_LIMIT_CPU_CORES", '1')), 1), len(os.sched_getaffinity(0))) if sys.platform.startswith("linux") else os.cpu_count() # 只支持linux，window和mac不支持
def __init__(self, sandbox=False):
self.sandbox = sandbox
if sandbox:
self.user = 'sandbox'
else:
self.user = getpass.getuser()
def __init__(self):
pass
@staticmethod
def init_sandbox_dir():
if not ToolExecutor.enable_sandbox:
if not _enable_sandbox:
# 不启用sandbox就不初始化目录
return
try:
@ -57,7 +54,7 @@ class ToolExecutor:
if CONFIG.get("SANDBOX_TMP_DIR_ENABLED", '0') == "1":
os.system("chmod g+rwx /tmp")
# 初始化sandbox配置文件
sandbox_lib_path = os.path.dirname(f'{ToolExecutor.sandbox_path}/lib/sandbox.so')
sandbox_lib_path = os.path.dirname(f'{_sandbox_path}/lib/sandbox.so')
sandbox_conf_file_path = f'{sandbox_lib_path}/.sandbox.conf'
if os.path.exists(sandbox_conf_file_path):
os.remove(sandbox_conf_file_path)
@ -70,8 +67,12 @@ class ToolExecutor:
with open(sandbox_conf_file_path, "w") as f:
f.write(f"SANDBOX_PYTHON_BANNED_HOSTS={banned_hosts}\n")
f.write(f"SANDBOX_PYTHON_ALLOW_SUBPROCESS={allow_subprocess}\n")
os.system(f"chmod -R 550 {ToolExecutor.sandbox_path}")
os.system(f"chmod -R 550 {_sandbox_path}")
try:
init_sandbox_dir()
except Exception as e:
maxkb_logger.error(f'Exception: {e}', exc_info=True)
def exec_code(self, code_str, keywords, function_name=None):
_id = str(uuid.uuid7())
@ -79,7 +80,7 @@ class ToolExecutor:
err = '{"code":500,"msg":str(e),"data":None}'
action_function = f'({function_name !a}, locals_v.get({function_name !a}))' if function_name else 'locals_v.popitem()'
python_paths = CONFIG.get_sandbox_python_package_paths().split(',')
set_run_user = f'os.setgid({pwd.getpwnam(self.user).pw_gid});os.setuid({pwd.getpwnam(self.user).pw_uid});' if self.sandbox else ''
set_run_user = f'os.setgid({pwd.getpwnam(_run_user).pw_gid});os.setuid({pwd.getpwnam(_run_user).pw_uid});' if _enable_sandbox else ''
_exec_code = f"""
try:
import os, sys, json, base64, builtins
@ -98,7 +99,7 @@ try:
exec_result=f(**keywords)
builtins.print("\\n{_id}:"+base64.b64encode(json.dumps({success}, default=str).encode()).decode(), flush=True)
except Exception as e:
if isinstance(e, MemoryError): e = Exception("Cannot allocate more memory: exceeded the limit of {ToolExecutor.process_limit_mem_mb} MB.")
if isinstance(e, MemoryError): e = Exception("Cannot allocate more memory: exceeded the limit of {_process_limit_mem_mb} MB.")
builtins.print("\\n{_id}:"+base64.b64encode(json.dumps({err}, default=str).encode()).decode(), flush=True)
"""
maxkb_logger.debug(f"Sandbox execute code: {_exec_code}")
@ -184,7 +185,7 @@ except Exception as e:
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, params)
set_run_user = f'os.setgid({pwd.getpwnam(self.user).pw_gid});os.setuid({pwd.getpwnam(self.user).pw_uid});' if self.sandbox else ''
set_run_user = f'os.setgid({pwd.getpwnam(_run_user).pw_gid});os.setuid({pwd.getpwnam(_run_user).pw_uid});' if _enable_sandbox else ''
return f"""
import os, sys, logging
logging.basicConfig(level=logging.WARNING)
@ -208,9 +209,9 @@ exec({dedent(code)!a})
'-c',
f'import base64,gzip; exec(gzip.decompress(base64.b64decode(\'{compressed_and_base64_encoded_code_str}\')).decode())',
],
'cwd': ToolExecutor.sandbox_path,
'cwd': _sandbox_path,
'env': {
'LD_PRELOAD': f'{ToolExecutor.sandbox_path}/lib/sandbox.so',
'LD_PRELOAD': f'{_sandbox_path}/lib/sandbox.so',
},
'transport': 'stdio',
}
@ -218,31 +219,26 @@ exec({dedent(code)!a})
def _exec(self, execute_file):
kwargs = {'cwd': BASE_DIR, 'env': {
'LD_PRELOAD': f'{ToolExecutor.sandbox_path}/lib/sandbox.so',
'LD_PRELOAD': f'{_sandbox_path}/lib/sandbox.so',
}}
try:
subprocess_result = subprocess.run(
[sys.executable, execute_file],
timeout=ToolExecutor.process_timeout_seconds,
timeout=_process_limit_timeout_seconds,
text=True,
capture_output=True,
**kwargs,
preexec_fn=(lambda: None if (not self.sandbox or not sys.platform.startswith("linux")) else (
resource.setrlimit(resource.RLIMIT_AS, (ToolExecutor.process_limit_mem_mb * 1024 * 1024,) * 2),
os.sched_setaffinity(0, set(random.sample(list(os.sched_getaffinity(0)), ToolExecutor.process_limit_cpu_cores)))
preexec_fn=(lambda: None if (not _enable_sandbox or not sys.platform.startswith("linux")) else (
resource.setrlimit(resource.RLIMIT_AS, (_process_limit_mem_mb * 1024 * 1024,) * 2),
os.sched_setaffinity(0, set(random.sample(list(os.sched_getaffinity(0)), _process_limit_cpu_cores)))
))
)
return subprocess_result
except subprocess.TimeoutExpired:
raise Exception(_(f"Process execution timed out after {ToolExecutor.process_timeout_seconds} seconds."))
raise Exception(_(f"Process execution timed out after {_process_limit_timeout_seconds} seconds."))
def validate_mcp_transport(self, code_str):
servers = json.loads(code_str)
for server, config in servers.items():
if config.get('transport') not in ['sse', 'streamable_http']:
raise Exception(_('Only support transport=sse or transport=streamable_http'))
try:
ToolExecutor.init_sandbox_dir()
except Exception as e:
maxkb_logger.error(f'Exception: {e}', exc_info=True)

2
apps/knowledge/serializers/knowledge_workflow.py
View File

@ -28,7 +28,7 @@ from system_manage.serializers.user_resource_permission import UserResourcePermi
from tools.models import Tool
from users.models import User
tool_executor = ToolExecutor(CONFIG.get('SANDBOX'))
tool_executor = ToolExecutor()
class KnowledgeWorkflowModelSerializer(serializers.ModelSerializer):

2
apps/tools/serializers/tool.py
View File

@ -38,7 +38,7 @@ from system_manage.serializers.user_resource_permission import UserResourcePermi
from tools.models import Tool, ToolScope, ToolFolder, ToolType
from users.serializers.user import is_workspace_manage
tool_executor = ToolExecutor(CONFIG.get('SANDBOX'))
tool_executor = ToolExecutor()
class ToolInstance: