Improve type checking on graph init and invoke/stream (#4932)

.gitignore

@@ -181,3 +181,4 @@ Chinook.db
 .vercel
 .turbo
 .editorconfig
+.scratch

libs/langgraph/langgraph/graph/state.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import inspect
 import logging
 import typing
@@ -10,9 +12,11 @@ from types import FunctionType
 from typing import (
     Any,
     Callable,
+    Generic,
     Literal,
     NamedTuple,
     Optional,
+    Protocol,
     Union,
     cast,
     get_args,
@@ -23,7 +27,7 @@ from typing import (
 
 from langchain_core.runnables import Runnable, RunnableConfig
 from pydantic import BaseModel
-from typing_extensions import Self
+from typing_extensions import Self, TypeAlias
 
 from langgraph.cache.base import BaseCache
 from langgraph.channels.base import BaseChannel
@@ -65,14 +69,23 @@ from langgraph.pregel.write import (
     ChannelWriteTupleEntry,
 )
 from langgraph.store.base import BaseStore
-from langgraph.types import All, CachePolicy, Checkpointer, Command, RetryPolicy, Send
+from langgraph.types import (
+    All,
+    CachePolicy,
+    Checkpointer,
+    Command,
+    RetryPolicy,
+    Send,
+    StreamWriter,
+)
+from langgraph.typing import InputT, StateT, StateT_contra, Unset
 from langgraph.utils.fields import (
     get_cached_annotated_keys,
     get_field_default,
     get_update_as_tuples,
 )
 from langgraph.utils.pydantic import create_model
-from langgraph.utils.runnable import RunnableLike, coerce_to_runnable
+from langgraph.utils.runnable import coerce_to_runnable
 
 logger = logging.getLogger(__name__)
 
@@ -89,17 +102,77 @@ def _warn_invalid_state_schema(schema: Union[type[Any], Any]) -> None:
     )
 
 
-def _get_node_name(node: RunnableLike) -> str:
-    if isinstance(node, Runnable):
-        return node.get_name()
-    elif callable(node):
+class _StateNode(Protocol[StateT_contra]):
+    def __call__(self, state: StateT_contra) -> Any: ...
+
+
+class _NodeWithConfig(Protocol[StateT_contra]):
+    def __call__(self, state: StateT_contra, config: RunnableConfig) -> Any: ...
+
+
+class _NodeWithWriter(Protocol[StateT_contra]):
+    def __call__(self, state: StateT_contra, *, writer: StreamWriter) -> Any: ...
+
+
+class _NodeWithStore(Protocol[StateT_contra]):
+    def __call__(self, state: StateT_contra, *, store: BaseStore) -> Any: ...
+
+
+class _NodeWithWriterStore(Protocol[StateT_contra]):
+    def __call__(
+        self, state: StateT_contra, *, writer: StreamWriter, store: BaseStore
+    ) -> Any: ...
+
+
+class _NodeWithConfigWriter(Protocol[StateT_contra]):
+    def __call__(
+        self, state: StateT_contra, *, config: RunnableConfig, writer: StreamWriter
+    ) -> Any: ...
+
+
+class _NodeWithConfigStore(Protocol[StateT_contra]):
+    def __call__(
+        self, state: StateT_contra, *, config: RunnableConfig, store: BaseStore
+    ) -> Any: ...
+
+
+class _NodeWithConfigWriterStore(Protocol[StateT_contra]):
+    def __call__(
+        self,
+        state: StateT_contra,
+        *,
+        config: RunnableConfig,
+        writer: StreamWriter,
+        store: BaseStore,
+    ) -> Any: ...
+
+
+# TODO: we probably don't want to explicitly support the config / store signatures once
+# we move to adding a context arg. Maybe what we do is we add support for kwargs with param spec
+# this is purely for typing purposes though, so can easily change in the coming weeks.
+StateNode: TypeAlias = Union[
+    _StateNode[StateT_contra],
+    _NodeWithConfig[StateT_contra],
+    _NodeWithWriter[StateT_contra],
+    _NodeWithStore[StateT_contra],
+    _NodeWithWriterStore[StateT_contra],
+    _NodeWithConfigWriter[StateT_contra],
+    _NodeWithConfigStore[StateT_contra],
+    _NodeWithConfigWriterStore[StateT_contra],
+]
+
+
+def _get_node_name(node: StateNode) -> str:
+    try:
         return getattr(node, "__name__", node.__class__.__name__)
-    else:
+    except AttributeError:
         raise TypeError(f"Unsupported node type: {type(node)}")
 
 
 class StateNodeSpec(NamedTuple):
-    runnable: Runnable
+    # TODO: rename this callable, also move away from NamedTuple so that we can use
+    # a generic StateNode, so maybe a dataclass
+    runnable: StateNode
     metadata: Optional[dict[str, Any]]
     input: type[Any]
     retry_policy: Optional[Union[RetryPolicy, Sequence[RetryPolicy]]]
@@ -108,7 +181,7 @@ class StateNodeSpec(NamedTuple):
     defer: bool = False
 
 
-class StateGraph:
+class StateGraph(Generic[StateT, InputT]):
     """A graph whose nodes communicate by reading and writing to a shared state.
     The signature of each node is State -> Partial<State>.
 
@@ -169,16 +242,14 @@ class StateGraph:
 
     def __init__(
         self,
-        state_schema: type[Any],
-        config_schema: Optional[type[Any]] = None,
+        state_schema: type[StateT],
+        config_schema: type[Any] | None = None,
         *,
-        input: Optional[type[Any]] = None,
-        output: Optional[type[Any]] = None,
+        input: type[InputT] | None = None,
+        output: type[Any] | None = None,
     ) -> None:
-        if input is None:
-            input = state_schema
-        if output is None:
-            output = state_schema
+        input = input or state_schema
+        output = output or state_schema
 
         self.nodes = {}
         self.edges = set[tuple[str, str]]()
@@ -238,7 +309,7 @@ class StateGraph:
     @overload
     def add_node(
         self,
-        node: RunnableLike,
+        node: StateNode[StateT],
         *,
         defer: bool = False,
         metadata: Optional[dict[str, Any]] = None,
@@ -256,7 +327,7 @@ class StateGraph:
     def add_node(
         self,
         node: str,
-        action: RunnableLike,
+        action: StateNode[StateT],
         *,
         defer: bool = False,
         metadata: Optional[dict[str, Any]] = None,
@@ -270,8 +341,8 @@ class StateGraph:
 
     def add_node(
         self,
-        node: Union[str, RunnableLike],
-        action: Optional[RunnableLike] = None,
+        node: Union[str, StateNode[StateT]],
+        action: Optional[StateNode[StateT]] = None,
         *,
         defer: bool = False,
         metadata: Optional[dict[str, Any]] = None,
@@ -417,7 +488,7 @@ class StateGraph:
         if input is not None:
             self._add_schema(input)
         self.nodes[node] = StateNodeSpec(
-            coerce_to_runnable(action, name=node, trace=False),
+            coerce_to_runnable(action, name=node, trace=False),  # type: ignore
             metadata,
             input=input or self.schema,
             retry_policy=retry,
@@ -531,12 +602,12 @@ class StateGraph:
 
     def add_sequence(
         self,
-        nodes: Sequence[Union[RunnableLike, tuple[str, RunnableLike]]],
+        nodes: Sequence[Union[StateNode[StateT], tuple[str, StateNode[StateT]]]],
     ) -> Self:
         """Add a sequence of nodes that will be executed in the provided order.
 
         Args:
-            nodes: A sequence of RunnableLike objects (e.g. a LangChain Runnable or a callable) or (name, RunnableLike) tuples.
+            nodes: A sequence of StateNodes (callables that accept a state arg) or (name, StateNode) tuples.
                 If no names are provided, the name will be inferred from the node object (e.g. a runnable or a callable name).
                 Each node will be executed in the order provided.
 
@@ -669,6 +740,32 @@ class StateGraph:
         self.compiled = True
         return self
 
+    @overload
+    def compile(
+        self: StateGraph[StateT, Unset],
+        checkpointer: Checkpointer = None,
+        *,
+        cache: Optional[BaseCache] = None,
+        store: Optional[BaseStore] = None,
+        interrupt_before: Optional[Union[All, list[str]]] = None,
+        interrupt_after: Optional[Union[All, list[str]]] = None,
+        debug: bool = False,
+        name: Optional[str] = None,
+    ) -> CompiledStateGraph[StateT, StateT]: ...
+
+    @overload
+    def compile(
+        self: StateGraph[StateT, InputT],
+        checkpointer: Checkpointer = None,
+        *,
+        cache: Optional[BaseCache] = None,
+        store: Optional[BaseStore] = None,
+        interrupt_before: Optional[Union[All, list[str]]] = None,
+        interrupt_after: Optional[Union[All, list[str]]] = None,
+        debug: bool = False,
+        name: Optional[str] = None,
+    ) -> CompiledStateGraph[StateT, InputT]: ...
+
     def compile(
         self,
         checkpointer: Checkpointer = None,
@@ -679,7 +776,7 @@ class StateGraph:
         interrupt_after: Optional[Union[All, list[str]]] = None,
         debug: bool = False,
         name: Optional[str] = None,
-    ) -> "CompiledStateGraph":
+    ) -> Union[CompiledStateGraph[StateT, StateT], CompiledStateGraph[StateT, InputT]]:
         """Compiles the state graph into a `CompiledStateGraph` object.
 
         The compiled graph implements the `Runnable` interface and can be invoked,
@@ -731,7 +828,8 @@ class StateGraph:
             ]
         )
 
-        compiled = CompiledStateGraph(
+        ResolvedInputT: Union[type[InputT], type[StateT]] = self.input or self.schema
+        compiled = CompiledStateGraph[StateT, ResolvedInputT](  # type: ignore[valid-type]
             builder=self,
             schema_to_mapper={},
             config_type=self.config_schema,
@@ -779,14 +877,14 @@ class StateGraph:
         return compiled.validate()
 
 
-class CompiledStateGraph(Pregel):
-    builder: StateGraph
+class CompiledStateGraph(Pregel[InputT], Generic[StateT, InputT]):
+    builder: StateGraph[StateT, InputT]
     schema_to_mapper: dict[type[Any], Optional[Callable[[Any], Any]]]
 
     def __init__(
         self,
         *,
-        builder: StateGraph,
+        builder: StateGraph[StateT, InputT],
         schema_to_mapper: dict[type[Any], Optional[Callable[[Any], Any]]],
         **kwargs: Any,
     ) -> None:
@@ -915,7 +1013,7 @@ class CompiledStateGraph(Pregel):
                 metadata=node.metadata,
                 retry_policy=node.retry_policy,
                 cache_policy=node.cache_policy,
-                bound=node.runnable,
+                bound=node.runnable,  # type: ignore[arg-type]
             )
         else:
             raise RuntimeError

libs/langgraph/langgraph/pregel/__init__.py

@@ -8,7 +8,7 @@ import weakref
 from collections import defaultdict, deque
 from collections.abc import AsyncIterator, Iterator, Mapping, Sequence
 from functools import partial
-from typing import Any, Callable, Union, cast, get_type_hints
+from typing import Any, Callable, Generic, Union, cast, get_type_hints
 from uuid import UUID, uuid5
 
 from langchain_core.globals import get_debug
@@ -107,6 +107,7 @@ from langgraph.types import (
     StreamChunk,
     StreamMode,
 )
+from langgraph.typing import InputT
 from langgraph.utils.config import (
     ensure_config,
     merge_configs,
@@ -297,7 +298,7 @@ class NodeBuilder:
         )
 
 
-class Pregel(PregelProtocol):
+class Pregel(PregelProtocol[InputT], Generic[InputT]):
     """Pregel manages the runtime behavior for LangGraph applications.
 
     ## Overview
@@ -738,7 +739,8 @@ class Pregel(PregelProtocol):
         }
 
     def copy(self, update: dict[str, Any] | None = None) -> Self:
