120 lines
3.1 KiB
Python
120 lines
3.1 KiB
Python
"""LangGraph compilation for deterministic workflow definitions."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from typing import Annotated, Callable, Literal, TypedDict
|
|
|
|
from langgraph.graph import END, START, StateGraph
|
|
from langgraph.graph.state import CompiledStateGraph
|
|
|
|
from gyxx_flow.core.context import RunContext
|
|
|
|
from .model import StepDefinition, WorkflowDefinition
|
|
|
|
GraphStepStatus = Literal["success", "failed", "skipped"]
|
|
|
|
|
|
class GraphAttemptResult(TypedDict):
|
|
"""JSON-serializable attempt result stored in graph state."""
|
|
|
|
attempt: int
|
|
status: GraphStepStatus
|
|
exit_code: int | None
|
|
error: str | None
|
|
reason: str | None
|
|
|
|
|
|
class GraphStepResult(TypedDict):
|
|
"""JSON-serializable step result stored in graph state."""
|
|
|
|
step_id: str
|
|
status: GraphStepStatus
|
|
attempts: list[GraphAttemptResult]
|
|
final_exit_code: int | None
|
|
critical: bool
|
|
reason: str | None
|
|
|
|
|
|
class WorkflowGraphState(TypedDict):
|
|
"""Serializable state passed between one workflow's LangGraph nodes."""
|
|
|
|
workflow_id: str
|
|
run_id: str
|
|
business_date: str
|
|
started_at: str
|
|
shadow: bool
|
|
dry_run: bool
|
|
step_results: Annotated[
|
|
dict[str, GraphStepResult],
|
|
_merge_step_results,
|
|
]
|
|
|
|
|
|
GraphStepRunner = Callable[
|
|
[StepDefinition, WorkflowGraphState],
|
|
GraphStepResult,
|
|
]
|
|
|
|
|
|
def initial_workflow_graph_state(
|
|
context: RunContext,
|
|
*,
|
|
dry_run: bool,
|
|
) -> WorkflowGraphState:
|
|
"""Create a JSON-serializable initial state from the public run context."""
|
|
|
|
return {
|
|
"workflow_id": context.workflow_id,
|
|
"run_id": context.run_id,
|
|
"business_date": context.business_date.isoformat(),
|
|
"started_at": context.started_at.isoformat(),
|
|
"shadow": context.shadow,
|
|
"dry_run": dry_run,
|
|
"step_results": {},
|
|
}
|
|
|
|
|
|
def compile_workflow_graph(
|
|
workflow: WorkflowDefinition,
|
|
*,
|
|
run_step: GraphStepRunner,
|
|
) -> CompiledStateGraph:
|
|
"""Compile dependencies into a real LangGraph DAG."""
|
|
|
|
builder = StateGraph(WorkflowGraphState)
|
|
ordered_steps = workflow.ordered_steps()
|
|
dependent_ids = {
|
|
dependency
|
|
for step in ordered_steps
|
|
for dependency in step.depends_on
|
|
}
|
|
for step in ordered_steps:
|
|
builder.add_node(step.step_id, _step_node(step, run_step))
|
|
if step.depends_on:
|
|
builder.add_edge(list(step.depends_on), step.step_id)
|
|
else:
|
|
builder.add_edge(START, step.step_id)
|
|
if step.step_id not in dependent_ids:
|
|
builder.add_edge(step.step_id, END)
|
|
return builder.compile()
|
|
|
|
|
|
def _step_node(
|
|
step: StepDefinition,
|
|
run_step: GraphStepRunner,
|
|
) -> Callable[[WorkflowGraphState], dict[str, object]]:
|
|
def execute(state: WorkflowGraphState) -> dict[str, object]:
|
|
result = run_step(step, state)
|
|
return {"step_results": {step.step_id: result}}
|
|
|
|
return execute
|
|
|
|
|
|
def _merge_step_results(
|
|
current: dict[str, GraphStepResult],
|
|
update: dict[str, GraphStepResult],
|
|
) -> dict[str, GraphStepResult]:
|
|
"""Merge branch-local results emitted in the same LangGraph superstep."""
|
|
|
|
return {**current, **update}
|