refactor Even more

2026-04-28 13:35:47 +07:00
parent 2ae027bb9a
commit 166405d88b
5 changed files with 189 additions and 124 deletions
@@ -0,0 +1,63 @@
+from __future__ import annotations
+
+from typing import Any
+
+from .model import Workflow
+from .run_state import ExecutionFrame, FrameStatus, RunState, RunStatus, TraceEntry
+from .schema_tools import validate_payload_against_schema
+from .state_ops import project_output
+from .tokens import END
+
+
+def append_trace(
+    run: RunState,
+    *,
+    frame_id: str,
+    node_id: str,
+    step_type: str,
+    resolved_input: dict[str, Any],
+    outcome: str,
+    next_node_id: str,
+    output: dict[str, Any],
+    state_changes: dict[str, Any],
+) -> None:
+    run.trace.append(
+        TraceEntry(
+            frame_id=frame_id,
+            node_id=node_id,
+            step_type=step_type,
+            resolved_input=resolved_input,
+            outcome=outcome,
+            next_node_id=next_node_id,
+            output=output,
+            state_changes=state_changes,
+        )
+    )
+
+
+def advance_frame(
+    run: RunState,
+    frame: ExecutionFrame,
+    *,
+    outcome: str,
+    next_node_id: str,
+) -> None:
+    frame.prior_outcome = outcome
+    frame.activated_incoming_edge = frame.node_id
+    frame.node_id = next_node_id
+    if next_node_id == END:
+        frame.status = FrameStatus.COMPLETED
+        frame.finished_at_node_id = END
+    else:
+        frame.finished_at_node_id = None
+    run.sync_from_current_frame()
+
+
+def finalize_run(workflow: Workflow, run: RunState) -> RunState:
+    run.output = project_output(workflow, run.state)
+    validate_payload_against_schema(
+        workflow.output_schema, run.output, "workflow output"
+    )
+    run.status = RunStatus.COMPLETED
+    run.current_node_id = END
+    return run
@@ -4,9 +4,10 @@ from typing import Any

 from .conditions import safe_resolve_path
 from .errors import WorkflowExecutionError
+from .flow_ops import advance_frame, append_trace
 from .frame_ops import frame_context_values
 from .model import ForeachNode, Workflow
