From 70c5a22e4b034cf86260481b2f08f1affc92dfe8 Mon Sep 17 00:00:00 2001
From: William Fu-Hinthorn <13333726+hinthornw@users.noreply.github.com>
Date: Mon, 20 May 2024 09:26:56 -0700
Subject: [PATCH] Warn invalid state

---
 langgraph/graph/state.py | 15 +++++++++++++
 tests/test_state.py      | 48 ++++++++++++++++++++++++++++++++++++++++
 2 files changed, 63 insertions(+)
 create mode 100644 tests/test_state.py

diff --git a/langgraph/graph/state.py b/langgraph/graph/state.py
index 99b6c0018..f71f1c042 100644
--- a/langgraph/graph/state.py
+++ b/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)
diff --git a/tests/test_state.py b/tests/test_state.py
new file mode 100644
index 000000000..ecf35eb2a
--- /dev/null
+++ b/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)