langgraph: use input schema from conditional edge (#2516)

Currently we ignore the input schema in the branch and instead use the
input schema from the previous node (or overall graph schema)

This change makes the input schema to branches respected. This means
that if you try to pass extra keys and they're NOT in the input schema,
you will receive an error. If you don't provide an annotation in the
router, it will fall back to the previous node's input schema / full
graph state schema

Alternative solution is to just ignore the input schema in the router
altogether (including ignoring the schema from previous node / full
graph), but personally I find it more confusing.

---------

Co-authored-by: Nuno Campos <nuno@langchain.dev>

libs/langgraph/langgraph/graph/branch.py

@@ -0,0 +1,215 @@
+import asyncio
+from inspect import (
+    isfunction,
+    ismethod,
+    signature,
+)
+from types import FunctionType
+from typing import (
+    Any,
+    Awaitable,
+    Callable,
+    Hashable,
+    Literal,
+    NamedTuple,
+    Optional,
+    Sequence,
+    Type,
+    Union,
+    cast,
+    get_args,
+    get_origin,
+    get_type_hints,
+)
+
+from langchain_core.runnables import (
+    Runnable,
+    RunnableConfig,
+    RunnableLambda,
+)
+
+from langgraph.constants import END, START
+from langgraph.errors import InvalidUpdateError
+from langgraph.pregel.write import ChannelWrite
+from langgraph.types import Send
+from langgraph.utils.runnable import (
+    RunnableCallable,
+)
+
+
+def _get_branch_path_input_schema(
+    path: Union[
+        Callable[..., Union[Hashable, list[Hashable]]],
+        Callable[..., Awaitable[Union[Hashable, list[Hashable]]]],
+        Runnable[Any, Union[Hashable, list[Hashable]]],
+    ],
+) -> Optional[Type[Any]]:
+    input = None
+    # detect input schema annotation in the branch callable
+    try:
+        callable_: Optional[
+            Union[
+                Callable[..., Union[Hashable, list[Hashable]]],
+                Callable[..., Awaitable[Union[Hashable, list[Hashable]]]],
+            ]
+        ] = None
+        if isinstance(path, (RunnableCallable, RunnableLambda)):
+            if isfunction(path.func) or ismethod(path.func):
+                callable_ = path.func
+            elif (callable_method := getattr(path.func, "__call__", None)) and ismethod(
+                callable_method
+            ):
+                callable_ = callable_method
+            elif isfunction(path.afunc) or ismethod(path.afunc):
+                callable_ = path.afunc
+            elif (
+                callable_method := getattr(path.afunc, "__call__", None)
+            ) and ismethod(callable_method):
+                callable_ = callable_method
+        elif callable(path):
+            callable_ = path
+
+        if callable_ is not None and (hints := get_type_hints(callable_)):
+            first_parameter_name = next(
+                iter(signature(cast(FunctionType, callable_)).parameters.keys())
+            )
+            if input_hint := hints.get(first_parameter_name):
+                if isinstance(input_hint, type) and get_type_hints(input_hint):
+                    input = input_hint
+    except (TypeError, StopIteration):
+        pass
+
+    return input
+
+
+class Branch(NamedTuple):
+    path: Runnable[Any, Union[Hashable, list[Hashable]]]
+    ends: Optional[dict[Hashable, str]]
+    then: Optional[str] = None
+    input_schema: Optional[Type[Any]] = None
+
+    @classmethod
+    def from_path(
+        cls,
+        path: Runnable[Any, Union[Hashable, list[Hashable]]],
+        path_map: Optional[Union[dict[Hashable, str], list[str]]],
+        then: Optional[str] = None,
+        infer_schema: bool = False,
+    ) -> "Branch":
+        # coerce path_map to a dictionary
+        path_map_: Optional[dict[Hashable, str]] = None
+        try:
+            if isinstance(path_map, dict):
+                path_map_ = path_map.copy()
+            elif isinstance(path_map, list):
+                path_map_ = {name: name for name in path_map}
+            else:
+                # find func
+                func: Optional[Callable] = None
+                if isinstance(path, (RunnableCallable, RunnableLambda)):
+                    func = path.func or path.afunc
+                if func is not None:
+                    # find callable method
+                    if (cal := getattr(path, "__call__", None)) and ismethod(cal):
+                        func = cal
+                    # get the return type
+                    if rtn_type := get_type_hints(func).get("return"):
+                        if get_origin(rtn_type) is Literal:
+                            path_map_ = {name: name for name in get_args(rtn_type)}
+        except Exception:
+            pass
+        # infer input schema
+        input_schema = _get_branch_path_input_schema(path) if infer_schema else None
+        # create branch
+        return cls(path=path, ends=path_map_, then=then, input_schema=input_schema)
+
+    def run(
+        self,
+        writer: Callable[
+            [Sequence[Union[str, Send]], RunnableConfig], Optional[ChannelWrite]
+        ],
+        reader: Optional[Callable[[RunnableConfig], Any]] = None,
+    ) -> RunnableCallable:
+        return ChannelWrite.register_writer(
+            RunnableCallable(
+                func=self._route,
+                afunc=self._aroute,
+                writer=writer,
+                reader=reader,
+                name=None,
+                trace=False,
+            )
+        )
+
+    def _route(
+        self,
+        input: Any,
+        config: RunnableConfig,
+        *,
+        reader: Optional[Callable[[RunnableConfig], Any]],
+        writer: Callable[
+            [Sequence[Union[str, Send]], RunnableConfig], Optional[ChannelWrite]
+        ],
+    ) -> Runnable:
+        if reader:
+            value = reader(config)
+            # passthrough additional keys from node to branch
+            # only doable when using dict states
+            if (
+                isinstance(value, dict)
+                and isinstance(input, dict)
+                and self.input_schema is None
+            ):
+                value = {**input, **value}
+        else:
+            value = input
+        result = self.path.invoke(value, config)
+        return self._finish(writer, input, result, config)
+
+    async def _aroute(
+        self,
+        input: Any,
+        config: RunnableConfig,
+        *,
+        reader: Optional[Callable[[RunnableConfig], Any]],
+        writer: Callable[
+            [Sequence[Union[str, Send]], RunnableConfig], Optional[ChannelWrite]
+        ],
+    ) -> Runnable:
+        if reader:
+            value = await asyncio.to_thread(reader, config)
+            # passthrough additional keys from node to branch
+            # only doable when using dict states
+            if (
+                isinstance(value, dict)
+                and isinstance(input, dict)
+                and self.input_schema is None
+            ):
+                value = {**input, **value}
+        else:
+            value = input
+        result = await self.path.ainvoke(value, config)
+        return self._finish(writer, input, result, config)
+
+    def _finish(
+        self,
+        writer: Callable[
+            [Sequence[Union[str, Send]], RunnableConfig], Optional[ChannelWrite]
+        ],
+        input: Any,
+        result: Any,
+        config: RunnableConfig,
+    ) -> Union[Runnable, Any]:
+        if not isinstance(result, (list, tuple)):
+            result = [result]
+        if self.ends:
+            destinations: Sequence[Union[Send, str]] = [
+                r if isinstance(r, Send) else self.ends[r] for r in result
+            ]
+        else:
+            destinations = cast(Sequence[Union[Send, str]], result)
+        if any(dest is None or dest == START for dest in destinations):
+            raise ValueError("Branch did not return a valid destination")
+        if any(p.node == END for p in destinations if isinstance(p, Send)):
+            raise InvalidUpdateError("Cannot send a packet to the END node")
+        return writer(destinations, config) or input

libs/langgraph/langgraph/graph/graph.py

@@ -1,4 +1,3 @@
-import asyncio
 import logging
 from collections import defaultdict
 from typing import (
@@ -6,15 +5,11 @@ from typing import (
     Awaitable,
     Callable,
     Hashable,
-    Literal,
     NamedTuple,
     Optional,
     Sequence,
     Union,
     cast,
-    get_args,
-    get_origin,
-    get_type_hints,
     overload,
 )
 
@@ -34,12 +29,12 @@ from langgraph.constants import (
     TAG_HIDDEN,
     Send,
 )
-from langgraph.errors import InvalidUpdateError
+from langgraph.graph.branch import Branch
 from langgraph.pregel import Channel, Pregel
 from langgraph.pregel.read import PregelNode
 from langgraph.pregel.write import ChannelWrite, ChannelWriteEntry
 from langgraph.types import All, Checkpointer
-from langgraph.utils.runnable import RunnableCallable, RunnableLike, coerce_to_runnable
+from langgraph.utils.runnable import RunnableLike, coerce_to_runnable
 
 logger = logging.getLogger(__name__)
 
@@ -50,95 +45,6 @@ class NodeSpec(NamedTuple):
     ends: Optional[Union[tuple[str, ...], dict[str, str]]] = EMPTY_SEQ
 
 
-class Branch(NamedTuple):
-    path: Runnable[Any, Union[Hashable, list[Hashable]]]
-    ends: Optional[dict[Hashable, str]]
-    then: Optional[str] = None
-
-    def run(
-        self,
-        writer: Callable[
-            [Sequence[Union[str, Send]], RunnableConfig], Optional[ChannelWrite]
-        ],
-        reader: Optional[Callable[[RunnableConfig], Any]] = None,
-    ) -> RunnableCallable:
-        return ChannelWrite.register_writer(
-            RunnableCallable(
-                func=self._route,
-                afunc=self._aroute,
-                writer=writer,
-                reader=reader,
-                name=None,
-                trace=False,
-            )
-        )
-
-    def _route(
-        self,
-        input: Any,
-        config: RunnableConfig,
-        *,
-        reader: Optional[Callable[[RunnableConfig], Any]],
-        writer: Callable[
-            [Sequence[Union[str, Send]], RunnableConfig], Optional[ChannelWrite]
-        ],
-    ) -> Runnable:
-        if reader:
-            value = reader(config)
-            # passthrough additional keys from node to branch
-            # only doable when using dict states
-            if isinstance(value, dict) and isinstance(input, dict):
-                value = {**input, **value}
-        else:
-            value = input
-        result = self.path.invoke(value, config)
-        return self._finish(writer, input, result, config)
-
-    async def _aroute(
-        self,
-        input: Any,
-        config: RunnableConfig,
-        *,
-        reader: Optional[Callable[[RunnableConfig], Any]],
-        writer: Callable[
-            [Sequence[Union[str, Send]], RunnableConfig], Optional[ChannelWrite]
-        ],
-    ) -> Runnable:
-        if reader:
-            value = await asyncio.to_thread(reader, config)
-            # passthrough additional keys from node to branch
-            # only doable when using dict states
-            if isinstance(value, dict) and isinstance(input, dict):
-                value = {**input, **value}
-        else:
-            value = input
-        result = await self.path.ainvoke(value, config)
-        return self._finish(writer, input, result, config)
-
-    def _finish(
-        self,
-        writer: Callable[
-            [Sequence[Union[str, Send]], RunnableConfig], Optional[ChannelWrite]
-        ],
-        input: Any,
-        result: Any,
-        config: RunnableConfig,
-    ) -> Union[Runnable, Any]:
-        if not isinstance(result, (list, tuple)):
-            result = [result]
-        if self.ends:
-            destinations: Sequence[Union[Send, str]] = [
-                r if isinstance(r, Send) else self.ends[r] for r in result
-            ]
-        else:
-            destinations = cast(Sequence[Union[Send, str]], result)
-        if any(dest is None or dest == START for dest in destinations):
-            raise ValueError("Branch did not return a valid destination")
-        if any(p.node == END for p in destinations if isinstance(p, Send)):
-            raise InvalidUpdateError("Cannot send a packet to the END node")
-        return writer(destinations, config) or input
-
-
 class Graph:
     def __init__(self) -> None:
         self.nodes: dict[str, NodeSpec] = {}
@@ -267,25 +173,7 @@ class Graph:
                 "Adding an edge to a graph that has already been compiled. This will "
                 "not be reflected in the compiled graph."
             )
-        # coerce path_map to a dictionary
-        try:
-            if isinstance(path_map, dict):
-                path_map_ = path_map.copy()
-            elif isinstance(path_map, list):
-                path_map_ = {name: name for name in path_map}
-            elif isinstance(path, Runnable):
-                path_map_ = None
-            elif rtn_type := get_type_hints(path.__call__).get(  # type: ignore[operator]
-                "return"
-            ) or get_type_hints(path).get("return"):
-                if get_origin(rtn_type) is Literal:
-                    path_map_ = {name: name for name in get_args(rtn_type)}
-                else:
-                    path_map_ = None
-            else:
-                path_map_ = None
-        except Exception:
-            path_map_ = None
+
         # find a name for the condition
         path = coerce_to_runnable(path, name=None, trace=True)
         name = path.name or "condition"
@@ -295,7 +183,7 @@ class Graph:
                 f"Branch with name `{path.name}` already exists for node " f"`{source}`"
             )
         # save it
