route routes state to dedicated return nodes
also this may blow up the raw graph
This commit is contained in:
@@ -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)
|
||||
@@ -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(
|
||||
|
||||
@@ -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."
|
||||
|
||||
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user