route routes state to dedicated return nodes

also this may blow up the raw graph
This commit is contained in:
lda
2026-05-17 10:08:22 +07:00 Verified
parent 5859718200
commit a8b5923dd8
4 changed files with 150 additions and 3 deletions
+85
View File
@@ -0,0 +1,85 @@
from __future__ import annotations
from typing import Literal
from pydantic import BaseModel
from wf_authoring import WorkflowBuilder, node, state
class ToolInput(BaseModel):
text: str
class ToolOutput(BaseModel):
status: Literal["done", "needs_input", "failed"]
message: str
class WrapperState(BaseModel):
status: str
message: str
class WrapperOutput(BaseModel):
message: str
@node
def raw_tool(input: ToolInput) -> ToolOutput:
"""Stand in for a thin upstream MCP tool wrapper returning provider status."""
if input.text.endswith("?"):
return ToolOutput(status="needs_input", message="Need clarification")
if not input.text.strip():
return ToolOutput(status="failed", message="No text supplied")
return ToolOutput(status="done", message=input.text.upper())
@node
def done(input: WrapperOutput) -> WrapperOutput:
"""Expose a normalized success payload."""
return input
@node
def needs_input(input: WrapperOutput) -> WrapperOutput:
"""Expose a normalized clarification payload."""
return input
@node
def failed(input: WrapperOutput) -> WrapperOutput:
"""Expose a normalized failure payload."""
return input
def build_wrapper() -> WorkflowBuilder:
"""Build a node-like wrapper graph around a status-returning raw tool."""
graph = WorkflowBuilder(
name="status_wrapper",
input_schema=ToolInput,
state_schema=WrapperState,
output_schema=WrapperOutput,
)
tool = graph.use(raw_tool)
graph.route(
state("status"),
{
"done": graph.use(done, id="done"),
"needs_input": graph.use(needs_input, id="needs_input"),
},
default=graph.use(failed, id="failed"),
)
graph.set_entry_point(tool)
graph.connect(tool, "ok", "condition")
graph.connect("done", "ok", "__end__")
graph.connect("needs_input", "ok", "__end__")
graph.connect("failed", "ok", "__end__")
return graph
if __name__ == "__main__":
workflow = build_wrapper()
for text in ("hello", "clarify?", ""):
run = workflow.execute({"text": text})
print(text, run.status.value, run.output)
+32 -1
View File
@@ -20,7 +20,7 @@ from wf_core import (
from wf_core.errors import WorkflowExecutionError
from wf_core.models.conditions import Condition as CoreCondition
from ..dsl import Expr, PathArg, compile_condition
from ..dsl import Expr, PathArg, PathExpr, compile_condition
from ..nodes.callables import SyncRegistryHandler
from ..nodes.registry import build_registry
from ..schemas import SchemaLike, StateSchemaLike, schema_ref_from, state_schema_from
@@ -210,6 +210,37 @@ class WorkflowBuilder:
resolved_targets[outcome] = cast(StepRef, resolved)
return resolved_targets
def route(
self,
value: PathExpr,
cases: Mapping[object, BranchRef],
*,
default: BranchRef,
) -> dict[object, StepRef]:
"""Route graph data by equality checks compiled to condition nodes.
`branch()` wires outcomes already produced by a node. `route()` is the
companion for ordinary graph data: it compares one path against case
values, expands those checks into a condition chain, and wires the first
matching target plus a required fallback.
"""
resolved_targets: dict[object, StepRef] = {}
default_target = self.use(default) if is_node_spec(default) else default
previous_condition: ConditionNode | None = None
for case_value, target in cases.items():
condition = self.condition(check=value.eq(case_value))
resolved = self.use(target) if is_node_spec(target) else target
if previous_condition is not None:
self.connect(previous_condition, "false", condition)
self.connect(condition, "true", cast(StepRef, resolved))
previous_condition = condition
resolved_targets[case_value] = cast(StepRef, resolved)
if previous_condition is None:
raise ValueError("WorkflowBuilder.route requires at least one case")
self.connect(previous_condition, "false", cast(StepRef, default_target))
resolved_targets["default"] = cast(StepRef, default_target)
return resolved_targets
def compile(self) -> Workflow:
if self.start is None:
raise WorkflowExecutionError(
-1
View File
@@ -307,7 +307,6 @@ def _artifact_capability_id(artifact: WorkflowArtifact) -> str:
return f"workflow.{artifact.id}.v{artifact.version}"
# this feels like a hack
def _parse_artifact_capability_id(qualified_name: str) -> tuple[str, int] | None:
"""Parse the stable `workflow.<artifact_id>.v<version>` capability name."""
prefix = "workflow."
+33 -1
View File
@@ -2,7 +2,7 @@ from __future__ import annotations
import pytest
from wf_authoring import WorkflowBuilder
from wf_authoring import WorkflowBuilder, state
from tests.authoring.helpers import (
AutoBindInput,
@@ -64,3 +64,35 @@ def test_builder_branch_warns_on_empty_branch_map() -> None:
targets = builder.branch(router, {})
assert targets == {}
def test_builder_route_expands_state_value_cases_into_condition_chain() -> None:
builder = WorkflowBuilder(
name="route_demo",
input_schema=AutoBindInput,
state_schema=AutoBindState,
output_schema=AutoBindOutput,
)
left = builder.use(auto_bind_node, id="left")
right = builder.use(auto_bind_node, id="right")
fallback = builder.use(auto_bind_node, id="fallback")
targets = builder.route(
state("value"),
{"left": left, "right": right},
default=fallback,
)
assert [node.id for node in builder.nodes if node.type == "condition"] == [
"condition",
"condition_2",
]
assert [(edge.from_, edge.outcome, edge.to) for edge in builder.edges] == [
("condition", "true", "left"),
("condition", "false", "condition_2"),
("condition_2", "true", "right"),
("condition_2", "false", "fallback"),
]
assert targets["left"] is left
assert targets["right"] is right
assert targets["default"] is fallback