-        self.branches[source][name] = Branch(path, path_map_, then)
+        self.branches[source][name] = Branch.from_path(path, path_map, then, False)
         return self
 
     def set_entry_point(self, key: str) -> Self:

libs/langgraph/langgraph/graph/state.py

@@ -7,7 +7,9 @@ from inspect import isclass, isfunction, ismethod, signature
 from types import FunctionType
 from typing import (
     Any,
+    Awaitable,
     Callable,
+    Hashable,
     Literal,
     NamedTuple,
     Optional,
@@ -40,7 +42,14 @@ from langgraph.errors import (
     ParentCommand,
     create_error_message,
 )
-from langgraph.graph.graph import END, START, Branch, CompiledGraph, Graph, Send
+from langgraph.graph.branch import Branch
+from langgraph.graph.graph import (
+    END,
+    START,
+    CompiledGraph,
+    Graph,
+    Send,
+)
 from langgraph.managed.base import (
     ChannelKeyPlaceholder,
     ChannelTypePlaceholder,
@@ -461,6 +470,57 @@ class StateGraph(Graph):
         self.waiting_edges.add((tuple(start_key), end_key))
         return self
 
+    def add_conditional_edges(
+        self,
+        source: str,
+        path: Union[
+            Callable[..., Union[Hashable, list[Hashable]]],
+            Callable[..., Awaitable[Union[Hashable, list[Hashable]]]],
+            Runnable[Any, Union[Hashable, list[Hashable]]],
+        ],
+        path_map: Optional[Union[dict[Hashable, str], list[str]]] = None,
+        then: Optional[str] = None,
+    ) -> Self:
+        """Add a conditional edge from the starting node to any number of destination nodes.
+
+        Args:
+            source (str): The starting node. This conditional edge will run when
+                exiting this node.
+            path (Union[Callable, Runnable]): The callable that determines the next
+                node or nodes. If not specifying `path_map` it should return one or
+                more nodes. If it returns END, the graph will stop execution.
+            path_map (Optional[dict[Hashable, str]]): Optional mapping of paths to node
+                names. If omitted the paths returned by `path` should be node names.
+            then (Optional[str]): The name of a node to execute after the nodes
+                selected by `path`.
+
+        Returns:
+            Self: The instance of the graph, allowing for method chaining.
+
+        Note: Without typehints on the `path` function's return value (e.g., `-> Literal["foo", "__end__"]:`)
+            or a path_map, the graph visualization assumes the edge could transition to any node in the graph.
+
+        """  # noqa: E501
+        if self.compiled:
+            logger.warning(
+                "Adding an edge to a graph that has already been compiled. This will "
+                "not be reflected in the compiled graph."
+            )
+
+        # find a name for the condition
+        path = coerce_to_runnable(path, name=None, trace=True)
+        name = path.name or "condition"
+        # validate the condition
+        if name in self.branches[source]:
+            raise ValueError(
+                f"Branch with name `{path.name}` already exists for node " f"`{source}`"
+            )
+        # save it
+        self.branches[source][name] = Branch.from_path(path, path_map, then, True)
+        if schema := self.branches[source][name].input_schema:
+            self._add_schema(schema)
+        return self
+
     def add_sequence(
         self,
         nodes: Sequence[Union[RunnableLike, tuple[str, RunnableLike]]],
@@ -826,12 +886,12 @@ class CompiledStateGraph(CompiledGraph):
                     config, cast(Sequence[Union[Send, ChannelWriteEntry]], writes)
                 )
 
-        # attach branch publisher
-        schema = (
+        schema = branch.input_schema or (
             self.builder.nodes[start].input
             if start in self.builder.nodes
             else self.builder.schema
         )
+        # attach branch publisher
         self.nodes[start] |= branch.run(
             branch_writer,
             _get_state_reader(self.builder, schema) if with_reader else None,

libs/langgraph/tests/test_large_cases.py

@@ -2827,7 +2827,7 @@ def test_state_graph_packets(
         }
 
     # Define decision-making logic
-    def should_continue(data: AgentState) -> str:
+    def should_continue(data: dict) -> str:
         assert isinstance(data["session"], httpx.Client)
         assert (
             data["something_extra"] == "hi there"

libs/langgraph/tests/test_state.py

@@ -1,4 +1,5 @@
 import inspect
+import operator
 import warnings
 from dataclasses import dataclass, field
 from typing import Annotated as Annotated2
@@ -328,3 +329,74 @@ def test__get_node_name() -> None:
 
     # class method
     assert _get_node_name(MyClass().class_method) == "class_method"
+
+
+def test_input_schema_conditional_edge():
+    class OverallState(TypedDict):
+        foo: Annotated[int, operator.add]
+        bar: str
+
+    class PrivateState(TypedDict):
+        baz: str
+
+    builder = StateGraph(OverallState)
+
+    def node_1(state: OverallState):
+        return {"foo": 1, "baz": "bar"}
+
+    def node_2(state: PrivateState):
+        return {"foo": 1, "bar": state["baz"], "something_else": "meow"}
+
+    def node_3(state: OverallState):
+        return {"foo": 1}
+
+    def router(state: OverallState):
+        assert state == {"foo": 2, "bar": "bar"}
+        if state["foo"] == 2:
+            return "node_3"
+        else:
+            return "__end__"
+
+    builder.add_node(node_1)
+    builder.add_node(node_2)
+    builder.add_node(node_3)
+    builder.add_conditional_edges("node_2", router)
+    builder.add_edge("__start__", "node_1")
+    builder.add_edge("node_1", "node_2")
+    graph = builder.compile()
+    assert graph.invoke({"foo": 0}) == {"foo": 3, "bar": "bar"}
+
+
+def test_private_input_schema_conditional_edge():
+    class OverallState(TypedDict):
+        foo: Annotated[int, operator.add]
+        bar: str
+
+    class RouterState(TypedDict):
+        baz: str
+
+    class Node2State(TypedDict):
+        foo: Annotated[int, operator.add]
+        baz: str
+
+    builder = StateGraph(OverallState)
+
+    def node_1(state: OverallState):
+        return {"foo": 1, "baz": "meow"}
+
+    def node_2(state: Node2State):
+        return {"foo": 1, "bar": state["baz"]}
+
+    def router(state: RouterState):
+        assert state == {"baz": "meow"}
+        if state["baz"] == "meow":
+            return "node_2"
+        else:
+            return "__end__"
+
+    builder.add_node(node_1)
+    builder.add_node(node_2)
+    builder.add_conditional_edges("node_1", router)
+    builder.add_edge("__start__", "node_1")
+    graph = builder.compile()
+    assert graph.invoke({"foo": 0}) == {"foo": 2, "bar": "meow"}