diff --git a/README.md b/README.md index 6cb6caa..93a919c 100644 --- a/README.md +++ b/README.md @@ -269,7 +269,7 @@ Three of those need their edges stated, because the gap is where people get hurt **Budgets.** Tokens are charged without the node's cooperation: a LangChain callback is installed for the duration of every node, so any chat model invoked on that thread reports usage to the run's meter — including calls buried inside library code the node merely calls — and the ceiling is enforced at the node boundary. `max_seconds` is an interrupt, not a poll: SIGALRM on the main thread, an asynchronous exception otherwise, so a node parked in `time.sleep` or on a provider's socket is cut off at the deadline. Where it stops short: spend a provider never reports cannot be charged, a model invoked on a thread the node started itself is outside the callback's context, and an async exception cannot unwind a thread sitting inside a C call — it lands when that call returns. Even then the deadline holds at the node boundary: a node that overran does not get its writes into state. -**Routing.** The routers are code, which is the property that matters: no model output is ever consulted to pick an edge. `add_conditional_edge` checks the mapping where it is declared — an empty mapping is refused, every target must name a node the graph has or `END`, and a router annotated with what it returns (a `Literal`, an `Enum`) has those members held against the mapping's keys. Where it stops short: a router that declares nothing is not second-guessed, so the key it returns is only known when it returns one. That case is no longer a bare `KeyError` from inside LangGraph's branch machinery — it raises `GraphRoutingError` naming the node, the key and the keys there were — but it is still discovered by a run rather than by `add_conditional_edge`. +**Routing.** The routers are code, which is the property that matters: no model output is ever consulted to pick an edge. `add_conditional_edge` now validates conditional-edge mappings at declaration time: empty mappings are rejected, mapping targets must name an existing node or END, and if the router has a Literal/Enum return annotation, every allowed return value must be present as a mapping key. GraphARC still does not try to infer arbitrary runtime router return values from callable logic. **Typing.** Writes are checked in both directions: the dict a node returns is validated field by field against the state schema before it lands, and the state is validated again when the next node receives it. A value that doesn't fit raises `StateTypeError` naming the node, the field, the declared type and what arrived — and that includes the last node before `END`, so a bad type no longer escapes into the result. The validated value is what gets written, so a schema that says `int` means the result holds an `int`. The remaining gap is narrow and worth stating exactly: write-time validation is built from each field's *annotation*, so constraints carried in the annotation (`Annotated[int, Field(gt=0)]`) do bite, but a validator the state model declares for itself — `@field_validator`, `@model_validator` — is not run on a write. A node returning `{"slug": "NOT-LOWER"}` into a field whose validator demands lowercase is accepted, even though constructing the model directly with that value raises; the violation surfaces only when a later node receives the state and the whole model is rebuilt, which means one written by the last node before `END` still reaches the result. The write *allowlist* is GraphARC's; the *types* are Pydantic's. diff --git a/grapharc/runtime/graph.py b/grapharc/runtime/graph.py index a7b722f..0af024a 100644 --- a/grapharc/runtime/graph.py +++ b/grapharc/runtime/graph.py @@ -36,8 +36,6 @@ from __future__ import annotations import asyncio -import copy -import functools import inspect import threading import time @@ -46,7 +44,7 @@ from contextlib import asynccontextmanager from dataclasses import replace from enum import Enum -from typing import Any, Literal, get_args, get_origin, get_type_hints +from typing import Any, Literal, get_args, get_origin from langchain_core.runnables import RunnableConfig from langgraph.graph import END, START, StateGraph @@ -426,15 +424,61 @@ def add_conditional_edge( raise GraphCycleError( f"graph {self.name!r} is dag=True: conditional edges are not allowed" ) - self._check_mapping(source, router, mapping) - self._graph.add_conditional_edges( - source, self._checked_router(source, router, mapping), mapping - ) + self._validate_conditional_edge(source, router, mapping) + self._graph.add_conditional_edges(source, router, mapping) self._conditional_edges.extend( (source, target) for target in dict.fromkeys(mapping.values()) ) return self + def _validate_conditional_edge( + self, + source: str, + router: Callable[[Any], str], + mapping: dict[str, str], + ) -> None: + """Validate a conditional-edge mapping at declaration time.""" + if not mapping: + raise GraphRoutingError( + f"conditional edge from node {source!r} has an empty mapping" + ) + + for key, target in mapping.items(): + if target == END or target in self._nodes: + continue + raise GraphRoutingError( + f"conditional edge from node {source!r} uses mapping key {key!r} -> " + f"{target!r}, but {target!r} is not a node of graph {self.name!r}; " + f"valid destinations: {self._destinations()}" + ) + + try: + signature = inspect.signature(router) + except (TypeError, ValueError): + return + + annotation = signature.return_annotation + if annotation is inspect.Signature.empty: + return + + if get_origin(annotation) is Literal: + allowed = get_args(annotation) + elif isinstance(annotation, type) and issubclass(annotation, Enum): + allowed = [member.value for member in annotation] + else: + return + + if not all(isinstance(value, str) for value in allowed): + return + + missing = [value for value in allowed if value not in mapping] + if missing: + missing_repr = ", ".join(repr(value) for value in missing) + raise GraphRoutingError( + f"conditional edge from node {source!r} has router return values " + f"{missing_repr}, but mapping is missing those keys" + ) + def add_fanout_edge( self, source: str, diff --git a/tests/test_add_conditional_edge.py b/tests/test_add_conditional_edge.py new file mode 100644 index 0000000..56d56c5 --- /dev/null +++ b/tests/test_add_conditional_edge.py @@ -0,0 +1,30 @@ +import pytest +from pydantic import BaseModel + +from grapharc.runtime.graph import GraphARC, GraphRoutingError + + +def test_add_conditional_edge_rejects_unknown_target_at_add_time(): + class State(BaseModel): + value: str = "" + + arc = GraphARC(state_schema=State, name="test") + + arc.add_node("start", lambda state: {"value": "ok"}, writes={"value"}) + arc.add_node("done", lambda state: {"value": "ok"}, writes={"value"}) + + with pytest.raises(GraphRoutingError, match="not a node of graph") as exc_info: + arc.add_conditional_edge("start", lambda _: "go", {"go": "missing"}) + print(f"\nMESSAGE: {exc_info.value}") + + +def test_add_conditional_edge_accepts_known_target_mapping(): + class State(BaseModel): + value: str = "" + + arc = GraphARC(state_schema=State, name="test") + + arc.add_node("start", lambda state: {"value": "ok"}, writes={"value"}) + arc.add_node("done", lambda state: {"value": "ok"}, writes={"value"}) + + arc.add_conditional_edge("start", lambda _: "go", {"go": "done"}) \ No newline at end of file