-        attrs = {**self.__dict__, **(update or {})}
+        attrs = {k: v for k, v in self.__dict__.items() if k != "__orig_class__"}
+        attrs.update(update or {})
         return self.__class__(**attrs)
 
     def with_config(self, config: RunnableConfig | None = None, **kwargs: Any) -> Self:
@@ -2279,7 +2281,7 @@ class Pregel(PregelProtocol):
 
     def stream(
         self,
-        input: dict[str, Any] | Any,
+        input: InputT,
         config: RunnableConfig | None = None,
         *,
         stream_mode: StreamMode | list[StreamMode] | None = None,
@@ -2500,7 +2502,7 @@ class Pregel(PregelProtocol):
 
     async def astream(
         self,
-        input: dict[str, Any] | Any,
+        input: InputT,
         config: RunnableConfig | None = None,
         *,
         stream_mode: StreamMode | list[StreamMode] | None = None,
@@ -2736,7 +2738,7 @@ class Pregel(PregelProtocol):
 
     def invoke(
         self,
-        input: dict[str, Any] | Any,
+        input: InputT,
         config: RunnableConfig | None = None,
         *,
         stream_mode: StreamMode = "values",
@@ -2802,7 +2804,7 @@ class Pregel(PregelProtocol):
 
     async def ainvoke(
         self,
-        input: dict[str, Any] | Any,
+        input: InputT,
         config: RunnableConfig | None = None,
         *,
         stream_mode: StreamMode = "values",

libs/langgraph/langgraph/pregel/protocol.py

@@ -1,21 +1,19 @@
+from __future__ import annotations
+
 from abc import ABC, abstractmethod
 from collections.abc import AsyncIterator, Iterator, Sequence
-from typing import (
-    Any,
-    Optional,
-    Union,
-)
+from typing import Any, Generic, Optional, Union
 
 from langchain_core.runnables import Runnable, RunnableConfig
 from langchain_core.runnables.graph import Graph as DrawableGraph
 from typing_extensions import Self
 
 from langgraph.pregel.types import All, StateSnapshot, StateUpdate, StreamMode
+from langgraph.typing import InputT
 
 
-class PregelProtocol(
-    Runnable[Union[dict[str, Any], Any], Union[dict[str, Any], Any]], ABC
-):
+# TODO: remove Runnable inheritance here!
+class PregelProtocol(Runnable[InputT, Any], Generic[InputT], ABC):
     @abstractmethod
     def with_config(
         self, config: Optional[RunnableConfig] = None, **kwargs: Any
@@ -100,7 +98,7 @@ class PregelProtocol(
     @abstractmethod
     def stream(
         self,
-        input: Union[dict[str, Any], Any],
+        input: InputT,
         config: Optional[RunnableConfig] = None,
         *,
         stream_mode: Optional[Union[StreamMode, list[StreamMode]]] = None,
@@ -112,7 +110,7 @@ class PregelProtocol(
     @abstractmethod
     def astream(
         self,
-        input: Union[dict[str, Any], Any],
+        input: InputT,
         config: Optional[RunnableConfig] = None,
         *,
         stream_mode: Optional[Union[StreamMode, list[StreamMode]]] = None,
@@ -124,7 +122,7 @@ class PregelProtocol(
     @abstractmethod
     def invoke(
         self,
-        input: Union[dict[str, Any], Any],
+        input: InputT,
         config: Optional[RunnableConfig] = None,
         *,
         interrupt_before: Optional[Union[All, Sequence[str]]] = None,
@@ -134,7 +132,7 @@ class PregelProtocol(
     @abstractmethod
     async def ainvoke(
         self,
-        input: Union[dict[str, Any], Any],
+        input: InputT,
         config: Optional[RunnableConfig] = None,
         *,
         interrupt_before: Optional[Union[All, Sequence[str]]] = None,

libs/langgraph/langgraph/typing.py

@@ -0,0 +1,70 @@
+from __future__ import annotations
+
+from dataclasses import Field
+from typing import (
+    Any,
+    ClassVar,
+    Protocol,
+    Union,
+)
+
+from pydantic import BaseModel
+from typing_extensions import TypeAlias, TypeVar
+
+
+class _TypedDictLikeV1(Protocol):
+    """Protocol to represent types that behave like TypedDicts
+
+    Version 1: using `ClassVar` for keys."""
+
+    __required_keys__: ClassVar[frozenset[str]]
+    __optional_keys__: ClassVar[frozenset[str]]
+
+
+class _TypedDictLikeV2(Protocol):
+    """Protocol to represent types that behave like TypedDicts
+
+    Version 2: not using `ClassVar` for keys."""
+
+    __required_keys__: frozenset[str]
+    __optional_keys__: frozenset[str]
+
+
+class _DataclassLike(Protocol):
+    """Protocol to represent types that behave like dataclasses.
+
+    Inspired by the private _DataclassT from dataclasses that uses a similar protocol as a bound."""
+
+    __dataclass_fields__: ClassVar[dict[str, Field[Any]]]
+
+
+StateLike: TypeAlias = Union[
+    _TypedDictLikeV1, _TypedDictLikeV2, _DataclassLike, BaseModel
+]
+"""Type alias for state-like types.
+
+It can either be a `TypedDict`, `dataclass`, or Pydantic `BaseModel`.
+Note: we cannot use either `TypedDict` or `dataclass` directly due to limitations in type checking."""
+
+
+class Unset:
+    """Sentinel representing an unset value."""
+
+
+UNSET = Unset()
+
+StateT = TypeVar("StateT", bound=StateLike)
+"""Type variable used to represent the state in a graph."""
+
+StateT_co = TypeVar("StateT_co", bound=StateLike, covariant=True)
+
+StateT_contra = TypeVar("StateT_contra", bound=StateLike, contravariant=True)
+
+InputT = TypeVar("InputT", bound=Union[StateLike, Unset], default=Unset)
+"""Type variable used to represent the input to a graph.
+
+In practice, InputT might be represented by either the `input_type` or `state_type` of a graph.
+If `input_type` is not specified, it defaults to `StateType`."""
+
+OutputT = TypeVar("OutputT", bound=Union[StateLike, Unset], default=Unset)
+"""Type variable used to represent the output of a graph."""

libs/langgraph/tests/test_state.py

@@ -6,7 +6,7 @@ from typing import Annotated, Any, Optional
 from typing import Annotated as Annotated2
 
 import pytest
-from langchain_core.runnables import RunnableConfig, RunnableLambda
+from langchain_core.runnables import RunnableConfig
 from pydantic import BaseModel
 from typing_extensions import NotRequired, Required, TypedDict
 
@@ -241,14 +241,6 @@ def test_state_schema_default_values(kw_only_: bool):
 
 
 def test__get_node_name() -> None:
-    # default runnable name
-    assert _get_node_name(RunnableLambda(func=lambda x: x)) == "RunnableLambda"
-    # custom runnable name
-    assert (
-        _get_node_name(RunnableLambda(name="my_runnable", func=lambda x: x))
-        == "my_runnable"
-    )
-
     # lambda
     assert _get_node_name(lambda x: x) == "<lambda>"
 

libs/langgraph/tests/test_type_checking.py

@@ -0,0 +1,105 @@
+from dataclasses import dataclass
+from operator import add
+from typing import Annotated, Any
+
+from langchain_core.runnables import RunnableConfig
+from pydantic import BaseModel
+from typing_extensions import TypedDict
+
+from langgraph.graph import StateGraph
+
+
+def test_typed_dict_state() -> None:
+    class TypedDictState(TypedDict):
+        info: Annotated[list[str], add]
+
+    graph_builder = StateGraph(TypedDictState)
+
+    def valid(state: TypedDictState) -> Any: ...
+
+    def valid_with_config(state: TypedDictState, config: RunnableConfig) -> Any: ...
+
+    def invalid() -> Any: ...
+
+    def invalid_node() -> Any: ...
+
+    graph_builder.add_node("valid", valid)
+    graph_builder.add_node("invalid", valid_with_config)
+    graph_builder.add_node("invalid_node", invalid_node)  # type: ignore[call-overload]
+    graph_builder.set_entry_point("valid")
+    graph = graph_builder.compile()
+
+    graph.invoke({"info": ["hello", "world"]})
+    graph.invoke({"invalid": "lalala"})  # type: ignore[arg-type]
+
+
+def test_dataclass_state() -> None:
+    @dataclass
+    class DataclassState:
+        info: Annotated[list[str], add]
+
+    def valid(state: DataclassState) -> Any: ...
+
+    def valid_with_config(state: DataclassState, config: RunnableConfig) -> Any: ...
+
+    def invalid() -> Any: ...
+
+    graph_builder = StateGraph(DataclassState)
+    graph_builder.add_node("valid", valid)
+    graph_builder.add_node("invalid", valid_with_config)
+    graph_builder.add_node("invalid_node", invalid)  # type: ignore[call-overload]
+
+    graph_builder.set_entry_point("valid")
+    graph = graph_builder.compile()
+
+    graph.invoke(DataclassState(info=["hello", "world"]))
+    graph.invoke({"invalid": 1})  # type: ignore[arg-type]
+    graph.invoke({"info": ["hello", "world"]})  # type: ignore[arg-type]
+
+
+def test_base_model_state() -> None:
+    class PydanticState(BaseModel):
+        info: Annotated[list[str], add]
+
+    def valid(state: PydanticState) -> Any: ...
+
+    def valid_with_config(state: PydanticState, config: RunnableConfig) -> Any: ...
+
+    def invalid() -> Any: ...
+
+    graph_builder = StateGraph(PydanticState)
+    graph_builder.add_node("valid", valid)
+    graph_builder.add_node("invalid", valid_with_config)
+    graph_builder.add_node("invalid_node", invalid)  # type: ignore[call-overload]
+
+    graph_builder.set_entry_point("valid")
+    graph = graph_builder.compile()
+
+    graph.invoke(PydanticState(info=["hello", "world"]))
+    graph.invoke({"invalid": 1})  # type: ignore[arg-type]
+    graph.invoke({"info": ["hello", "world"]})  # type: ignore[arg-type]
+
+
+def test_plain_class_not_allowed() -> None:
+    class NotAllowed:
+        info: Annotated[list[str], add]
+
+    StateGraph(NotAllowed)  # type: ignore[type-var]
+
+
+def test_input_state_specified() -> None:
+    class InputState(TypedDict):
+        something: int
+
+    class State(InputState):
+        info: Annotated[list[str], add]
+
+    def valid(state: State) -> Any: ...
+
+    new_builder = StateGraph(State, input=InputState)
+    new_builder.add_node("valid", valid)
+    new_builder.set_entry_point("valid")
+    new_graph = new_builder.compile()
+
+    new_graph.invoke({"something": 1})
+    new_graph.invoke({"something": 2, "info": ["hello", "world"]})  # type: ignore[arg-type]