graph: improve generics on StateGraph etc + move typing utils to private file (#4982)

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
 

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)

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,

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)

libs/langgraph/langgraph/_typing.py

@@ -0,0 +1,54 @@
+"""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 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."""

libs/langgraph/langgraph/func/__init__.py

@@ -19,6 +19,7 @@ from typing import (
 
 from typing_extensions import Unpack
 
+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
@@ -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
 
 
@@ -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,

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 UNSET, 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, OutputT, StateT, StateT_contra
 from langgraph.utils.fields import (
     get_cached_annotated_keys,
     get_field_default,
@@ -174,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
@@ -181,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<State>.
 
@@ -239,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]]:
@@ -313,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,
@@ -332,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,
@@ -348,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,
@@ -363,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)
@@ -407,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,
@@ -415,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):
@@ -460,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(
@@ -470,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)
@@ -498,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,
@@ -749,32 +783,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 +793,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,
@@ -821,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)
             ]
         )
@@ -837,23 +845,22 @@ 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, 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",
@@ -886,14 +893,16 @@ class StateGraph(Generic[StateT, InputT]):
         return compiled.validate()
 
 
-class CompiledStateGraph(Pregel[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:
@@ -903,7 +912,7 @@ class CompiledStateGraph(Pregel[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"),
@@ -913,7 +922,7 @@ class CompiledStateGraph(Pregel[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"),
@@ -923,7 +932,7 @@ class CompiledStateGraph(Pregel[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:
@@ -991,7 +1000,7 @@ class CompiledStateGraph(Pregel[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:
@@ -1077,7 +1086,7 @@ class CompiledStateGraph(Pregel[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

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, OutputT, 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, OutputT], Generic[StateT, InputT, OutputT]):
     """Pregel manages the runtime behavior for LangGraph applications.
 
     ## Overview

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, OutputT, StateT
 
 
 # TODO: remove Runnable inheritance here!
-class PregelProtocol(Runnable[InputT, Any], Generic[InputT], ABC):
+class PregelProtocol(Runnable[InputT, Any], Generic[StateT, InputT, OutputT], ABC):
     @abstractmethod
     def with_config(
         self, config: RunnableConfig | None = None, **kwargs: Any

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."""

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]

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")
 

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)

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)

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()

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()

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: