From 12ad47e4e8c19b795f34e4d98865c9274457a1ab Mon Sep 17 00:00:00 2001 From: William FH <13333726+hinthornw@users.noreply.github.com> Date: Mon, 21 Apr 2025 14:53:58 -0700 Subject: [PATCH] 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> --- .../langgraph/langgraph/graph/schema_utils.py | 20 +++- libs/langgraph/tests/test_pregel.py | 96 ++++++++++++++++--- 2 files changed, 101 insertions(+), 15 deletions(-) diff --git a/libs/langgraph/langgraph/graph/schema_utils.py b/libs/langgraph/langgraph/graph/schema_utils.py index 35ccada7b..c1e6eae5e 100644 --- a/libs/langgraph/langgraph/graph/schema_utils.py +++ b/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: diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index c458c7dbc..5bac66252 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/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