feat: refactor MCP server code generation and add execution method

This commit is contained in:
CaptainB 2025-08-12 17:13:30 +08:00
parent f2459ed124
commit 1d7a5af626
2 changed files with 28 additions and 2 deletions

View File

@ -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)

View File

@ -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: