feat: 支持 Python 代码节点 main 返回输出

- 自动传入节点输入并映射 main 返回字典,兼容历史 _result

- 限制脚本输出缓冲并补充运行时与编辑器测试

- 同步 Python 帮助、占位提示和独立可执行补全片段
This commit is contained in:
2026-08-03 11:16:03 +08:00
parent 51dbfd41b6
commit 1bf755f6c4
5 changed files with 267 additions and 17 deletions

View File

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

View File

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