Add generics tests

Signed-off-by: William Fu-Hinthorn <13333726+hinthornw@users.noreply.github.com>

libs/langgraph/langgraph/graph/schema_utils.py

@@ -70,6 +70,17 @@ class SchemaCoercionMapper:
                 for n, f in schema.__fields__.items()
             }
             self._construct = schema.construct
+            unhandled_attrs = (
+                "__pre_root_validators__",
+                "__post_root_validators__",
+                "__validators__",
+            )
+            if any(getattr(schema, c, None) for c in unhandled_attrs):
+                self.coerce: Callable[[Any, Any], Union[BaseModelV1, BaseModel]] = (
+                    lambda v, _: schema(**v)
+                )
+            else:
+                self.coerce = self._coerce
 
         elif issubclass(schema, BaseModel):
             self._fields = {
@@ -77,6 +88,13 @@ class SchemaCoercionMapper:
                 for n, f in schema.model_fields.items()
             }
             self._construct: Callable[..., Any] = schema.model_construct  # type: ignore
+            unhandled_attrs = ("validators", "field_validators", "root_validators")
+            if (decorators := getattr(schema, "__pydantic_decorators__", None)) and any(
+                getattr(decorators, attr, None) for attr in unhandled_attrs
+            ):
+                self.coerce = lambda v, _: schema.model_validate(v)
+            else:
+                self.coerce = self._coerce
 
         else:
             raise TypeError("Schema is neither a Pydantic v1 nor v2 model.")
@@ -86,7 +104,7 @@ class SchemaCoercionMapper:
     def __call__(self, input_data: Any, depth: Optional[int] = None) -> Any:
         return self.coerce(input_data, depth)
 
-    def coerce(self, input_data: Any, depth: Optional[int] = None) -> Any:
+    def _coerce(self, input_data: Any, depth: Optional[int] = None) -> Any:
         if depth is None:
             depth = self.max_depth
         if not isinstance(input_data, dict) or depth <= 0:

libs/langgraph/tests/test_pregel.py

@@ -24,12 +24,14 @@ from typing import (
     Any,
     Dict,
     Generator,
+    Generic,
     Iterator,
     List,
     Literal,
     Optional,
     Sequence,
     Tuple,
+    TypeVar,
     Union,
     get_type_hints,
 )
