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