From 5920d8aa92fb8a76c7629a65acac5480387de0a5 Mon Sep 17 00:00:00 2001 From: Sydney Runkle Date: Fri, 6 Jun 2025 12:58:19 -0400 Subject: [PATCH 1/2] 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.""" From b7354521537175aed60c6b8bacac22049ada8ca6 Mon Sep 17 00:00:00 2001 From: Sydney Runkle <54324534+sydney-runkle@users.noreply.github.com> Date: Fri, 6 Jun 2025 19:44:56 -0400 Subject: [PATCH 2/2] deprecate `input` and `output` in favor of `input_schema` and `output_schema` (#4983) --- docs/docs/concepts/low_level.md | 4 +- docs/docs/concepts/server-mcp.md | 2 +- docs/docs/how-tos/graph-api.ipynb | 6 +- libs/langgraph/bench/fanout_to_subgraph.py | 4 +- libs/langgraph/langgraph/_typing.py | 7 + libs/langgraph/langgraph/func/__init__.py | 6 +- libs/langgraph/langgraph/graph/state.py | 121 +++++++++++------- libs/langgraph/langgraph/pregel/__init__.py | 4 +- libs/langgraph/langgraph/pregel/protocol.py | 4 +- libs/langgraph/tests/test_deprecation.py | 26 ++++ libs/langgraph/tests/test_large_cases.py | 2 +- libs/langgraph/tests/test_pregel.py | 30 +++-- libs/langgraph/tests/test_pregel_async.py | 2 +- libs/langgraph/tests/test_state.py | 2 +- libs/langgraph/tests/test_type_checking.py | 2 +- .../langgraph/prebuilt/chat_agent_executor.py | 28 ++-- 16 files changed, 164 insertions(+), 86 deletions(-) diff --git a/docs/docs/concepts/low_level.md b/docs/docs/concepts/low_level.md index 945a9cc41..f277b742e 100644 --- a/docs/docs/concepts/low_level.md +++ b/docs/docs/concepts/low_level.md @@ -89,7 +89,7 @@ def node_3(state: PrivateState) -> OutputState: # Read from PrivateState, write to OutputState return {"graph_output": state["bar"] + " Lance"} -builder = StateGraph(OverallState,input=InputState,output=OutputState) +builder = StateGraph(OverallState,input_schema=InputState,output_schema=OutputState) builder.add_node("node_1", node_1) builder.add_node("node_2", node_2) builder.add_node("node_3", node_3) @@ -107,7 +107,7 @@ There are two subtle and important points to note here: 1. We pass `state: InputState` as the input schema to `node_1`. But, we write out to `foo`, a channel in `OverallState`. How can we write out to a state channel that is not included in the input schema? This is because a node _can write to any state channel in the graph state._ The graph state is the union of the state channels defined at initialization, which includes `OverallState` and the filters `InputState` and `OutputState`. -2. We initialize the graph with `StateGraph(OverallState,input=InputState,output=OutputState)`. So, how can we write to `PrivateState` in `node_2`? How does the graph gain access to this schema if it was not passed in the `StateGraph` initialization? We can do this because _nodes can also declare additional state channels_ as long as the state schema definition exists. In this case, the `PrivateState` schema is defined, so we can add `bar` as a new state channel in the graph and write to it. +2. We initialize the graph with `StateGraph(OverallState,input_schema=InputState,output_schema=OutputState)`. So, how can we write to `PrivateState` in `node_2`? How does the graph gain access to this schema if it was not passed in the `StateGraph` initialization? We can do this because _nodes can also declare additional state channels_ as long as the state schema definition exists. In this case, the `PrivateState` schema is defined, so we can add `bar` as a new state channel in the graph and write to it. ### Reducers diff --git a/docs/docs/concepts/server-mcp.md b/docs/docs/concepts/server-mcp.md index 7d4280b59..658e69399 100644 --- a/docs/docs/concepts/server-mcp.md +++ b/docs/docs/concepts/server-mcp.md @@ -94,7 +94,7 @@ def answer_node(state: InputState): return {"answer": "bye", "question": state["question"]} # Build the graph with explicit schemas -builder = StateGraph(OverallState, input=InputState, output=OutputState) +builder = StateGraph(OverallState, input_schema=InputState, output_schema=OutputState) builder.add_node(answer_node) builder.add_edge(START, "answer_node") builder.add_edge("answer_node", END) diff --git a/docs/docs/how-tos/graph-api.ipynb b/docs/docs/how-tos/graph-api.ipynb index 9adfa27ac..9d890beae 100644 --- a/docs/docs/how-tos/graph-api.ipynb +++ b/docs/docs/how-tos/graph-api.ipynb @@ -439,7 +439,7 @@ }, { "cell_type": "code", - "execution_count": 6, + "execution_count": null, "id": "6ec0eb77-874e-443e-8c73-93125b515106", "metadata": {}, "outputs": [ @@ -478,7 +478,7 @@ "\n", "\n", "# Build the graph with input and output schemas specified\n", - "builder = StateGraph(OverallState, input=InputState, output=OutputState)\n", + "builder = StateGraph(OverallState, input_schema=InputState, output_schema=OutputState)\n", "builder.add_node(answer_node) # Add the answer node\n", "builder.add_edge(START, \"answer_node\") # Define the starting edge\n", "builder.add_edge(\"answer_node\", END) # Define the ending edge\n", @@ -3430,7 +3430,7 @@ "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", - "version": "3.12.9" + "version": "3.9.6" } }, "nbformat": 4, diff --git a/libs/langgraph/bench/fanout_to_subgraph.py b/libs/langgraph/bench/fanout_to_subgraph.py index 107366ca3..ade894acc 100644 --- a/libs/langgraph/bench/fanout_to_subgraph.py +++ b/libs/langgraph/bench/fanout_to_subgraph.py @@ -37,7 +37,7 @@ def fanout_to_subgraph() -> StateGraph: return END if state["jokes"][0].endswith(" a" * 10) else "bump" # subgraph - subgraph = StateGraph(JokeState, input=JokeInput, output=JokeOutput) + subgraph = StateGraph(JokeState, input_schema=JokeInput, output_schema=JokeOutput) subgraph.add_node("edit", edit) subgraph.add_node("generate", generate) subgraph.add_node("bump", bump) @@ -87,7 +87,7 @@ def fanout_to_subgraph_sync() -> StateGraph: return END if state["jokes"][0].endswith(" a" * 10) else "bump" # subgraph - subgraph = StateGraph(JokeState, input=JokeInput, output=JokeOutput) + subgraph = StateGraph(JokeState, input_schema=JokeInput, output_schema=JokeOutput) subgraph.add_node("edit", edit) subgraph.add_node("generate", generate) subgraph.add_node("bump", bump) diff --git a/libs/langgraph/langgraph/_typing.py b/libs/langgraph/langgraph/_typing.py index 970ba63ba..79b5478d0 100644 --- a/libs/langgraph/langgraph/_typing.py +++ b/libs/langgraph/langgraph/_typing.py @@ -43,5 +43,12 @@ Note: we cannot use either `TypedDict` or `dataclass` directly due to limitation """ +class Unset: + """A sentinel value to represent an unset type.""" + + +UNSET: Unset = Unset() + + 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 6ebd98293..4c7595c7e 100644 --- a/libs/langgraph/langgraph/func/__init__.py +++ b/libs/langgraph/langgraph/func/__init__.py @@ -19,7 +19,7 @@ from typing import ( from typing_extensions import Unpack -from langgraph._typing import DeprecatedKwargs +from langgraph._typing import UNSET, DeprecatedKwargs from langgraph.cache.base import BaseCache from langgraph.channels.ephemeral_value import EphemeralValue from langgraph.channels.last_value import LastValue @@ -176,7 +176,7 @@ def task( await add_one.ainvoke([1, 2, 3]) # Returns [2, 3, 4] ``` """ - if (retry := kwargs.get("retry")) is not None: + if (retry := kwargs.get("retry", UNSET)) is not UNSET: warnings.warn( "`retry` is deprecated and will be removed. Please use `retry_policy` instead.", category=LangGraphDeprecatedSinceV10, @@ -380,7 +380,7 @@ class entrypoint: **kwargs: Unpack[DeprecatedKwargs], ) -> None: """Initialize the entrypoint decorator.""" - if (retry := kwargs.get("retry")) is not None: + if (retry := kwargs.get("retry", UNSET)) is not UNSET: warnings.warn( "`retry` is deprecated and will be removed. Please use `retry_policy` instead.", category=LangGraphDeprecatedSinceV10, diff --git a/libs/langgraph/langgraph/graph/state.py b/libs/langgraph/langgraph/graph/state.py index 2eeea574f..b4210a351 100644 --- a/libs/langgraph/langgraph/graph/state.py +++ b/libs/langgraph/langgraph/graph/state.py @@ -28,7 +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._typing import UNSET, DeprecatedKwargs from langgraph.cache.base import BaseCache from langgraph.channels.base import BaseChannel from langgraph.channels.binop import BinaryOperatorAggregate @@ -78,7 +78,7 @@ from langgraph.types import ( Send, StreamWriter, ) -from langgraph.typing import InputT, StateT, StateT_contra +from langgraph.typing import InputT, OutputT, StateT, StateT_contra from langgraph.utils.fields import ( get_cached_annotated_keys, get_field_default, @@ -175,6 +175,8 @@ class StateNodeSpec(NamedTuple): # a generic StateNode, so maybe a dataclass runnable: StateNode metadata: dict[str, Any] | None + # TODO: rename to input_schema, though we really just want to modify this structure to + # be a dataclass input: type[Any] retry_policy: RetryPolicy | Sequence[RetryPolicy] | None cache_policy: CachePolicy | None @@ -182,7 +184,7 @@ class StateNodeSpec(NamedTuple): defer: bool = False -class StateGraph(Generic[StateT, InputT]): +class StateGraph(Generic[StateT, InputT, OutputT]): """A graph whose nodes communicate by reading and writing to a shared state. The signature of each node is State -> Partial. @@ -240,34 +242,57 @@ class StateGraph(Generic[StateT, InputT]): channels: dict[str, BaseChannel] managed: dict[str, ManagedValueSpec] schemas: dict[type[Any], dict[str, BaseChannel | ManagedValueSpec]] + waiting_edges: set[tuple[tuple[str, ...], str]] + + compiled: bool + state_schema: type[StateT] + input_schema: type[InputT] + output_schema: type[OutputT] def __init__( self, state_schema: type[StateT], config_schema: type[Any] | None = None, *, - input: type[InputT] | None = None, - output: type[Any] | None = None, + input_schema: type[InputT] | None = None, + output_schema: type[OutputT] | None = None, + **kwargs: Unpack[DeprecatedKwargs], ) -> None: - input = input or state_schema - output = output or state_schema + if (input_ := kwargs.get("input", UNSET)) is not UNSET: + warnings.warn( + "`input` is deprecated and will be removed. Please use `input_schema` instead.", + category=LangGraphDeprecatedSinceV10, + stacklevel=2, + ) + if input_schema is None: + input_schema = cast(Union[type[InputT], None], input_) + + if (output := kwargs.get("output", UNSET)) is not UNSET: + warnings.warn( + "`output` is deprecated and will be removed. Please use `output_schema` instead.", + category=LangGraphDeprecatedSinceV10, + stacklevel=2, + ) + if output_schema is None: + output_schema = cast(Union[type[OutputT], None], output) self.nodes = {} - self.edges = set[tuple[str, str]]() + self.edges = set() self.branches = defaultdict(dict) - self.support_multiple_edges = False - self.compiled = False self.schemas = {} self.channels = {} self.managed = {} - self.schema = state_schema - self.input = input - self.output = output - self._add_schema(state_schema) - self._add_schema(input, allow_managed=False) - self._add_schema(output, allow_managed=False) + self.compiled = False + self.waiting_edges = set() + + self.state_schema = state_schema + self.input_schema = cast(type[InputT], input_schema or state_schema) + self.output_schema = cast(type[OutputT], output_schema or state_schema) self.config_schema = config_schema - self.waiting_edges: set[tuple[tuple[str, ...], str]] = set() + + self._add_schema(self.state_schema) + self._add_schema(self.input_schema, allow_managed=False) + self._add_schema(self.output_schema, allow_managed=False) @property def _all_edges(self) -> set[tuple[str, str]]: @@ -314,7 +339,7 @@ class StateGraph(Generic[StateT, InputT]): *, defer: bool = False, metadata: dict[str, Any] | None = None, - input: type[Any] | None = None, + input_schema: type[Any] | None = None, retry_policy: RetryPolicy | Sequence[RetryPolicy] | None = None, cache_policy: CachePolicy | None = None, destinations: dict[str, str] | tuple[str, ...] | None = None, @@ -333,7 +358,7 @@ class StateGraph(Generic[StateT, InputT]): *, defer: bool = False, metadata: dict[str, Any] | None = None, - input: type[Any] | None = None, + input_schema: type[Any] | None = None, retry_policy: RetryPolicy | Sequence[RetryPolicy] | None = None, cache_policy: CachePolicy | None = None, destinations: dict[str, str] | tuple[str, ...] | None = None, @@ -349,7 +374,7 @@ class StateGraph(Generic[StateT, InputT]): *, defer: bool = False, metadata: dict[str, Any] | None = None, - input: type[Any] | None = None, + input_schema: type[Any] | None = None, retry_policy: RetryPolicy | Sequence[RetryPolicy] | None = None, cache_policy: CachePolicy | None = None, destinations: dict[str, str] | tuple[str, ...] | None = None, @@ -364,7 +389,7 @@ class StateGraph(Generic[StateT, InputT]): Will be used as the node function or runnable if `node` is a string (node name). defer: Whether to defer the execution of the node until the run is about to end. metadata: The metadata associated with the node. (default: None) - input: The input schema for the node. (default: the graph's input schema) + input_schema: The input schema for the node. (default: the graph's state schema) retry_policy: The retry policy for the node. (default: None) If a sequence is provided, the first matching policy will be applied. cache_policy: The cache policy for the node. (default: None) @@ -408,7 +433,7 @@ class StateGraph(Generic[StateT, InputT]): Returns: Self: The instance of the state graph, allowing for method chaining. """ - if (retry := kwargs.get("retry")) is not None: + if (retry := kwargs.get("retry", UNSET)) is not UNSET: warnings.warn( "`retry` is deprecated and will be removed. Please use `retry_policy` instead.", category=LangGraphDeprecatedSinceV10, @@ -416,6 +441,14 @@ class StateGraph(Generic[StateT, InputT]): if retry_policy is None: retry_policy = retry # type: ignore[assignment] + if (input_ := kwargs.get("input", UNSET)) is not UNSET: + warnings.warn( + "`input` is deprecated and will be removed. Please use `input_schema` instead.", + category=LangGraphDeprecatedSinceV10, + ) + if input_schema is None: + input_schema = cast(Union[type[InputT], None], input_) + if not isinstance(node, str): action = node if isinstance(action, Runnable): @@ -461,7 +494,7 @@ class StateGraph(Generic[StateT, InputT]): hints := get_type_hints(getattr(action, "__call__")) or get_type_hints(action) ): - if input is None: + if input_schema is None: first_parameter_name = next( iter( inspect.signature( @@ -471,7 +504,7 @@ class StateGraph(Generic[StateT, InputT]): ) if input_hint := hints.get(first_parameter_name): if isinstance(input_hint, type) and get_type_hints(input_hint): - input = input_hint + input_schema = input_hint if rtn := hints.get("return"): # Handle Union types rtn_origin = get_origin(rtn) @@ -499,12 +532,12 @@ class StateGraph(Generic[StateT, InputT]): if destinations is not None: ends = destinations - if input is not None: - self._add_schema(input) + if input_schema is not None: + self._add_schema(input_schema) self.nodes[node] = StateNodeSpec( coerce_to_runnable(action, name=node, trace=False), # type: ignore metadata, - input=input or self.schema, + input=input_schema or self.state_schema, retry_policy=retry_policy, cache_policy=cache_policy, ends=ends, @@ -796,11 +829,11 @@ class StateGraph(Generic[StateT, InputT]): # prepare output channels output_channels = ( "__root__" - if len(self.schemas[self.output]) == 1 - and "__root__" in self.schemas[self.output] + if len(self.schemas[self.output_schema]) == 1 + and "__root__" in self.schemas[self.output_schema] else [ key - for key, val in self.schemas[self.output].items() + for key, val in self.schemas[self.output_schema].items() if not is_managed_value(val) ] ) @@ -812,22 +845,22 @@ class StateGraph(Generic[StateT, InputT]): ] ) - compiled = CompiledStateGraph[StateT, InputT]( + compiled = CompiledStateGraph[StateT, InputT, OutputT]( builder=self, schema_to_mapper={}, config_type=self.config_schema, input_model=( - self.input + self.input_schema if len(self.channels) > 1 - and isclass(self.input) - and issubclass(self.input, BaseModel) + and isclass(self.input_schema) + and issubclass(self.input_schema, BaseModel) else None ), nodes={}, channels={ **self.channels, **self.managed, - START: EphemeralValue(self.input), + START: EphemeralValue(self.input_schema), }, input_channels=START, stream_mode="updates", @@ -860,14 +893,16 @@ class StateGraph(Generic[StateT, InputT]): return compiled.validate() -class CompiledStateGraph(Pregel[StateT, InputT], Generic[StateT, InputT]): - builder: StateGraph[StateT, InputT] +class CompiledStateGraph( + Pregel[StateT, InputT, OutputT], Generic[StateT, InputT, OutputT] +): + builder: StateGraph[StateT, InputT, OutputT] schema_to_mapper: dict[type[Any], Callable[[Any], Any] | None] def __init__( self, *, - builder: StateGraph[StateT, InputT], + builder: StateGraph[StateT, InputT, OutputT], schema_to_mapper: dict[type[Any], Callable[[Any], Any] | None], **kwargs: Any, ) -> None: @@ -877,7 +912,7 @@ class CompiledStateGraph(Pregel[StateT, InputT], Generic[StateT, InputT]): def get_input_schema(self, config: RunnableConfig | None = None) -> type[BaseModel]: return _get_schema( - typ=self.builder.input, + typ=self.builder.input_schema, schemas=self.builder.schemas, channels=self.builder.channels, name=self.get_name("Input"), @@ -887,7 +922,7 @@ class CompiledStateGraph(Pregel[StateT, InputT], Generic[StateT, InputT]): self, config: RunnableConfig | None = None ) -> type[BaseModel]: return _get_schema( - typ=self.builder.output, + typ=self.builder.output_schema, schemas=self.builder.schemas, channels=self.builder.channels, name=self.get_name("Output"), @@ -897,7 +932,7 @@ class CompiledStateGraph(Pregel[StateT, InputT], Generic[StateT, InputT]): if key == START: output_keys = [ k - for k, v in self.builder.schemas[self.builder.input].items() + for k, v in self.builder.schemas[self.builder.input_schema].items() if not is_managed_value(v) ] else: @@ -965,7 +1000,7 @@ class CompiledStateGraph(Pregel[StateT, InputT], Generic[StateT, InputT]): writers=[ChannelWrite(write_entries)], ) elif node is not None: - input_schema = node.input if node else self.builder.schema + input_schema = node.input if node else self.builder._state_schema input_values = {k: k for k in self.builder.schemas[input_schema]} is_single_input = len(input_values) == 1 and "__root__" in input_values if input_schema in self.schema_to_mapper: @@ -1051,7 +1086,7 @@ class CompiledStateGraph(Pregel[StateT, InputT], Generic[StateT, InputT]): schema = branch.input_schema or ( self.builder.nodes[start].input if start in self.builder.nodes - else self.builder.schema + else self.builder.state_schema ) channels = list(self.builder.schemas[schema]) # get mapper diff --git a/libs/langgraph/langgraph/pregel/__init__.py b/libs/langgraph/langgraph/pregel/__init__.py index aa54c252e..5519ea237 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, StateT +from langgraph.typing import InputT, OutputT, StateT from langgraph.utils.config import ( ensure_config, merge_configs, @@ -298,7 +298,7 @@ class NodeBuilder: ) -class Pregel(PregelProtocol[StateT, InputT], Generic[StateT, InputT]): +class Pregel(PregelProtocol[StateT, InputT, OutputT], Generic[StateT, InputT, OutputT]): """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 ca4db0bfa..e55be3cbc 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, StateT +from langgraph.typing import InputT, OutputT, StateT # TODO: remove Runnable inheritance here! -class PregelProtocol(Runnable[InputT, Any], Generic[StateT, InputT], ABC): +class PregelProtocol(Runnable[InputT, Any], Generic[StateT, InputT, OutputT], ABC): @abstractmethod def with_config( self, config: RunnableConfig | None = None, **kwargs: Any diff --git a/libs/langgraph/tests/test_deprecation.py b/libs/langgraph/tests/test_deprecation.py index 7300b3026..456961390 100644 --- a/libs/langgraph/tests/test_deprecation.py +++ b/libs/langgraph/tests/test_deprecation.py @@ -40,3 +40,29 @@ def test_entrypoint_retry_arg() -> None: @entrypoint(retry=RetryPolicy()) # type: ignore[arg-type] def my_entrypoint(state: PlainState) -> PlainState: return state + + +def test_state_graph_input_schema() -> None: + with pytest.warns( + LangGraphDeprecatedSinceV10, + match="`input` is deprecated and will be removed. Please use `input_schema` instead.", + ): + StateGraph(PlainState, input=PlainState) # type: ignore[arg-type] + + +def test_state_graph_output_schema() -> None: + with pytest.warns( + LangGraphDeprecatedSinceV10, + match="`output` is deprecated and will be removed. Please use `output_schema` instead.", + ): + StateGraph(PlainState, output=PlainState) # type: ignore[arg-type] + + +def test_add_node_input_schema() -> None: + builder = StateGraph(PlainState) + + with pytest.warns( + LangGraphDeprecatedSinceV10, + match="`input` is deprecated and will be removed. Please use `input_schema` instead.", + ): + builder.add_node("test_node", lambda state: state, input=PlainState) # type: ignore[arg-type] diff --git a/libs/langgraph/tests/test_large_cases.py b/libs/langgraph/tests/test_large_cases.py index 647e72bcb..c19202bf8 100644 --- a/libs/langgraph/tests/test_large_cases.py +++ b/libs/langgraph/tests/test_large_cases.py @@ -570,7 +570,7 @@ def test_conditional_state_graph( workflow = StateGraph(AgentState) workflow.add_node("agent", agent) - workflow.add_node("tools", execute_tools, input=ToolState) + workflow.add_node("tools", execute_tools, input_schema=ToolState) workflow.set_entry_point("agent") diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index 8c75d80dc..f8e99fa0e 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -288,7 +288,7 @@ def test_node_schemas_custom_output() -> None: "now": 123, } - builder = StateGraph(State, output=Output) + builder = StateGraph(State, output_schema=Output) builder.add_node("a", node_a) builder.add_node("b", node_b) builder.add_node("c", node_c) @@ -301,7 +301,7 @@ def test_node_schemas_custom_output() -> None: "messages": [_AnyIdHumanMessage(content="hello")], } - builder = StateGraph(State, output=Output) + builder = StateGraph(State, output_schema=Output) builder.add_node("a", node_a) builder.add_node("b", node_b) builder.add_node("c", node_c) @@ -2492,7 +2492,7 @@ def test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2( assert isinstance(data, State) return "retriever_two" - workflow = StateGraph(State, input=Input, output=Output) + workflow = StateGraph(State, input_schema=Input, output_schema=Output) workflow.add_node("rewrite_query", rewrite_query) workflow.add_node("analyzer_one", analyzer_one) @@ -2621,7 +2621,7 @@ def test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic_inp assert isinstance(data, State) return "retriever_two" - workflow = StateGraph(State, input=Input, output=Output) + workflow = StateGraph(State, input_schema=Input, output_schema=Output) workflow.add_node("rewrite_query", rewrite_query) workflow.add_node("analyzer_one", analyzer_one) @@ -6185,14 +6185,17 @@ def test_multiple_subgraphs(sync_checkpointer: BaseCheckpointSaver) -> None: return {"result": state["a"] + state["b"]} add_subgraph = ( - StateGraph(State, output=Output).add_node(add).add_edge(START, "add").compile() + StateGraph(State, output_schema=Output) + .add_node(add) + .add_edge(START, "add") + .compile() ) def multiply(state): return {"result": state["a"] * state["b"]} multiply_subgraph = ( - StateGraph(State, output=Output) + StateGraph(State, output_schema=Output) .add_node(multiply) .add_edge(START, "multiply") .compile() @@ -6205,7 +6208,7 @@ def test_multiple_subgraphs(sync_checkpointer: BaseCheckpointSaver) -> None: return another_result parent_call_same_subgraph = ( - StateGraph(State, output=Output) + StateGraph(State, output_schema=Output) .add_node(call_same_subgraph) .add_edge(START, "call_same_subgraph") .compile(checkpointer=sync_checkpointer) @@ -6227,7 +6230,7 @@ def test_multiple_subgraphs(sync_checkpointer: BaseCheckpointSaver) -> None: } parent_call_multiple_subgraphs = ( - StateGraph(State, output=Output) + StateGraph(State, output_schema=Output) .add_node(call_multiple_subgraphs) .add_edge(START, "call_multiple_subgraphs") .compile(checkpointer=sync_checkpointer) @@ -6301,14 +6304,17 @@ def test_multiple_subgraphs_mixed_entrypoint( return {"result": state["a"] + state["b"]} add_subgraph = ( - StateGraph(State, output=Output).add_node(add).add_edge(START, "add").compile() + StateGraph(State, output_schema=Output) + .add_node(add) + .add_edge(START, "add") + .compile() ) def multiply(state): return {"result": state["a"] * state["b"]} multiply_subgraph = ( - StateGraph(State, output=Output) + StateGraph(State, output_schema=Output) .add_node(multiply) .add_edge(START, "multiply") .compile() @@ -6377,7 +6383,7 @@ def test_multiple_subgraphs_mixed_state_graph( return {"result": another_result} parent_call_same_subgraph = ( - StateGraph(State, output=Output) + StateGraph(State, output_schema=Output) .add_node(call_same_subgraph) .add_edge(START, "call_same_subgraph") .compile(checkpointer=sync_checkpointer) @@ -6399,7 +6405,7 @@ def test_multiple_subgraphs_mixed_state_graph( } parent_call_multiple_subgraphs = ( - StateGraph(State, output=Output) + StateGraph(State, output_schema=Output) .add_node(call_multiple_subgraphs) .add_edge(START, "call_multiple_subgraphs") .compile(checkpointer=sync_checkpointer) diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index d35b5be1d..f3fb807be 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/libs/langgraph/tests/test_pregel_async.py @@ -4310,7 +4310,7 @@ async def test_in_one_fan_out_state_graph_waiting_edge_custom_state_class( assert isinstance(data, State) return "retriever_two" - workflow = StateGraph(State, input=Input, output=Output) + workflow = StateGraph(State, input_schema=Input, output_schema=Output) workflow.add_node("rewrite_query", rewrite_query) workflow.add_node("analyzer_one", analyzer_one) diff --git a/libs/langgraph/tests/test_state.py b/libs/langgraph/tests/test_state.py index e02ddc169..fe1958b08 100644 --- a/libs/langgraph/tests/test_state.py +++ b/libs/langgraph/tests/test_state.py @@ -153,7 +153,7 @@ def test_state_schema_optional_values(total_: bool): class State(InputState): # this would be ignored val4: dict - builder = StateGraph(State, input=InputState, output=OutputState) + builder = StateGraph(State, input_schema=InputState, output_schema=OutputState) builder.add_node("n", lambda x: x) builder.add_edge("__start__", "n") graph = builder.compile() diff --git a/libs/langgraph/tests/test_type_checking.py b/libs/langgraph/tests/test_type_checking.py index 146c022e4..0b7ee0679 100644 --- a/libs/langgraph/tests/test_type_checking.py +++ b/libs/langgraph/tests/test_type_checking.py @@ -96,7 +96,7 @@ def test_input_state_specified() -> None: def valid(state: State) -> Any: ... - new_builder = StateGraph(State, input=InputState) + new_builder = StateGraph(State, input_schema=InputState) new_builder.add_node("valid", valid) new_builder.set_entry_point("valid") new_graph = new_builder.compile() diff --git a/libs/prebuilt/langgraph/prebuilt/chat_agent_executor.py b/libs/prebuilt/langgraph/prebuilt/chat_agent_executor.py index 6ed2b7f2b..9bf35a941 100644 --- a/libs/prebuilt/langgraph/prebuilt/chat_agent_executor.py +++ b/libs/prebuilt/langgraph/prebuilt/chat_agent_executor.py @@ -591,11 +591,11 @@ def create_react_agent( workflow = StateGraph(state_schema, config_schema=config_schema) workflow.add_node( "agent", - RunnableCallable(call_model, acall_model), - input=input_schema, + RunnableCallable(call_model, acall_model), # type: ignore[call-overload] + input_schema=input_schema, ) if pre_model_hook is not None: - workflow.add_node("pre_model_hook", pre_model_hook) + workflow.add_node("pre_model_hook", pre_model_hook) # type: ignore[arg-type] workflow.add_edge("pre_model_hook", "agent") entrypoint = "pre_model_hook" else: @@ -604,14 +604,15 @@ def create_react_agent( workflow.set_entry_point(entrypoint) if post_model_hook is not None: - workflow.add_node("post_model_hook", post_model_hook) + workflow.add_node("post_model_hook", post_model_hook) # type: ignore[arg-type] workflow.add_edge("agent", "post_model_hook") if response_format is not None: workflow.add_node( "generate_structured_response", - RunnableCallable( - generate_structured_response, agenerate_structured_response + RunnableCallable( # type: ignore[call-overload] + generate_structured_response, + agenerate_structured_response, ), ) if post_model_hook is not None: @@ -658,14 +659,16 @@ def create_react_agent( # Define the two nodes we will cycle between workflow.add_node( - "agent", RunnableCallable(call_model, acall_model), input=input_schema + "agent", + RunnableCallable(call_model, acall_model), # type: ignore[call-overload] + input_schema=input_schema, ) - workflow.add_node("tools", tool_node) + workflow.add_node("tools", tool_node) # type: ignore[call-overload] # Optionally add a pre-model hook node that will be called # every time before the "agent" (LLM-calling node) if pre_model_hook is not None: - workflow.add_node("pre_model_hook", pre_model_hook) + workflow.add_node("pre_model_hook", pre_model_hook) # type: ignore[arg-type] workflow.add_edge("pre_model_hook", "agent") entrypoint = "pre_model_hook" else: @@ -680,7 +683,7 @@ def create_react_agent( # Add a post model hook node if post_model_hook is provided if post_model_hook is not None: - workflow.add_node("post_model_hook", post_model_hook) + workflow.add_node("post_model_hook", post_model_hook) # type: ignore[arg-type] agent_paths.append("post_model_hook") workflow.add_edge("agent", "post_model_hook") else: @@ -690,8 +693,9 @@ def create_react_agent( if response_format is not None: workflow.add_node( "generate_structured_response", - RunnableCallable( - generate_structured_response, agenerate_structured_response + RunnableCallable( # type: ignore[call-overload] + generate_structured_response, + agenerate_structured_response, ), ) if post_model_hook is not None: