From 37c215b5a2deb16ae042eba0380407bf1a8d0dd5 Mon Sep 17 00:00:00 2001 From: Sydney Runkle <54324534+sydney-runkle@users.noreply.github.com> Date: Tue, 3 Jun 2025 15:21:51 -0400 Subject: [PATCH] Improve type checking on graph `init` and `invoke`/`stream` (#4932) --- .gitignore | 1 + libs/langgraph/langgraph/graph/state.py | 160 ++++++++++++++++---- libs/langgraph/langgraph/pregel/__init__.py | 16 +- libs/langgraph/langgraph/pregel/protocol.py | 22 ++- libs/langgraph/langgraph/typing.py | 70 +++++++++ libs/langgraph/tests/test_state.py | 10 +- libs/langgraph/tests/test_type_checking.py | 105 +++++++++++++ 7 files changed, 325 insertions(+), 59 deletions(-) create mode 100644 libs/langgraph/langgraph/typing.py create mode 100644 libs/langgraph/tests/test_type_checking.py diff --git a/.gitignore b/.gitignore index 24170c7a9..48916100a 100644 --- a/.gitignore +++ b/.gitignore @@ -181,3 +181,4 @@ Chinook.db .vercel .turbo .editorconfig +.scratch diff --git a/libs/langgraph/langgraph/graph/state.py b/libs/langgraph/langgraph/graph/state.py index bbfdd439e..741f05f31 100644 --- a/libs/langgraph/langgraph/graph/state.py +++ b/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. @@ -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 diff --git a/libs/langgraph/langgraph/pregel/__init__.py b/libs/langgraph/langgraph/pregel/__init__.py index 54dcd24af..82c72baad 100644 --- a/libs/langgraph/langgraph/pregel/__init__.py +++ b/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", diff --git a/libs/langgraph/langgraph/pregel/protocol.py b/libs/langgraph/langgraph/pregel/protocol.py index 85bd724ea..391bf2ba5 100644 --- a/libs/langgraph/langgraph/pregel/protocol.py +++ b/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, diff --git a/libs/langgraph/langgraph/typing.py b/libs/langgraph/langgraph/typing.py new file mode 100644 index 000000000..7779f1c00 --- /dev/null +++ b/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.""" diff --git a/libs/langgraph/tests/test_state.py b/libs/langgraph/tests/test_state.py index 84616c4f1..e02ddc169 100644 --- a/libs/langgraph/tests/test_state.py +++ b/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) == "" diff --git a/libs/langgraph/tests/test_type_checking.py b/libs/langgraph/tests/test_type_checking.py new file mode 100644 index 000000000..146c022e4 --- /dev/null +++ b/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]