-from .run_state import ExecutionFrame, FrameStatus, RunState, TraceEntry
+from .run_state import ExecutionFrame, FrameStatus, RunState


 def step_foreach(
@@ -41,22 +42,18 @@ def step_foreach(
            raise WorkflowExecutionError(
                f"no edge found for node {frame.node_id!r} and outcome {outcome!r}"
            )
-        run.trace.append(
-            TraceEntry(
-                frame_id=frame.id,
-                node_id=frame.node_id,
-                step_type=step.type,
-                resolved_input={"count": len(iterable), "index": index},
-                outcome=outcome,
-                next_node_id=next_node_id,
-                output={},
-                state_changes={},
-            )
+        append_trace(
+            run,
+            frame_id=frame.id,
+            node_id=frame.node_id,
+            step_type=step.type,
+            resolved_input={"count": len(iterable), "index": index},
+            outcome=outcome,
+            next_node_id=next_node_id,
+            output={},
+            state_changes={},
        )
-        frame.prior_outcome = outcome
-        frame.activated_incoming_edge = frame.node_id
-        frame.node_id = next_node_id
-        run.sync_from_current_frame()
+        advance_frame(run, frame, outcome=outcome, next_node_id=next_node_id)
        return run

    loop_start = edge_map.get((frame.node_id, "loop"))
@@ -82,17 +79,16 @@ def step_foreach(
        parent_frame_id=frame.id,
        metadata=child_metadata,
    )
-    run.trace.append(
-        TraceEntry(
-            frame_id=frame.id,
-            node_id=frame.node_id,
-            step_type=step.type,
-            resolved_input={"item": item, "index": index},
-            outcome="loop",
-            next_node_id=loop_start,
-            output={},
-            state_changes={},
-        )
+    append_trace(
+        run,
+        frame_id=frame.id,
+        node_id=frame.node_id,
+        step_type=step.type,
+        resolved_input={"item": item, "index": index},
+        outcome="loop",
+        next_node_id=loop_start,
+        output={},
+        state_changes={},
    )
    run.current_frame_id = child_id
    run.sync_from_current_frame()
@@ -4,10 +4,10 @@ from typing import Any

 from .conditions import safe_resolve_path
 from .errors import WorkflowExecutionError
+from .flow_ops import advance_frame, append_trace
 from .model import InterruptNode, Workflow
-from .run_state import FrameStatus, InterruptRequest, RunState, TraceEntry
+from .run_state import InterruptRequest, RunState
 from .state_ops import apply_mapped_state
-from .tokens import END


 def build_interrupt_request(
@@ -76,22 +76,16 @@ def resume_interrupt(
            f"no edge found for interrupt node {frame.node_id!r} and outcome {resume_outcome!r}"
        )

-    run.trace.append(
-        TraceEntry(
-            frame_id=frame.id,
-            node_id=frame.node_id,
-            step_type=step.type,
-            resolved_input=resume_payload,
-            outcome=resume_outcome,
-            next_node_id=next_node_id,
-            output=resume_payload,
-            state_changes=state_changes,
-        )
+    append_trace(
+        run,
+        frame_id=frame.id,
+        node_id=frame.node_id,
+        step_type=step.type,
+        resolved_input=resume_payload,
+        outcome=resume_outcome,
+        next_node_id=next_node_id,
+        output=resume_payload,
+        state_changes=state_changes,
    )
-    frame.prior_outcome = resume_outcome
-    frame.activated_incoming_edge = frame.node_id
-    frame.node_id = next_node_id
-    frame.status = FrameStatus.RUNNING if next_node_id != END else FrameStatus.COMPLETED
-    frame.finished_at_node_id = END if next_node_id == END else None
    run.interrupt = None
-    run.sync_from_current_frame()
+    advance_frame(run, frame, outcome=resume_outcome, next_node_id=next_node_id)
@@ -2,11 +2,11 @@ from __future__ import annotations

 from typing import Any

-from .conditions import eval_condition
 from .errors import WorkflowExecutionError
 from .foreach_ops import step_foreach
-from .frame_ops import collapse_completed_frames, frame_context_values
-from .interrupt_ops import build_interrupt_request, resume_interrupt
+from .flow_ops import advance_frame, append_trace, finalize_run
+from .frame_ops import collapse_completed_frames
+from .interrupt_ops import resume_interrupt
 from .model import (
    ConditionNode,
    ForeachNode,
@@ -22,10 +22,13 @@ from .run_state import (
    FrameStatus,
    RunState,
    RunStatus,
-    TraceEntry,
 )
 from .schema_tools import validate_payload_against_schema
-from .state_ops import project_output
+from .step_handlers import (
+    handle_condition_step,
+    handle_interrupt_step,
+    handle_join_step,
+)
 from .tokens import END

 __all__ = [
@@ -112,12 +115,7 @@ def resume_workflow(
        )
        collapse_completed_frames(run)
        if run.current_node_id == END:
-            run.output = project_output(workflow, run.state)
-            validate_payload_against_schema(
-                workflow.output_schema, run.output, "workflow output"
-            )
-            run.status = RunStatus.COMPLETED
-            return run
+            return finalize_run(workflow, run)

    run.status = RunStatus.RUNNING
    run.error = None
@@ -138,13 +136,7 @@ def resume_workflow(
        if run.status == RunStatus.INTERRUPTED:
            return run

-    run.output = project_output(workflow, run.state)
-    validate_payload_against_schema(
-        workflow.output_schema, run.output, "workflow output"
-    )
-    run.status = RunStatus.COMPLETED
-    run.current_node_id = END
-    return run
+    return finalize_run(workflow, run)


 def step_workflow(
@@ -187,49 +179,11 @@ def step_workflow(
        step_result = execute_node_use(workflow, run, step, node_def, registry)
        outcome = step_result["outcome"]
    elif isinstance(step, ConditionNode):
-        predicate = eval_condition(
-            step.check,
-            run.state,
-            run.workflow_input,
-            frame.prior_outcome,
-        )
-        outcome = "true" if predicate else "false"
-        step_result = {
-            "resolved_input": {},
-            "output": {"predicate": predicate},
-            "state_changes": {},
-        }
+        outcome, step_result = handle_condition_step(run, step)
    elif isinstance(step, JoinNode):
-        outcome = "done"
-        step_result = {
-            "resolved_input": {},
-            "output": {},
-            "state_changes": {},
-        }
+        outcome, step_result = handle_join_step()
    elif isinstance(step, InterruptNode):
-        interrupt_request = build_interrupt_request(
-            step,
-            frame_id=frame.id,
-            state=run.state,
-            workflow_input=run.workflow_input,
-            context=frame_context_values(frame),
-        )
-        run.interrupt = interrupt_request
-        run.status = RunStatus.INTERRUPTED
-        frame.status = FrameStatus.INTERRUPTED
-        run.trace.append(
-            TraceEntry(
-                frame_id=frame.id,
-                node_id=frame.node_id,
-                step_type=step.type,
-                resolved_input=interrupt_request.payload,
-                outcome="interrupt",
-                next_node_id=frame.node_id,
-                output=interrupt_request.payload,
-                state_changes={},
-            )
-        )
-        return run
+        return handle_interrupt_step(run, step)
    elif isinstance(step, ForeachNode):
        return step_foreach(workflow, run, step, edge_map)
    else:
@@ -241,24 +195,16 @@ def step_workflow(
            f"no edge found for node {frame.node_id!r} and outcome {outcome!r}"
        )

-    run.trace.append(
-        TraceEntry(
-            frame_id=frame.id,
-            node_id=frame.node_id,
-            step_type=step.type,
-            resolved_input=step_result["resolved_input"],
-            outcome=outcome,
-            next_node_id=next_node_id,
-            output=step_result["output"],
-            state_changes=step_result["state_changes"],
-        )
+    append_trace(
+        run,
+        frame_id=frame.id,
+        node_id=frame.node_id,
+        step_type=step.type,
+        resolved_input=step_result["resolved_input"],
+        outcome=outcome,
+        next_node_id=next_node_id,
+        output=step_result["output"],
+        state_changes=step_result["state_changes"],
    )
-
-    frame.prior_outcome = outcome
-    frame.activated_incoming_edge = frame.node_id
-    frame.node_id = next_node_id
-    if next_node_id == END:
-        frame.status = FrameStatus.COMPLETED
-        frame.finished_at_node_id = END
-    run.sync_from_current_frame()
+    advance_frame(run, frame, outcome=outcome, next_node_id=next_node_id)
    return run
@@ -0,0 +1,66 @@
+from __future__ import annotations
+
+from typing import Any
+
+from .conditions import eval_condition
+from .flow_ops import append_trace
+from .frame_ops import frame_context_values
+from .interrupt_ops import build_interrupt_request
+from .model import ConditionNode, InterruptNode, JoinNode
+from .run_state import FrameStatus, RunState, RunStatus
+
+
+def handle_condition_step(
+    run: RunState,
+    step: ConditionNode,
+) -> tuple[str, dict[str, Any]]:
+    frame = run.current_frame()
+    predicate = eval_condition(
+        step.check,
+        run.state,
+        run.workflow_input,
+        frame.prior_outcome,
+    )
+    outcome = "true" if predicate else "false"
+    return outcome, {
+        "resolved_input": {},
+        "output": {"predicate": predicate},
+        "state_changes": {},
+    }
+
+
+def handle_join_step() -> tuple[str, dict[str, Any]]:
+    return "done", {
+        "resolved_input": {},
+        "output": {},
+        "state_changes": {},
+    }
+
+
+def handle_interrupt_step(
+    run: RunState,
+    step: InterruptNode,
+) -> RunState:
+    frame = run.current_frame()
+    interrupt_request = build_interrupt_request(
+        step,
+        frame_id=frame.id,
+        state=run.state,
+        workflow_input=run.workflow_input,
+        context=frame_context_values(frame),
+    )
+    run.interrupt = interrupt_request
+    run.status = RunStatus.INTERRUPTED
+    frame.status = FrameStatus.INTERRUPTED
+    append_trace(
+        run,
+        frame_id=frame.id,
+        node_id=frame.node_id,
+        step_type=step.type,
+        resolved_input=interrupt_request.payload,
+        outcome="interrupt",
+        next_node_id=frame.node_id,
+        output=interrupt_request.payload,
+        state_changes={},
+    )
+    return run