@@ -37,11 +39,7 @@ from typing import (
 import httpx
 import pytest
 from langchain_core.language_models import GenericFakeChatModel
-from langchain_core.runnables import (
-    RunnableConfig,
-    RunnableLambda,
-    RunnablePassthrough,
-)
+from langchain_core.runnables import RunnableConfig, RunnableLambda, RunnablePassthrough
 from langchain_core.runnables.graph import Edge
 from langsmith import traceable
 from pytest_mock import MockerFixture
@@ -1339,22 +1337,26 @@ def test_pending_writes_resume(
             "configurable": {
                 "thread_id": "1",
                 "checkpoint_ns": "",
-                "checkpoint_id": checkpoints[2].config["configurable"]["checkpoint_id"]
-                if checkpoint_during
-                else AnyStr(),
+                "checkpoint_id": (
+                    checkpoints[2].config["configurable"]["checkpoint_id"]
+                    if checkpoint_during
+                    else AnyStr()
+                ),
             }
         },
-        pending_writes=UnsortedSequence(
-            (AnyStr(), "value", 2),
-            (AnyStr(), "__error__", 'ConnectionError("I\'m not good")'),
-            (AnyStr(), "value", 3),
-        )
-        if checkpoint_during
-        else UnsortedSequence(
-            (AnyStr(), "value", 2),
-            (AnyStr(), "__error__", 'ConnectionError("I\'m not good")'),
-            # the write against the previous checkpoint is not saved, as it is
-            # produced in a run where only the next checkpoint (the last) is saved
+        pending_writes=(
+            UnsortedSequence(
+                (AnyStr(), "value", 2),
+                (AnyStr(), "__error__", 'ConnectionError("I\'m not good")'),
+                (AnyStr(), "value", 3),
+            )
+            if checkpoint_during
+            else UnsortedSequence(
+                (AnyStr(), "value", 2),
+                (AnyStr(), "__error__", 'ConnectionError("I\'m not good")'),
+                # the write against the previous checkpoint is not saved, as it is
+                # produced in a run where only the next checkpoint (the last) is saved
+            )
         ),
     )
     if not checkpoint_during:
@@ -3105,10 +3107,10 @@ def test_nested_pydantic_models(version: str) -> None:
     # Import necessary modules
 
     if version == "v1":
-        from pydantic.v1 import (  # type: ignore
+        from pydantic.v1 import (
             BaseModel,
             ByteSize,
-            Field,
+            Field,  # type: ignore
             SecretStr,
             confloat,
             conint,
@@ -3116,10 +3118,10 @@ def test_nested_pydantic_models(version: str) -> None:
             constr,
         )
     else:
-        from pydantic import (  # type: ignore
+        from pydantic import (
             BaseModel,
             ByteSize,
-            Field,
+            Field,  # type: ignore
             SecretStr,
             confloat,
             conint,
@@ -3355,6 +3357,107 @@ def test_nested_pydantic_models(version: str) -> None:
     assert {**new_inputs, **update} == graph.invoke(new_inputs.copy())
 
 
+def test_pydantic_state_field_validator():
+    from pydantic import BaseModel, field_validator, model_validator
+
+    class State(BaseModel):
+        name: str
+        text: str = ""
+        only_root: int = 13
+
+        @field_validator("name", mode="after")
+        @classmethod
+        def validate_name(cls, value):
+            if value[0].islower():
+                raise ValueError("Name must start with a capital letter")
+            return "Validated " + value
+
+        @model_validator(mode="before")
+        @classmethod
+        def validate_amodel(cls, values: "State"):
+            return values | {"only_root": 392}
+
+    input_state = {"name": "John"}
+
+    def process_node(state: State):
+        assert State.model_validate(input_state) == state
+        return {"text": "Hello, " + state.name + "!"}
+
+    builder = StateGraph(state_schema=State)
+    builder.add_node("process", process_node)
+    builder.add_edge(START, "process")
+    builder.add_edge("process", END)
+    g = builder.compile()
+    res = g.invoke(input_state)
+    assert res["text"] == "Hello, Validated John!"
+
+
+def test_pydantic_v1_state_root_validator():
+    from pydantic.v1 import BaseModel, root_validator
+
+    class State(BaseModel):
+        name: str
+        text: str = ""
+        only_root: int = 13
+
+        @root_validator(pre=True)
+        @classmethod
+        def validate(cls, values: dict):
+            values["name"] = "Validated " + values["name"]
+            return values | {"only_root": 396}
+
+    input_state = {"name": "John"}
+
+    def process_node(state: State):
+        assert State(**input_state) == state
+        return {"text": "Hello, " + state.name + "!"}
+
+    builder = StateGraph(state_schema=State)
+    builder.add_node("process", process_node)
+    builder.add_edge(START, "process")
+    builder.add_edge("process", END)
+    g = builder.compile()
+    res = g.invoke(input_state)
+    assert res["text"] == "Hello, Validated John!"
+
+
+def test_pydantic_generics():
+    from pydantic import BaseModel
+
+    class A(BaseModel):
+        a: str
+
+    class B(BaseModel):
+        b: str
+
+    AorB = TypeVar("AorB", A, B)
+
+    class C(BaseModel, Generic[AorB]):
+        c: AorB
+
+    class State(BaseModel):
+        text: str
+        count: int
+        c: C[A]
+
+    input_state = {"text": "1", "count": 0, "c": {"c": {"a": "1"}}}
+    expected_input = State.model_validate(input_state)
+
+    def process_node(state: State):
+        assert state == expected_input
+        new_text = ", the type of c is " + str(type(state.c.c))
+
+        return {"text": state.text + new_text, "count": state.count + 1}
+
+    builder = StateGraph(State)
+    builder.add_node("process", process_node)
+    builder.add_edge(START, "process")
+    g = builder.compile()
+
+    g.invoke(input_state)
+    g.invoke(expected_input)
+
+
 @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC)
 def test_in_one_fan_out_state_graph_waiting_edge_plus_regular(
     request: pytest.FixtureRequest, checkpointer_name: str