Use model_validate if needed (#4363)

If the state schema uses validators, skip the model construct
optimization.

For context, pydantic state can be significantly slower to run than
typed dict and dataclass states due to the full recursive validation.

We have some optimizations to reduce the impact of this (using cached
validators with model_construct), but this doesn't handle things like
field_validator.

We prefer correctness over performance, obviously.

Resolves: https://github.com/langchain-ai/langgraph/issues/4074

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

@@ -1339,22 +1339,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:
@@ -3355,6 +3359,70 @@ 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!"
+
+
 @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