feat: 支持 Python 代码节点 main 返回输出
- 自动传入节点输入并映射 main 返回字典,兼容历史 _result - 限制脚本输出缓冲并补充运行时与编辑器测试 - 同步 Python 帮助、占位提示和独立可执行补全片段
This commit is contained in:
@@ -14,10 +14,19 @@ import java.io.IOException;
|
||||
import java.io.InputStream;
|
||||
import java.io.OutputStream;
|
||||
import java.nio.charset.StandardCharsets;
|
||||
import java.util.ArrayList;
|
||||
import java.util.Collections;
|
||||
import java.util.HashMap;
|
||||
import java.util.Map;
|
||||
import java.util.concurrent.TimeUnit;
|
||||
|
||||
/**
|
||||
* 基于独立 Python 子进程的代码节点执行器。
|
||||
*
|
||||
* <p>优先保留历史 {@code _result} 输出约定;当脚本未写入
|
||||
* {@code _result} 且声明了 {@code main} 函数时,自动传入节点参数并使用
|
||||
* {@code main} 返回的字典作为节点输出。</p>
|
||||
*/
|
||||
public class PythonRuntimeEngine implements CodeRuntimeEngine {
|
||||
|
||||
private static final String PYTHON_HELPER_SCRIPT = String.join("\n",
|
||||
@@ -27,6 +36,40 @@ public class PythonRuntimeEngine implements CodeRuntimeEngine {
|
||||
"import sys",
|
||||
"import traceback",
|
||||
"",
|
||||
"class _BoundedTextWriter(io.TextIOBase):",
|
||||
" def __init__(self, max_bytes):",
|
||||
" super().__init__()",
|
||||
" self._max_bytes = max(0, max_bytes)",
|
||||
" self._written_bytes = 0",
|
||||
" self._buffer = io.StringIO()",
|
||||
" self.truncated = False",
|
||||
"",
|
||||
" def write(self, text):",
|
||||
" if not isinstance(text, str):",
|
||||
" raise TypeError('write() argument must be str')",
|
||||
" if not text:",
|
||||
" return 0",
|
||||
" remaining = self._max_bytes - self._written_bytes",
|
||||
" if remaining <= 0:",
|
||||
" self.truncated = True",
|
||||
" return len(text)",
|
||||
" candidate = text[:remaining]",
|
||||
" encoded = candidate.encode('utf-8')",
|
||||
" if len(encoded) > remaining:",
|
||||
" candidate = encoded[:remaining].decode('utf-8', errors='ignore')",
|
||||
" encoded = candidate.encode('utf-8')",
|
||||
" self._buffer.write(candidate)",
|
||||
" self._written_bytes += len(encoded)",
|
||||
" if len(candidate) < len(text):",
|
||||
" self.truncated = True",
|
||||
" return len(text)",
|
||||
"",
|
||||
" def getvalue(self):",
|
||||
" return self._buffer.getvalue()",
|
||||
"",
|
||||
" def writable(self):",
|
||||
" return True",
|
||||
"",
|
||||
"def _read_payload():",
|
||||
" raw = sys.stdin.read()",
|
||||
" if not raw:",
|
||||
@@ -39,23 +82,40 @@ public class PythonRuntimeEngine implements CodeRuntimeEngine {
|
||||
" context = payload.get('context')",
|
||||
" if not isinstance(context, dict):",
|
||||
" context = {}",
|
||||
" parameter_names = payload.get('parameterNames')",
|
||||
" if not isinstance(parameter_names, list):",
|
||||
" parameter_names = []",
|
||||
" max_output_bytes = payload.get('maxOutputBytes')",
|
||||
" if not isinstance(max_output_bytes, int):",
|
||||
" max_output_bytes = 0",
|
||||
"",
|
||||
" local_ctx = dict(context)",
|
||||
" local_ctx['_result'] = {}",
|
||||
" local_ctx['__builtins__'] = __builtins__",
|
||||
"",
|
||||
" captured_stdout = io.StringIO()",
|
||||
" captured_stderr = io.StringIO()",
|
||||
" captured_stdout = _BoundedTextWriter(max_output_bytes)",
|
||||
" captured_stderr = _BoundedTextWriter(max_output_bytes)",
|
||||
" main_invoked = False",
|
||||
"",
|
||||
" try:",
|
||||
" with contextlib.redirect_stdout(captured_stdout), contextlib.redirect_stderr(captured_stderr):",
|
||||
" exec(code, {'__builtins__': __builtins__}, local_ctx)",
|
||||
" exec(code, local_ctx, local_ctx)",
|
||||
" result = local_ctx.get('_result')",
|
||||
" if isinstance(result, dict) and not result:",
|
||||
" main_func = local_ctx.get('main')",
|
||||
" if callable(main_func):",
|
||||
" main_invoked = True",
|
||||
" main_input = {name: context.get(name) for name in parameter_names if isinstance(name, str)}",
|
||||
" result = main_func(main_input)",
|
||||
" except Exception:",
|
||||
" sys.stderr.write(traceback.format_exc())",
|
||||
" sys.exit(2)",
|
||||
"",
|
||||
" result = local_ctx.get('_result')",
|
||||
" if not isinstance(result, dict):",
|
||||
" sys.stderr.write('Python 脚本执行结束后,_result 必须是 dict\\n')",
|
||||
" if main_invoked:",
|
||||
" sys.stderr.write(\"Python main 函数必须返回 dict,例如:return {'result': value}\\n\")",
|
||||
" else:",
|
||||
" sys.stderr.write('Python 脚本执行结束后,_result 必须是 dict\\n')",
|
||||
" sys.exit(3)",
|
||||
"",
|
||||
" if captured_stdout.getvalue():",
|
||||
@@ -130,8 +190,37 @@ public class PythonRuntimeEngine implements CodeRuntimeEngine {
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* 执行 Python 脚本,兼容直接调用旧接口的场景。
|
||||
*
|
||||
* @param code Python 脚本
|
||||
* @param node 当前代码节点
|
||||
* @param chain 当前工作流
|
||||
* @return 脚本输出结果
|
||||
*/
|
||||
@Override
|
||||
public Map<String, Object> execute(String code, CodeNode node, Chain chain) {
|
||||
Map<String, Object> parameterValues =
|
||||
chain.getExecutionState().resolveParameters(node);
|
||||
return execute(code, node, chain, parameterValues);
|
||||
}
|
||||
|
||||
/**
|
||||
* 使用调用方已解析的参数执行 Python 脚本。
|
||||
*
|
||||
* <p>该方法作为普通重载保留,可同时兼容仅声明三参数接口的旧版
|
||||
* easy-agents;新版接口会按同签名方法直接分派,避免重复解析参数。</p>
|
||||
*
|
||||
* @param code Python 脚本
|
||||
* @param node 当前代码节点
|
||||
* @param chain 当前工作流
|
||||
* @param parameterValues 已解析的节点参数
|
||||
* @return 脚本输出结果
|
||||
*/
|
||||
public Map<String, Object> execute(String code,
|
||||
CodeNode node,
|
||||
Chain chain,
|
||||
Map<String, Object> parameterValues) {
|
||||
if (StringUtil.noText(command)) {
|
||||
throw new RuntimeException("python 执行命令为空,请检查配置 node.code-engine.python.command");
|
||||
}
|
||||
@@ -158,10 +247,16 @@ public class PythonRuntimeEngine implements CodeRuntimeEngine {
|
||||
stdoutThread.start();
|
||||
stderrThread.start();
|
||||
|
||||
Map<String, Object> context = buildContext(chain, node);
|
||||
Map<String, Object> context = buildContext(chain, parameterValues);
|
||||
Map<String, Object> payload = new HashMap<>();
|
||||
payload.put("code", code);
|
||||
payload.put("context", context);
|
||||
payload.put(
|
||||
"parameterNames",
|
||||
parameterValues == null
|
||||
? Collections.emptyList()
|
||||
: new ArrayList<>(parameterValues.keySet()));
|
||||
payload.put("maxOutputBytes", maxOutputBytes);
|
||||
|
||||
String payloadText = JSON.toJSONString(payload);
|
||||
try (OutputStream stdin = process.getOutputStream()) {
|
||||
@@ -217,7 +312,15 @@ public class PythonRuntimeEngine implements CodeRuntimeEngine {
|
||||
}
|
||||
}
|
||||
|
||||
private Map<String, Object> buildContext(Chain chain, CodeNode node) {
|
||||
/**
|
||||
* 构建传入 Python 子进程的上下文。
|
||||
*
|
||||
* @param chain 当前工作流
|
||||
* @param parameterValues 已解析的节点参数
|
||||
* @return Python 脚本上下文
|
||||
*/
|
||||
private Map<String, Object> buildContext(
|
||||
Chain chain, Map<String, Object> parameterValues) {
|
||||
Map<String, Object> context = new HashMap<>();
|
||||
ChainState chainState =
|
||||
chain.getExecutionState();
|
||||
@@ -230,8 +333,6 @@ public class PythonRuntimeEngine implements CodeRuntimeEngine {
|
||||
}
|
||||
});
|
||||
|
||||
Map<String, Object> parameterValues =
|
||||
chainState.resolveParameters(node);
|
||||
if (parameterValues != null && !parameterValues.isEmpty()) {
|
||||
context.putAll(parameterValues);
|
||||
}
|
||||
|
||||
@@ -31,6 +31,106 @@ public class PythonRuntimeEngineTest {
|
||||
Assert.assertEquals("dict", result.get("env_type"));
|
||||
}
|
||||
|
||||
@Test
|
||||
public void testExecuteMainSuccess() {
|
||||
Assume.assumeTrue(PythonRuntimeEngine.probe(PYTHON_COMMAND, 1500L).isAvailable());
|
||||
|
||||
Chain chain = createChain();
|
||||
CodeNode node = (CodeNode) chain.getDefinition().getNodeById("code-test");
|
||||
chain.getState().getMemory().put("ambient", "memory-value");
|
||||
PythonRuntimeEngine engine = new PythonRuntimeEngine(
|
||||
PYTHON_COMMAND,
|
||||
3000L,
|
||||
65536,
|
||||
System.getProperty("java.io.tmpdir"));
|
||||
|
||||
Map<String, Object> result = engine.execute(
|
||||
String.join("\n",
|
||||
"import math",
|
||||
"def main(inputs):",
|
||||
" return {",
|
||||
" 'answer': inputs['input'],",
|
||||
" 'score': 95,",
|
||||
" 'sqrt': math.sqrt(16),",
|
||||
" 'has_ambient': 'ambient' in inputs",
|
||||
" }"),
|
||||
node,
|
||||
chain,
|
||||
Map.of("input", "hello"));
|
||||
|
||||
Assert.assertEquals("hello", result.get("answer"));
|
||||
Assert.assertEquals(95, ((Number) result.get("score")).intValue());
|
||||
Assert.assertEquals(4D, ((Number) result.get("sqrt")).doubleValue(), 0D);
|
||||
Assert.assertEquals(Boolean.FALSE, result.get("has_ambient"));
|
||||
}
|
||||
|
||||
@Test
|
||||
public void testLegacyResultTakesPrecedenceOverMain() {
|
||||
Assume.assumeTrue(PythonRuntimeEngine.probe(PYTHON_COMMAND, 1500L).isAvailable());
|
||||
|
||||
Chain chain = createChain();
|
||||
CodeNode node = (CodeNode) chain.getDefinition().getNodeById("code-test");
|
||||
PythonRuntimeEngine engine = new PythonRuntimeEngine(
|
||||
PYTHON_COMMAND,
|
||||
3000L,
|
||||
65536,
|
||||
System.getProperty("java.io.tmpdir"));
|
||||
|
||||
Map<String, Object> result = engine.execute(
|
||||
String.join("\n",
|
||||
"_result['source'] = 'legacy'",
|
||||
"def main(inputs):",
|
||||
" return {'source': 'main'}"),
|
||||
node,
|
||||
chain,
|
||||
Map.of("input", "hello"));
|
||||
|
||||
Assert.assertEquals("legacy", result.get("source"));
|
||||
}
|
||||
|
||||
@Test
|
||||
public void testMainMustReturnDict() {
|
||||
Assume.assumeTrue(PythonRuntimeEngine.probe(PYTHON_COMMAND, 1500L).isAvailable());
|
||||
|
||||
Chain chain = createChain();
|
||||
CodeNode node = (CodeNode) chain.getDefinition().getNodeById("code-test");
|
||||
PythonRuntimeEngine engine = new PythonRuntimeEngine(
|
||||
PYTHON_COMMAND,
|
||||
3000L,
|
||||
65536,
|
||||
System.getProperty("java.io.tmpdir"));
|
||||
|
||||
assertExecuteFail(
|
||||
engine,
|
||||
node,
|
||||
chain,
|
||||
"def main(inputs):\n return 'invalid'",
|
||||
"Python main 函数必须返回 dict");
|
||||
}
|
||||
|
||||
@Test
|
||||
public void testCapturedOutputIsBounded() {
|
||||
Assume.assumeTrue(PythonRuntimeEngine.probe(PYTHON_COMMAND, 1500L).isAvailable());
|
||||
|
||||
Chain chain = createChain();
|
||||
CodeNode node = (CodeNode) chain.getDefinition().getNodeById("code-test");
|
||||
PythonRuntimeEngine engine = new PythonRuntimeEngine(
|
||||
PYTHON_COMMAND,
|
||||
3000L,
|
||||
64,
|
||||
System.getProperty("java.io.tmpdir"));
|
||||
|
||||
Map<String, Object> result = engine.execute(
|
||||
String.join("\n",
|
||||
"import sys",
|
||||
"print('a' * 2048)",
|
||||
"_result['truncated'] = sys.stdout.truncated"),
|
||||
node,
|
||||
chain);
|
||||
|
||||
Assert.assertEquals(Boolean.TRUE, result.get("truncated"));
|
||||
}
|
||||
|
||||
@Test
|
||||
public void testSyntaxError() {
|
||||
Assume.assumeTrue(PythonRuntimeEngine.probe(PYTHON_COMMAND, 1500L).isAvailable());
|
||||
|
||||
Reference in New Issue
Block a user