91 lines
3.1 KiB
Python
91 lines
3.1 KiB
Python
from __future__ import annotations
|
|
|
|
from dataclasses import dataclass, field
|
|
from typing import Any
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class GraphNodeSpec:
|
|
node_id: str
|
|
node_type: str
|
|
metadata: dict[str, Any] = field(default_factory=dict)
|
|
|
|
def to_dict(self) -> dict[str, Any]:
|
|
return {"id": self.node_id, "type": self.node_type, "metadata": self.metadata}
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class GraphEdgeSpec:
|
|
source: str
|
|
target: str
|
|
condition: str = ""
|
|
metadata: dict[str, Any] = field(default_factory=dict)
|
|
|
|
def to_dict(self) -> dict[str, Any]:
|
|
return {"source": self.source, "target": self.target, "condition": self.condition, "metadata": self.metadata}
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class ExecutionGraphSpec:
|
|
name: str
|
|
version: int
|
|
graph_type: str
|
|
entry: str
|
|
nodes: dict[str, GraphNodeSpec]
|
|
edges: list[GraphEdgeSpec]
|
|
terminal_nodes: list[str]
|
|
metadata: dict[str, Any] = field(default_factory=dict)
|
|
|
|
def to_dict(self) -> dict[str, Any]:
|
|
return {
|
|
"name": self.name,
|
|
"version": self.version,
|
|
"graph_type": self.graph_type,
|
|
"entry": self.entry,
|
|
"nodes": {node_id: node.to_dict() for node_id, node in self.nodes.items()},
|
|
"edges": [edge.to_dict() for edge in self.edges],
|
|
"terminal_nodes": self.terminal_nodes,
|
|
"metadata": self.metadata,
|
|
}
|
|
|
|
@classmethod
|
|
def from_dict(cls, payload: dict[str, Any]) -> "ExecutionGraphSpec":
|
|
nodes = {
|
|
node_id: GraphNodeSpec(node_id=str(raw["id"]), node_type=str(raw["type"]), metadata=dict(raw.get("metadata", {})))
|
|
for node_id, raw in dict(payload["nodes"]).items()
|
|
}
|
|
edges = [
|
|
GraphEdgeSpec(
|
|
source=str(raw["source"]),
|
|
target=str(raw["target"]),
|
|
condition=str(raw.get("condition", "")),
|
|
metadata=dict(raw.get("metadata", {})),
|
|
)
|
|
for raw in list(payload["edges"])
|
|
]
|
|
return cls(
|
|
name=str(payload["name"]),
|
|
version=int(payload["version"]),
|
|
graph_type=str(payload["graph_type"]),
|
|
entry=str(payload["entry"]),
|
|
nodes=nodes,
|
|
edges=edges,
|
|
terminal_nodes=[str(item) for item in payload.get("terminal_nodes", [])],
|
|
metadata=dict(payload.get("metadata", {})),
|
|
)
|
|
|
|
def validate(self) -> None:
|
|
if self.entry not in self.nodes:
|
|
raise ValueError(f"Graph entry node does not exist: {self.entry}")
|
|
for terminal in self.terminal_nodes:
|
|
if terminal not in self.nodes:
|
|
raise ValueError(f"Terminal node does not exist: {terminal}")
|
|
for edge in self.edges:
|
|
if edge.source not in self.nodes:
|
|
raise ValueError(f"Edge source does not exist: {edge.source}")
|
|
if edge.target not in self.nodes:
|
|
raise ValueError(f"Edge target does not exist: {edge.target}")
|
|
|
|
|
|
def outgoing_edges(spec: ExecutionGraphSpec, node_id: str) -> list[GraphEdgeSpec]:
|
|
return [edge for edge in spec.edges if edge.source == node_id]
|