From 5920d8aa92fb8a76c7629a65acac5480387de0a5 Mon Sep 17 00:00:00 2001 From: Sydney Runkle Date: Fri, 6 Jun 2025 12:58:19 -0400 Subject: [PATCH] using StateT as default for InputT --- libs/langgraph/langgraph/_typing.py | 47 ++++++++++++++ libs/langgraph/langgraph/func/__init__.py | 2 +- libs/langgraph/langgraph/graph/state.py | 36 ++--------- libs/langgraph/langgraph/pregel/__init__.py | 4 +- libs/langgraph/langgraph/pregel/protocol.py | 4 +- libs/langgraph/langgraph/typing.py | 72 ++++----------------- 6 files changed, 71 insertions(+), 94 deletions(-) create mode 100644 libs/langgraph/langgraph/_typing.py diff --git a/libs/langgraph/langgraph/_typing.py b/libs/langgraph/langgraph/_typing.py new file mode 100644 index 000000000..970ba63ba --- /dev/null +++ b/libs/langgraph/langgraph/_typing.py @@ -0,0 +1,47 @@ +"""Private typing utilities for LangGraph.""" + +from __future__ import annotations + +from dataclasses import Field +from typing import Any, ClassVar, Protocol, Union + +from pydantic import BaseModel +from typing_extensions import TypeAlias, TypedDict + + +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 DeprecatedKwargs(TypedDict): + """TypedDict to use for extra keyword arguments, enabling type checking warnings for deprecated arguments.""" diff --git a/libs/langgraph/langgraph/func/__init__.py b/libs/langgraph/langgraph/func/__init__.py index e85ccecf8..6ebd98293 100644 --- a/libs/langgraph/langgraph/func/__init__.py +++ b/libs/langgraph/langgraph/func/__init__.py @@ -19,6 +19,7 @@ from typing import ( from typing_extensions import Unpack +from langgraph._typing import DeprecatedKwargs from langgraph.cache.base import BaseCache from langgraph.channels.ephemeral_value import EphemeralValue from langgraph.channels.last_value import LastValue @@ -37,7 +38,6 @@ from langgraph.pregel.read import PregelNode from langgraph.pregel.write import ChannelWrite, ChannelWriteEntry from langgraph.store.base import BaseStore from langgraph.types import _DC_KWARGS, CachePolicy, RetryPolicy, StreamMode -from langgraph.typing import DeprecatedKwargs from langgraph.warnings import LangGraphDeprecatedSinceV10 diff --git a/libs/langgraph/langgraph/graph/state.py b/libs/langgraph/langgraph/graph/state.py index b6b73310e..2eeea574f 100644 --- a/libs/langgraph/langgraph/graph/state.py +++ b/libs/langgraph/langgraph/graph/state.py @@ -28,6 +28,7 @@ from langchain_core.runnables import Runnable, RunnableConfig from pydantic import BaseModel from typing_extensions import Self, TypeAlias, Unpack +from langgraph._typing import DeprecatedKwargs from langgraph.cache.base import BaseCache from langgraph.channels.base import BaseChannel from langgraph.channels.binop import BinaryOperatorAggregate @@ -77,7 +78,7 @@ from langgraph.types import ( Send, StreamWriter, ) -from langgraph.typing import DeprecatedKwargs, InputT, StateT, StateT_contra, Unset +from langgraph.typing import InputT, StateT, StateT_contra from langgraph.utils.fields import ( get_cached_annotated_keys, get_field_default, @@ -749,32 +750,6 @@ class StateGraph(Generic[StateT, InputT]): self.compiled = True return self - @overload - def compile( - self: StateGraph[StateT, Unset], - checkpointer: Checkpointer = None, - *, - cache: BaseCache | None = None, - store: BaseStore | None = None, - interrupt_before: All | list[str] | None = None, - interrupt_after: All | list[str] | None = None, - debug: bool = False, - name: str | None = None, - ) -> CompiledStateGraph[StateT, StateT]: ... - - @overload - def compile( - self: StateGraph[StateT, InputT], - checkpointer: Checkpointer = None, - *, - cache: BaseCache | None = None, - store: BaseStore | None = None, - interrupt_before: All | list[str] | None = None, - interrupt_after: All | list[str] | None = None, - debug: bool = False, - name: str | None = None, - ) -> CompiledStateGraph[StateT, InputT]: ... - def compile( self, checkpointer: Checkpointer = None, @@ -785,7 +760,7 @@ class StateGraph(Generic[StateT, InputT]): interrupt_after: All | list[str] | None = None, debug: bool = False, name: str | None = None, - ) -> CompiledStateGraph[StateT, StateT] | CompiledStateGraph[StateT, InputT]: + ) -> CompiledStateGraph[StateT, InputT]: """Compiles the state graph into a `CompiledStateGraph` object. The compiled graph implements the `Runnable` interface and can be invoked, @@ -837,8 +812,7 @@ class StateGraph(Generic[StateT, InputT]): ] ) - ResolvedInputT: type[InputT] | type[StateT] = self.input or self.schema - compiled = CompiledStateGraph[StateT, ResolvedInputT]( # type: ignore[valid-type] + compiled = CompiledStateGraph[StateT, InputT]( builder=self, schema_to_mapper={}, config_type=self.config_schema, @@ -886,7 +860,7 @@ class StateGraph(Generic[StateT, InputT]): return compiled.validate() -class CompiledStateGraph(Pregel[InputT], Generic[StateT, InputT]): +class CompiledStateGraph(Pregel[StateT, InputT], Generic[StateT, InputT]): builder: StateGraph[StateT, InputT] schema_to_mapper: dict[type[Any], Callable[[Any], Any] | None] diff --git a/libs/langgraph/langgraph/pregel/__init__.py b/libs/langgraph/langgraph/pregel/__init__.py index 8fc14a0d1..aa54c252e 100644 --- a/libs/langgraph/langgraph/pregel/__init__.py +++ b/libs/langgraph/langgraph/pregel/__init__.py @@ -107,7 +107,7 @@ from langgraph.types import ( StreamChunk, StreamMode, ) -from langgraph.typing import InputT +from langgraph.typing import InputT, StateT from langgraph.utils.config import ( ensure_config, merge_configs, @@ -298,7 +298,7 @@ class NodeBuilder: ) -class Pregel(PregelProtocol[InputT], Generic[InputT]): +class Pregel(PregelProtocol[StateT, InputT], Generic[StateT, InputT]): """Pregel manages the runtime behavior for LangGraph applications. ## Overview diff --git a/libs/langgraph/langgraph/pregel/protocol.py b/libs/langgraph/langgraph/pregel/protocol.py index 9441f4514..ca4db0bfa 100644 --- a/libs/langgraph/langgraph/pregel/protocol.py +++ b/libs/langgraph/langgraph/pregel/protocol.py @@ -9,11 +9,11 @@ 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 +from langgraph.typing import InputT, StateT # TODO: remove Runnable inheritance here! -class PregelProtocol(Runnable[InputT, Any], Generic[InputT], ABC): +class PregelProtocol(Runnable[InputT, Any], Generic[StateT, InputT], ABC): @abstractmethod def with_config( self, config: RunnableConfig | None = None, **kwargs: Any diff --git a/libs/langgraph/langgraph/typing.py b/libs/langgraph/langgraph/typing.py index 7bac53a59..01ea27adf 100644 --- a/libs/langgraph/langgraph/typing.py +++ b/libs/langgraph/langgraph/typing.py @@ -1,57 +1,10 @@ from __future__ import annotations -from dataclasses import Field -from typing import ( - Any, - ClassVar, - Protocol, - Union, -) +from typing import Union -from pydantic import BaseModel -from typing_extensions import TypeAlias, TypedDict, TypeVar +from typing_extensions import 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() +from langgraph._typing import StateLike StateT = TypeVar("StateT", bound=StateLike) """Type variable used to represent the state in a graph.""" @@ -60,15 +13,18 @@ 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. +InputT = TypeVar("InputT", bound=StateLike, default=StateT) +"""Type variable used to represent the input to a state 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`.""" +Defaults to `StateT`. +""" -OutputT = TypeVar("OutputT", bound=Union[StateLike, Unset], default=Unset) -"""Type variable used to represent the output of a graph.""" +ResolvedInputT = TypeVar("ResolvedInputT", bound=StateLike) +"""Type variable used to represent the resolved input to a state graph. + +No default. +""" -class DeprecatedKwargs(TypedDict): - """TypedDict to use for extra keyword arguments, enabling type checking warnings for deprecated arguments.""" +OutputT = TypeVar("OutputT", bound=Union[StateLike, None], default=StateT) +"""Type variable used to represent the output of a state graph."""