84 lines
2.9 KiB
Python
84 lines
2.9 KiB
Python
|
"""
|
|||
|
|
TraceStudio 节点图运行时数据结构设计
|
|||
|
|
|
|||
|
|
本文件定义了节点图执行所需的核心数据结构，确保前后端一致。
|
|||
|
|
"""
|
|||
|
|
from enum import Enum
|
|||
|
|
from typing import Any, Dict, List, Optional
|
|||
|
|
from dataclasses import dataclass, field
|
|||
|
|
from app.core.node_base import NodeType
|
|||
|
|
|
|||
|
|
class DimensionMode(Enum):
|
|||
|
|
NONE = "none"
|
|||
|
|
EXPAND = "expand"
|
|||
|
|
COLLAPSE = "collapse"
|
|||
|
|
BROADCAST = "broadcast"
|
|||
|
|
|
|||
|
|
@dataclass
|
|||
|
|
class EdgeMetadata:
|
|||
|
|
source_node: str
|
|||
|
|
source_port: str
|
|||
|
|
target_node: str
|
|||
|
|
target_port: str
|
|||
|
|
dimension_mode: DimensionMode = DimensionMode.NONE
|
|||
|
|
|
|||
|
|
@dataclass
|
|||
|
|
class NodeMetadata:
|
|||
|
|
node_id: str
|
|||
|
|
# 严格约定：type 必须为 NodeType 的值（"input"/"normal"/"output"/"function"/...）
|
|||
|
|
type: str
|
|||
|
|
# 实现类名（必须提供，用于从 NodeRegistry 加载实现）
|
|||
|
|
class_name: str
|
|||
|
|
params: Dict[str, Any] = field(default_factory=dict)
|
|||
|
|
sub_workflow_nodes: Optional[List[Dict]] = None
|
|||
|
|
sub_workflow_edges: Optional[List[Dict]] = None
|
|||
|
|
|
|||
|
|
@dataclass
|
|||
|
|
class GraphRuntime:
|
|||
|
|
nodes: List[NodeMetadata]
|
|||
|
|
edges: List[EdgeMetadata]
|
|||
|
|
# 可扩展：全局上下文、运行参数等
|
|||
|
|
|
|||
|
|
def to_executor_format(self) -> Dict[str, Any]:
|
|||
|
|
"""
|
|||
|
|
将运行时 Graph 格式转换为 WorkflowExecutor 接受的 nodes/edges 列表。
|
|||
|
|
|
|||
|
|
- 将 EdgeMetadata 字段命名转换为 executor 期望的 snake_case（source, source_port, target, target_port）
|
|||
|
|
- 将 NodeMetadata 中的 sub_workflow_* 字段展开为 executor 可识别的键名（sub_workflow_nodes/sub_workflow_edges）
|
|||
|
|
- 保持 params 原样传递
|
|||
|
|
"""
|
|||
|
|
nodes_out = []
|
|||
|
|
for n in self.nodes:
|
|||
|
|
# 验证 type 为 NodeType 的合法值
|
|||
|
|
try:
|
|||
|
|
NodeType(n.type)
|
|||
|
|
except Exception:
|
|||
|
|
raise ValueError(f"NodeMetadata.type 必须为 NodeType 的值之一: {[t.value for t in NodeType]} (got {n.type})")
|
|||
|
|
|
|||
|
|
nd = {
|
|||
|
|
"id": n.node_id,
|
|||
|
|
"type": n.type,
|
|||
|
|
"class_name": n.class_name,
|
|||
|
|
"params": n.params or {}
|
|||
|
|
}
|
|||
|
|
# 支持两种约定：直接 sub_workflow_* 或嵌套 sub_workflow
|
|||
|
|
if n.sub_workflow_nodes:
|
|||
|
|
nd["sub_workflow_nodes"] = n.sub_workflow_nodes
|
|||
|
|
if n.sub_workflow_edges:
|
|||
|
|
nd["sub_workflow_edges"] = n.sub_workflow_edges
|
|||
|
|
nodes_out.append(nd)
|
|||
|
|
|
|||
|
|
edges_out = []
|
|||
|
|
for e in self.edges:
|
|||
|
|
edges_out.append({
|
|||
|
|
"source": e.source_node,
|
|||
|
|
"source_port": e.source_port,
|
|||
|
|
"target": e.target_node,
|
|||
|
|
"target_port": e.target_port,
|
|||
|
|
# executor currently ignores dimension_mode on input dicts,
|
|||
|
|
# but we include it for completeness (as string)
|
|||
|
|
"dimension_mode": e.dimension_mode.value if e.dimension_mode else DimensionMode.NONE.value
|
|||
|
|
})
|
|||
|
|
|
|||
|
|
return {"nodes": nodes_out, "edges": edges_out}
|