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.InputStream;
import java.io.OutputStream; import java.io.OutputStream;
import java.nio.charset.StandardCharsets; import java.nio.charset.StandardCharsets;
import java.util.ArrayList;
import java.util.Collections;
import java.util.HashMap; import java.util.HashMap;
import java.util.Map; import java.util.Map;
import java.util.concurrent.TimeUnit; import java.util.concurrent.TimeUnit;
/**
* 基于独立 Python 子进程的代码节点执行器。
*
* <p>优先保留历史 {@code _result} 输出约定;当脚本未写入
* {@code _result} 且声明了 {@code main} 函数时,自动传入节点参数并使用
* {@code main} 返回的字典作为节点输出。</p>
*/
public class PythonRuntimeEngine implements CodeRuntimeEngine { public class PythonRuntimeEngine implements CodeRuntimeEngine {
private static final String PYTHON_HELPER_SCRIPT = String.join("\n", private static final String PYTHON_HELPER_SCRIPT = String.join("\n",
@@ -27,6 +36,40 @@ public class PythonRuntimeEngine implements CodeRuntimeEngine {
"import sys", "import sys",
"import traceback", "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():", "def _read_payload():",
" raw = sys.stdin.read()", " raw = sys.stdin.read()",
" if not raw:", " if not raw:",
@@ -39,23 +82,40 @@ public class PythonRuntimeEngine implements CodeRuntimeEngine {
" context = payload.get('context')", " context = payload.get('context')",
" if not isinstance(context, dict):", " if not isinstance(context, dict):",
" context = {}", " 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 = dict(context)",
" local_ctx['_result'] = {}", " local_ctx['_result'] = {}",
" local_ctx['__builtins__'] = __builtins__",
"", "",
" captured_stdout = io.StringIO()", " captured_stdout = _BoundedTextWriter(max_output_bytes)",
" captured_stderr = io.StringIO()", " captured_stderr = _BoundedTextWriter(max_output_bytes)",
" main_invoked = False",
"", "",
" try:", " try:",
" with contextlib.redirect_stdout(captured_stdout), contextlib.redirect_stderr(captured_stderr):", " 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:", " except Exception:",
" sys.stderr.write(traceback.format_exc())", " sys.stderr.write(traceback.format_exc())",
" sys.exit(2)", " sys.exit(2)",
"", "",
" result = local_ctx.get('_result')",
" if not isinstance(result, dict):", " 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)", " sys.exit(3)",
"", "",
" if captured_stdout.getvalue():", " if captured_stdout.getvalue():",
@@ -130,8 +190,37 @@ public class PythonRuntimeEngine implements CodeRuntimeEngine {
} }
} }
/**
* 执行 Python 脚本,兼容直接调用旧接口的场景。
*
* @param code Python 脚本
* @param node 当前代码节点
* @param chain 当前工作流
* @return 脚本输出结果
*/
@Override @Override
public Map<String, Object> execute(String code, CodeNode node, Chain chain) { 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)) { if (StringUtil.noText(command)) {
throw new RuntimeException("python 执行命令为空,请检查配置 node.code-engine.python.command"); throw new RuntimeException("python 执行命令为空,请检查配置 node.code-engine.python.command");
} }
@@ -158,10 +247,16 @@ public class PythonRuntimeEngine implements CodeRuntimeEngine {
stdoutThread.start(); stdoutThread.start();
stderrThread.start(); stderrThread.start();
Map<String, Object> context = buildContext(chain, node); Map<String, Object> context = buildContext(chain, parameterValues);
Map<String, Object> payload = new HashMap<>(); Map<String, Object> payload = new HashMap<>();
payload.put("code", code); payload.put("code", code);
payload.put("context", context); payload.put("context", context);
payload.put(
"parameterNames",
parameterValues == null
? Collections.emptyList()
: new ArrayList<>(parameterValues.keySet()));
payload.put("maxOutputBytes", maxOutputBytes);
String payloadText = JSON.toJSONString(payload); String payloadText = JSON.toJSONString(payload);
try (OutputStream stdin = process.getOutputStream()) { 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<>(); Map<String, Object> context = new HashMap<>();
ChainState chainState = ChainState chainState =
chain.getExecutionState(); chain.getExecutionState();
@@ -230,8 +333,6 @@ public class PythonRuntimeEngine implements CodeRuntimeEngine {
} }
}); });
Map<String, Object> parameterValues =
chainState.resolveParameters(node);
if (parameterValues != null && !parameterValues.isEmpty()) { if (parameterValues != null && !parameterValues.isEmpty()) {
context.putAll(parameterValues); context.putAll(parameterValues);
} }

View File

