Codex
123 lines
4.0 KiB
Python
123 lines
4.0 KiB
Python
"""
|
||
合并节点
|
||
"""
|
||
import logging
|
||
from typing import Any, Dict, List
|
||
|
||
from ..base import BaseNode, NodeContext, NodeResult
|
||
from ..registry import NodeRegistry
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
|
||
@NodeRegistry.register
|
||
class MergeNode(BaseNode):
|
||
"""
|
||
合并节点
|
||
|
||
等待所有并行分支执行完成后，合并结果继续执行
|
||
"""
|
||
|
||
node_type = 'merge'
|
||
node_name = '合并'
|
||
node_category = 'logic'
|
||
node_icon = 'git-merge'
|
||
node_description = '等待所有并行分支完成后合并结果'
|
||
|
||
inputs = [
|
||
{
|
||
'name': 'branch_results',
|
||
'type': 'array',
|
||
'description': '各分支的执行结果',
|
||
},
|
||
]
|
||
|
||
outputs = [
|
||
{
|
||
'name': 'merged_result',
|
||
'type': 'object',
|
||
'description': '合并后的结果',
|
||
},
|
||
]
|
||
|
||
def execute(self, context: NodeContext) -> NodeResult:
|
||
"""
|
||
执行合并节点
|
||
|
||
从上下文中获取所有并行分支的结果并合并
|
||
"""
|
||
try:
|
||
# 获取合并模式
|
||
merge_mode = self.config.get('merge_mode', 'object')
|
||
output_variable = self.config.get('output_variable', 'merged_result')
|
||
|
||
# 从上下文获取并行分支结果
|
||
# 并行执行时，每个分支的结果会存储在 _parallel_results 中
|
||
parallel_results = context.get_variable('_parallel_results', {})
|
||
|
||
if merge_mode == 'object':
|
||
# 对象模式：将各分支结果合并为一个对象
|
||
merged = {}
|
||
for branch_id, result in parallel_results.items():
|
||
merged[branch_id] = result
|
||
elif merge_mode == 'array':
|
||
# 数组模式：将各分支结果合并为数组
|
||
merged = list(parallel_results.values())
|
||
elif merge_mode == 'first':
|
||
# 取第一个完成的结果
|
||
merged = list(parallel_results.values())[0] if parallel_results else None
|
||
elif merge_mode == 'concat':
|
||
# 字符串拼接模式
|
||
separator = self.config.get('separator', '\n')
|
||
merged = separator.join(str(v) for v in parallel_results.values())
|
||
else:
|
||
merged = parallel_results
|
||
|
||
return NodeResult(
|
||
success=True,
|
||
output=merged,
|
||
output_variables={
|
||
output_variable: merged,
|
||
'branch_count': len(parallel_results),
|
||
},
|
||
metadata={
|
||
'merge_mode': merge_mode,
|
||
'branch_ids': list(parallel_results.keys()),
|
||
},
|
||
)
|
||
|
||
except Exception as e:
|
||
logger.exception(f'合并节点执行失败: {e}')
|
||
return NodeResult(
|
||
success=False,
|
||
error=str(e),
|
||
)
|
||
|
||
@classmethod
|
||
def get_config_schema(cls) -> Dict[str, Any]:
|
||
"""获取配置 Schema"""
|
||
return {
|
||
'type': 'object',
|
||
'properties': {
|
||
'merge_mode': {
|
||
'type': 'string',
|
||
'title': '合并模式',
|
||
'description': '如何合并各分支的结果',
|
||
'enum': ['object', 'array', 'first', 'concat'],
|
||
'enumNames': ['对象（按分支ID）', '数组', '取第一个', '字符串拼接'],
|
||
'default': 'object',
|
||
},
|
||
'separator': {
|
||
'type': 'string',
|
||
'title': '分隔符',
|
||
'description': '字符串拼接模式的分隔符',
|
||
'default': '\n',
|
||
},
|
||
'output_variable': {
|
||
'type': 'string',
|
||
'title': '输出变量名',
|
||
'default': 'merged_result',
|
||
},
|
||
},
|
||
}
|