Build lightweight AI agent admin
This commit is contained in:
Codex
@@ -0,0 +1,182 @@
|
||||
"""
|
||||
代码执行节点
|
||||
"""
|
||||
import concurrent.futures
|
||||
import logging
|
||||
import time
|
||||
from typing import Any, Dict
|
||||
|
||||
from ..base import BaseNode, NodeContext, NodeResult
|
||||
from ..registry import NodeRegistry
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@NodeRegistry.register
|
||||
class CodeNode(BaseNode):
|
||||
"""
|
||||
代码执行节点
|
||||
|
||||
执行 Python 代码片段
|
||||
"""
|
||||
|
||||
node_type = 'code'
|
||||
node_name = '代码'
|
||||
node_category = 'logic'
|
||||
node_icon = 'code'
|
||||
node_description = '执行 Python 代码片段'
|
||||
|
||||
inputs = [
|
||||
{
|
||||
'name': 'inputs',
|
||||
'type': 'object',
|
||||
'description': '输入变量',
|
||||
},
|
||||
]
|
||||
|
||||
outputs = [
|
||||
{
|
||||
'name': 'result',
|
||||
'type': 'any',
|
||||
'description': '执行结果',
|
||||
},
|
||||
]
|
||||
|
||||
# 安全的内置函数白名单
|
||||
SAFE_BUILTINS = {
|
||||
'abs', 'all', 'any', 'bool', 'dict', 'enumerate', 'filter',
|
||||
'float', 'int', 'len', 'list', 'map', 'max', 'min', 'range',
|
||||
'round', 'set', 'sorted', 'str', 'sum', 'tuple', 'zip',
|
||||
'True', 'False', 'None',
|
||||
}
|
||||
|
||||
def execute(self, context: NodeContext) -> NodeResult:
|
||||
"""执行代码（线程池隔离，防止死循环卡死主流程）"""
|
||||
timeout = self.config.get('timeout', 30)
|
||||
|
||||
with concurrent.futures.ThreadPoolExecutor(max_workers=4) as executor:
|
||||
future = executor.submit(self._execute_sync, context)
|
||||
try:
|
||||
return future.result(timeout=timeout)
|
||||
except concurrent.futures.TimeoutError:
|
||||
future.cancel()
|
||||
return NodeResult(
|
||||
success=False,
|
||||
error=f'代码执行超时（{timeout}秒），请检查是否存在死循环',
|
||||
elapsed_time=timeout * 1000,
|
||||
)
|
||||
|
||||
def _execute_sync(self, context: NodeContext) -> NodeResult:
|
||||
"""同步执行代码（在子线程中运行）"""
|
||||
start_time = time.time()
|
||||
|
||||
try:
|
||||
code = self.config.get('code', '')
|
||||
input_variables = self.config.get('inputs', []) or self.config.get('input_variables', [])
|
||||
output_variable = self.config.get('output_variable', 'result')
|
||||
|
||||
if not code:
|
||||
return NodeResult(
|
||||
success=False,
|
||||
error='代码不能为空',
|
||||
)
|
||||
|
||||
safe_globals = {
|
||||
'__builtins__': {k: getattr(__builtins__, k) if hasattr(__builtins__, k) else __builtins__[k]
|
||||
for k in self.SAFE_BUILTINS if hasattr(__builtins__, k) or k in __builtins__},
|
||||
}
|
||||
|
||||
import json
|
||||
import re
|
||||
import math
|
||||
safe_globals['json'] = json
|
||||
safe_globals['re'] = re
|
||||
safe_globals['math'] = math
|
||||
|
||||
local_vars = {}
|
||||
for var_config in input_variables:
|
||||
if isinstance(var_config, dict):
|
||||
var_name = var_config.get('variable', '')
|
||||
default_value = var_config.get('default_value', None)
|
||||
if var_name:
|
||||
value = context.get_variable(var_name)
|
||||
if value is not None:
|
||||
local_vars[var_name] = value
|
||||
elif default_value:
|
||||
local_vars[var_name] = context.resolve_template(str(default_value))
|
||||
else:
|
||||
local_vars[var_name] = None
|
||||
elif isinstance(var_config, str):
|
||||
local_vars[var_config] = context.get_variable(var_config)
|
||||
|
||||
local_vars['user_input'] = context.user_input
|
||||
local_vars['variables'] = context.variables.copy()
|
||||
|
||||
exec(code, safe_globals, local_vars)
|
||||
|
||||
if 'main' in local_vars and callable(local_vars['main']):
|
||||
inputs_dict = {
|
||||
'user_input': context.user_input,
|
||||
**context.variables,
|
||||
**local_vars,
|
||||
}
|
||||
main_result = local_vars['main'](inputs_dict)
|
||||
if isinstance(main_result, dict):
|
||||
result = main_result.get('result', main_result)
|
||||
else:
|
||||
result = main_result
|
||||
else:
|
||||
result = local_vars.get('result', None)
|
||||
|
||||
elapsed_time = int((time.time() - start_time) * 1000)
|
||||
|
||||
return NodeResult(
|
||||
success=True,
|
||||
output=result,
|
||||
output_variables={output_variable: result},
|
||||
elapsed_time=elapsed_time,
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
logger.exception(f'代码节点执行失败: {e}')
|
||||
elapsed_time = int((time.time() - start_time) * 1000)
|
||||
return NodeResult(
|
||||
success=False,
|
||||
error=str(e),
|
||||
elapsed_time=elapsed_time,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def get_config_schema(cls) -> Dict[str, Any]:
|
||||
"""获取配置 Schema"""
|
||||
return {
|
||||
'type': 'object',
|
||||
'properties': {
|
||||
'code': {
|
||||
'type': 'string',
|
||||
'title': '代码',
|
||||
'description': 'Python 代码，结果存储在 result 变量中',
|
||||
'format': 'code',
|
||||
'default': '# 在这里编写代码\n# 可用变量: user_input, variables\n# 将结果赋值给 result\n\nresult = user_input.upper()',
|
||||
},
|
||||
'input_variables': {
|
||||
'type': 'array',
|
||||
'title': '输入变量',
|
||||
'items': {'type': 'string'},
|
||||
'description': '需要传入代码的变量名列表',
|
||||
},
|
||||
'output_variable': {
|
||||
'type': 'string',
|
||||
'title': '输出变量名',
|
||||
'default': 'code_result',
|
||||
},
|
||||
'timeout': {
|
||||
'type': 'integer',
|
||||
'title': '超时时间（秒）',
|
||||
'default': 30,
|
||||
'minimum': 1,
|
||||
'maximum': 300,
|
||||
},
|
||||
},
|
||||
'required': ['code'],
|
||||
}
|
||||
Reference in New Issue
Block a user