@@ -31,6 +31,106 @@ public class PythonRuntimeEngineTest {
Assert.assertEquals("dict", result.get("env_type")); 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 @Test
public void testSyntaxError() { public void testSyntaxError() {
Assume.assumeTrue(PythonRuntimeEngine.probe(PYTHON_COMMAND, 1500L).isAvailable()); Assume.assumeTrue(PythonRuntimeEngine.probe(PYTHON_COMMAND, 1500L).isAvailable());

View File

@@ -22,8 +22,8 @@
let currentNode = useNodesData(currentNodeId); let currentNode = useNodesData(currentNodeId);
const { addParameter } = useAddParameter(); const { addParameter } = useAddParameter();
const { updateNodeData } = useSvelteFlow(); const { updateNodeData } = useSvelteFlow();
const codeNodeHelp = `JavaScript 如何返回结果 const codeNodeHelp = `代码如何返回结果
- 定义 main 函数并返回对象,系统会自动执行。 - JavaScript 定义 main 函数并返回对象Python 定义 main 函数并返回 dict,系统会自动执行。
- main 接收由输入参数组成的对象。 - main 接收由输入参数组成的对象。
输出参数如何配置 输出参数如何配置
@@ -31,11 +31,15 @@
- 字段名与 main 返回对象的 key 保持一致。 - 字段名与 main 返回对象的 key 保持一致。
- 下游节点可引用代码节点ID.输出参数名。 - 下游节点可引用代码节点ID.输出参数名。
示例 JavaScript 示例
- 代码里写:
function main({ input }) { function main({ input }) {
return { answer: input, score: 95 }; return { answer: input, score: 95 };
} }
Python 示例
def main(inputs):
return {'answer': inputs.get('input'), 'score': 95}
- 输出参数配置answer(String)、score(Number) - 输出参数配置answer(String)、score(Number)
- 结束节点输出参数可引用代码节点ID.answer、代码节点ID.score - 结束节点输出参数可引用代码节点ID.answer、代码节点ID.score
@@ -53,6 +57,13 @@
const firstAvailable = engines.find((item) => item.selectable !== false); const firstAvailable = engines.find((item) => item.selectable !== false);
return firstAvailable?.value || 'js'; return firstAvailable?.value || 'js';
}); });
const codePlaceholder = $derived.by(() => {
const engine = String(data.engine || defaultEngine).trim().toLowerCase();
if (engine === 'python' || engine === 'py') {
return "请输入代码例如def main(inputs):\n return {'output': inputs.get('input')}";
}
return '请输入代码例如function main({ input }) { return { output: input }; }';
});
onMount(async () => { onMount(async () => {
const codeEngines = await options.provider?.codeEngine?.(); const codeEngines = await options.provider?.codeEngine?.();
@@ -118,7 +129,7 @@
<CodeScriptEditor <CodeScriptEditor
mode="textarea" mode="textarea"
rows={10} rows={10}
placeholder={'请输入代码例如function main({ input }) { return { output: input }; }'} placeholder={codePlaceholder}
style="width: 100%" style="width: 100%"
engine={(data.engine as string) || (defaultEngine as string)} engine={(data.engine as string) || (defaultEngine as string)}
parameters={editorParameters} parameters={editorParameters}

View File

@@ -71,6 +71,36 @@ describe('codeCompletion utils', () => {
); );
}); });
it('should prefer main return snippets for python output', () => {
const completions = createBusinessCompletions('python', []);
const mainCompletion = completions.find((item) => item.label === 'main');
const resultCompletion = completions.find(
(item) => item.label === 'result-object',
);
expect(mainCompletion).toMatchObject({
type: 'snippet',
detail: '代码节点入口',
boost: 1000,
});
expect(String(mainCompletion?.apply)).toContain('def main(inputs):');
expect(String(mainCompletion?.apply)).toContain('return {');
expect(String(resultCompletion?.apply)).toContain('def main(inputs):');
expect(String(resultCompletion?.apply)).toContain(
"return {'message': 'ok', 'data': inputs.get('data')}",
);
});
it('should provide standalone python snippets', () => {
const completions = createBusinessCompletions('python', []);
const snippets = completions.filter((item) => item.type === 'snippet');
expect(snippets).toHaveLength(4);
for (const snippet of snippets) {
expect(String(snippet.apply)).toMatch(/^def main\(inputs\):\n /);
}
});
it('should detect blocked nodes for comment/string/property contexts', () => { it('should detect blocked nodes for comment/string/property contexts', () => {
expect(shouldSkipBusinessCompletion('Comment')).toBe(true); expect(shouldSkipBusinessCompletion('Comment')).toBe(true);
expect(shouldSkipBusinessCompletion('LineComment')).toBe(true); expect(shouldSkipBusinessCompletion('LineComment')).toBe(true);

View File

@@ -60,21 +60,29 @@ const PYTHON_SNIPPETS: Array<{
insert: string; insert: string;
detail: string; detail: string;
}> = [ }> = [
{
label: 'main',
detail: '代码节点入口',
insert:
"def main(inputs):\n return {\n 'output': inputs.get('input'),\n }",
},
{ {
label: 'if-else', label: 'if-else',
detail: '条件分支', detail: '条件分支',
insert: insert:
"if condition:\n _result['value'] = value\nelse:\n _result['value'] = None", "def main(inputs):\n condition = inputs.get('condition')\n value = inputs.get('value')\n if condition:\n return {'value': value}\n return {'value': None}",
}, },
{ {
label: 'for-loop', label: 'for-loop',
detail: '遍历数组', detail: '遍历数组',
insert: "for item in items:\n # TODO\n pass\n_result['done'] = True", insert:
"def main(inputs):\n result = []\n for item in inputs.get('items', []):\n result.append(item)\n return {'result': result}",
}, },
{ {
label: 'result-object', label: 'result-object',
detail: '返回对象', detail: '返回对象',
insert: "_result['message'] = 'ok'\n_result['data'] = data", insert:
"def main(inputs):\n return {'message': 'ok', 'data': inputs.get('data')}",
}, },
]; ];