Skip to content

[#147 phase 3d] Typed-edge validation: catch branch-mapping mistakes via Literal[...] return annotations #238

@miguelgfierro

Description

@miguelgfierro

Follow-up from #147. When a state-pipeline node has a deterministic branch-decision return shape, that decision should be expressible in the type system so mypy/pyright catch wiring mistakes at static-check time.

Why

Today `.branch(source, router, mapping)` is validated at runtime — the router runs, returns a label, the mapping is looked up. If the mapping is wrong or has a typo, the pipeline fails at the FIRST invocation that hits that branch, often deep in production logs:

```python
def route(state):
return "escalate" if state.intent == "complaint" else "answer"

PipelineBuilder("agent", state=AgentState)
.add_node(classify).add_node(answer).add_node(escalate)
.branch(classify, route, {"answer": answer, "escallate": escalate}) # typo!
.build()

.build() succeeds. Failure only surfaces when an actual complaint arrives.

```

Pydantic-graph's standout feature is that node return annotations declare the outgoing edges: run() -> Answer | Escalate. We can offer the same static check inside our string-id API by inspecting `Literal[...]` return annotations on routers.

Proposed behaviour

When a router function is annotated with `Literal["a", "b", ...]` as its return type:

```python
def route(state) -> Literal["answer", "escalate"]:
return "escalate" if state.intent == "complaint" else "answer"
```

At `.build()`, the builder:

  1. Inspects `typing.get_type_hints(route)` to extract the `Literal` values.
  2. If `.branch(source, route)` was called WITHOUT a mapping, every node in the literal must exist in the DAG — otherwise `.build()` raises with a clear error.
  3. If `.branch(source, route, mapping=...)` was called, the mapping's keys must be a superset of (or equal to) the literal values — otherwise `.build()` raises.
  4. If the router has no annotation, the current runtime-only validation continues unchanged.

```python

Caught at .build() time, not first invocation:

PipelineError: Router for 'classify' declared Literal["answer", "escalate"] but mapping
keys are {"answer", "escallate"} — missing 'escalate', extra 'escallate'
```

Implementation notes

  • Inspect via `typing.get_type_hints(router).get("return")`. Recognize `Literal` via `typing.get_origin` / `typing.get_args`.
  • Add a build-time validator alongside the existing connectivity check in `StatePipeline._validate`.
  • Tests: a node with a correct `Literal` annotation builds; a typo in the mapping raises at `.build()`; an unannotated router preserves current behaviour; `Literal` values not in the DAG raise.

Scope estimate

~100 LOC + ~80 LOC tests. No new dependencies. No runtime cost (validation is build-only).

Out of scope

  • Inferring edges from `Union[NodeA, NodeB]` style — that's pydantic-graph's pattern and would require a different node shape.
  • Anything involving the `pyright` plugin API.

Metadata

Metadata

Assignees

No one assigned

    Labels

    enhancementNew feature or request

    Type

    No type
    No fields configured for issues without a type.

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions