114 lines
4.3 KiB
Python
114 lines
4.3 KiB
Python
from __future__ import annotations
|
|
|
|
from graph.models import ExecutionGraphDefinition, ExecutionGraphVersion, ExecutionGraphVersionStatus, GraphEdgeTraversal, GraphRun, GraphRunStatus
|
|
from graph.native_runtime import NativeGraphRuntime
|
|
from graph.registry import NodeHandlerRegistry, NodeResult
|
|
from graph.spec import ExecutionGraphSpec, GraphEdgeSpec, GraphNodeSpec
|
|
|
|
|
|
class FixedNode:
|
|
idempotent = True
|
|
replay_safe = True
|
|
destructive = False
|
|
|
|
def __init__(self, node_type: str, edge_result: str = "success") -> None:
|
|
self.node_type = node_type
|
|
self.edge_result = edge_result
|
|
self.calls = 0
|
|
|
|
def run(self, context: object) -> NodeResult:
|
|
self.calls += 1
|
|
return NodeResult("COMPLETE", self.edge_result, {"calls": self.calls}, {"model_requests": 1})
|
|
|
|
|
|
def persist_spec(spec: ExecutionGraphSpec) -> GraphRun:
|
|
definition = ExecutionGraphDefinition.objects.create(name=spec.name, graph_type=spec.graph_type)
|
|
version = ExecutionGraphVersion.objects.create(graph=definition, version=spec.version, status=ExecutionGraphVersionStatus.CHAMPION, graph_spec=spec.to_dict())
|
|
return GraphRun.objects.create(execution_graph_version=version, current_node=spec.entry)
|
|
|
|
|
|
def test_native_graph_runtime_executes_conditional_edges_and_persists_history() -> None:
|
|
spec = ExecutionGraphSpec(
|
|
name="fixture_graph",
|
|
version=1,
|
|
graph_type="FIXTURE",
|
|
entry="start",
|
|
nodes={
|
|
"start": GraphNodeSpec("start", "start"),
|
|
"success": GraphNodeSpec("success", "success"),
|
|
"fail": GraphNodeSpec("fail", "fail"),
|
|
},
|
|
edges=[GraphEdgeSpec("start", "success", "ok"), GraphEdgeSpec("start", "fail", "bad")],
|
|
terminal_nodes=["success", "fail"],
|
|
)
|
|
graph_run = persist_spec(spec)
|
|
registry = NodeHandlerRegistry()
|
|
start = FixedNode("start", "ok")
|
|
registry.register(start)
|
|
|
|
result = NativeGraphRuntime(registry).run_until_terminal_or_paused(graph_run)
|
|
|
|
assert result.status == GraphRunStatus.COMPLETE
|
|
assert result.current_node == "success"
|
|
assert start.calls == 1
|
|
assert list(GraphEdgeTraversal.objects.filter(graph_run=graph_run).values_list("source_node", "target_node", "condition")) == [("start", "success", "ok")]
|
|
assert graph_run.node_runs.get(node_id="start").telemetry["duration_ms"] >= 0
|
|
|
|
|
|
def test_native_graph_runtime_resume_does_not_repeat_completed_current_node() -> None:
|
|
spec = ExecutionGraphSpec(
|
|
name="resume_graph",
|
|
version=1,
|
|
graph_type="FIXTURE",
|
|
entry="first",
|
|
nodes={
|
|
"first": GraphNodeSpec("first", "first"),
|
|
"second": GraphNodeSpec("second", "second"),
|
|
"done": GraphNodeSpec("done", "done"),
|
|
},
|
|
edges=[GraphEdgeSpec("first", "second", "success"), GraphEdgeSpec("second", "done", "success")],
|
|
terminal_nodes=["done"],
|
|
)
|
|
graph_run = persist_spec(spec)
|
|
registry = NodeHandlerRegistry()
|
|
first = FixedNode("first")
|
|
second = FixedNode("second")
|
|
registry.register(first)
|
|
registry.register(second)
|
|
runtime = NativeGraphRuntime(registry)
|
|
|
|
runtime.run_until_terminal_or_paused(graph_run, interrupt_after="first")
|
|
graph_run.refresh_from_db()
|
|
assert graph_run.current_node == "first"
|
|
|
|
runtime.run_until_terminal_or_paused(graph_run)
|
|
|
|
graph_run.refresh_from_db()
|
|
assert graph_run.status == GraphRunStatus.COMPLETE
|
|
assert first.calls == 1
|
|
assert second.calls == 1
|
|
|
|
|
|
def test_native_graph_runtime_records_paused_runs() -> None:
|
|
class PauseNode(FixedNode):
|
|
def run(self, context: object) -> NodeResult:
|
|
self.calls += 1
|
|
return NodeResult("PAUSED", "awaiting", pause_reason="AWAITING_APPROVAL")
|
|
|
|
spec = ExecutionGraphSpec(
|
|
name="pause_graph",
|
|
version=1,
|
|
graph_type="FIXTURE",
|
|
entry="approval",
|
|
nodes={"approval": GraphNodeSpec("approval", "approval"), "done": GraphNodeSpec("done", "done")},
|
|
edges=[GraphEdgeSpec("approval", "done", "approved")],
|
|
terminal_nodes=["done"],
|
|
)
|
|
graph_run = persist_spec(spec)
|
|
registry = NodeHandlerRegistry()
|
|
registry.register(PauseNode("approval"))
|
|
|
|
result = NativeGraphRuntime(registry).run_until_terminal_or_paused(graph_run)
|
|
|
|
assert result.status == GraphRunStatus.PAUSED
|
|
assert result.failure_reason == "AWAITING_APPROVAL"
|