Warn invalid state

langgraph/graph/state.py

@@ -1,4 +1,6 @@
 import logging
+import typing
+import warnings
 from functools import partial
 from inspect import signature
 from typing import (
@@ -35,6 +37,18 @@ from langgraph.utils import RunnableCallable
 logger = logging.getLogger(__name__)
 
 
+def _warn_invalid_state_schema(schema: Union[Type[Any], Any]) -> None:
+    if isinstance(schema, type):
+        return
+    if typing.get_args(schema):
+        return
+    warnings.warn(
+        f"Invalid state_schema: {schema}. Expected a type or Annotated[type, reducer]. "
+        "Please provide a valid schema to ensure correct updates.\n"
+        " See: https://langchain-ai.github.io/langgraph/reference/graphs/#stategraph"
+    )
+
+
 class StateGraph(Graph):
     """A graph whose nodes communicate by reading and writing to a shared state.
     The signature of each node is State -> Partial<State>.
@@ -90,6 +104,7 @@ class StateGraph(Graph):
         self, state_schema: Type[Any], config_schema: Optional[Type[Any]] = None
     ) -> None:
         super().__init__()
+        _warn_invalid_state_schema(state_schema)
         self.schema = state_schema
         self.config_schema = config_schema
         self.channels, self.managed = _get_channels(state_schema)

tests/test_state.py

@@ -0,0 +1,48 @@
+from typing import Annotated as Annotated2
+from typing import Any
+
+import pytest
+from pydantic.v1 import BaseModel
+from typing_extensions import Annotated, TypedDict
+
+from langgraph.graph.state import _warn_invalid_state_schema
+
+
+class State(BaseModel):
+    foo: str
+    bar: int
+
+
+class State2(TypedDict):
+    foo: str
+    bar: int
+
+
+@pytest.mark.parametrize(
+    "schema",
+    [
+        {"foo": "bar"},
+        ["hi", lambda x, y: x + y],
+        State(foo="bar", bar=1),
+        State2(foo="bar", bar=1),
+    ],
+)
+def test_warns_invalid_schema(schema: Any):
+    with pytest.warns(UserWarning):
+        _warn_invalid_state_schema(schema)
+
+
+@pytest.mark.parametrize(
+    "schema",
+    [
+        Annotated[dict, lambda x, y: y],
+        Annotated2[list, lambda x, y: y],
+        dict,
+        State,
+        State2,
+    ],
+)
+def test_doesnt_warn_valid_schema(schema: Any):
+    # Assert the function does not raise a warning
+    with pytest.warns(None):
+        _warn_invalid_state_schema(schema)