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 d42309a8..5505a027 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 @@ -25,6 +25,7 @@ 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.tool_code import ToolExecutor +from maxkb.conf import PROJECT_DIR 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 @@ -280,7 +281,7 @@ class BaseChatNode(IChatNode): for tool_id in tool_ids: tool = QuerySet(Tool).filter(id=tool_id).first() executor = ToolExecutor() - code = executor.generate_mcp_server_code(tool.code) + code = executor.get_exec_code(tool.code) code_path = f'{executor.sandbox_path}/execute/{tool_id}.py' with open(code_path, 'w') as f: f.write(code) diff --git a/apps/common/utils/tool_code.py b/apps/common/utils/tool_code.py index f4b8e2ad..cfe41a56 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): self.validate_banned_keywords(_code) # 解析代码,提取导入语句和函数定义 @@ -116,6 +116,31 @@ except Exception as e: return "\n".join(code_parts) + def get_exec_code(self, code_str): + python_paths = CONFIG.get_sandbox_python_package_paths().split(',') + code = self._generate_mcp_server_code(code_str) + return f""" +try: + import os + import sys + import pickle + path_to_exclude = ['/opt/py3/lib/python3.11/site-packages', '/opt/maxkb-app/apps'] + sys.path = [p for p in sys.path if p not in path_to_exclude] + sys.path += {python_paths} + env = dict(os.environ) + for key in list(env.keys()): + if key in os.environ and (key.startswith('MAXKB') or key.startswith('POSTGRES') or key.startswith('PG') or key.startswith('REDIS') or key == 'PATH'): + del os.environ[key] + locals_v={'{}'} + globals_v=globals() + exec({dedent(code)!a}, globals_v, locals_v) + f_name, f = locals_v.popitem() + for local in locals_v: + globals_v[local] = locals_v[local] +except Exception as e: + pass +""" + def _exec_sandbox(self, _code, _id): exec_python_file = f'{self.sandbox_path}/execute/{_id}.py' with open(exec_python_file, 'w') as file: