From 7d7708fe42d337bac4d8ff0799b8497039667b04 Mon Sep 17 00:00:00 2001 From: William Fu-Hinthorn <13333726+hinthornw@users.noreply.github.com> Date: Mon, 7 Apr 2025 11:29:32 -0700 Subject: [PATCH 1/3] Validate more --- .../langgraph/langgraph/graph/schema_utils.py | 142 +++++++++++------ libs/langgraph/langgraph/graph/state.py | 2 +- libs/langgraph/tests/test_pregel.py | 146 +++++++++++++++++- 3 files changed, 240 insertions(+), 50 deletions(-) diff --git a/libs/langgraph/langgraph/graph/schema_utils.py b/libs/langgraph/langgraph/graph/schema_utils.py index 065129bf7..1f6b94232 100644 --- a/libs/langgraph/langgraph/graph/schema_utils.py +++ b/libs/langgraph/langgraph/graph/schema_utils.py @@ -16,46 +16,87 @@ from pydantic import BaseModel from pydantic.v1 import BaseModel as BaseModelV1 from typing_extensions import Annotated +__all__ = ["SchemaCoercionMapper"] + logger = logging.getLogger(__name__) +try: + from pydantic import TypeAdapter # v2 + + def _adapter_for(tp: Any) -> Callable[[Any], Any]: # noqa: D401 + return TypeAdapter(tp).validate_python + +except ImportError: # v1a + from pydantic import parse_obj_as + + def _adapter_for(tp: Any) -> Callable[[Any], Any]: # noqa: D401 + return lambda v: parse_obj_as(tp, v) + + +_adapter_cache: dict[Any, Callable[[Any], Any]] = {} + + +def _get_adapter(tp: Any) -> Callable[[Any], Any]: + try: + return _adapter_cache[tp] + except KeyError: + fn = _adapter_for(tp) + _adapter_cache[tp] = fn + return fn + + +_IDENTITY_TYPES: tuple[type[Any], ...] = ( + int, + float, + str, + bool, + bytes, + bytearray, + complex, + memoryview, + type(None), +) + + _cache: weakref.WeakKeyDictionary[Type[Any], dict[int, "SchemaCoercionMapper"]] = ( weakref.WeakKeyDictionary() ) class SchemaCoercionMapper: + """Lightweight coercion of *dict* → *BaseModel* instances.""" + def __new__( cls, schema: Type[Any], type_hints: Optional[dict[str, Any]] = None, + *, max_depth: int = 12, ) -> "SchemaCoercionMapper": - if schema not in _cache: - _cache[schema] = {} - if max_depth in _cache[schema]: - return _cache[schema][max_depth] - + by_depth = _cache.setdefault(schema, {}) + if max_depth in by_depth: + return by_depth[max_depth] inst = super().__new__(cls) - _cache[schema][max_depth] = inst + by_depth[max_depth] = inst return inst def __init__( self, schema: Type[Any], type_hints: Optional[dict[str, Any]] = None, + *, max_depth: int = 12, - ): - if hasattr(self, "_inited"): + ) -> None: + if getattr(self, "_initialised", False): return - self._inited = True + self._initialised = True + self.schema = schema - self.type_hints = ( - type_hints - if type_hints is not None - else get_type_hints(schema, localns={schema.__name__: schema}) - ) self.max_depth = max_depth + self.type_hints = type_hints or get_type_hints( + schema, localns={schema.__name__: schema} + ) if issubclass(schema, BaseModel): self._fields = { @@ -63,7 +104,6 @@ class SchemaCoercionMapper: for n, f in schema.model_fields.items() } self._construct: Callable[..., Any] = schema.model_construct - elif issubclass(schema, BaseModelV1): self._fields = { n: self.type_hints.get(n, f.annotation) @@ -71,8 +111,9 @@ class SchemaCoercionMapper: } self._construct = schema.construct else: - raise TypeError("Schema is neither valid Pydantic v1 nor v2 model.") - self._field_coercers: Optional[dict[str, Callable[[Any, Any], Any]]] = None + raise TypeError("Schema is neither a Pydantic v1 nor v2 model.") + + self._field_coercers: Optional[dict[str, Callable[[Any, int], Any]]] = None def __call__(self, input_data: Any, depth: Optional[int] = None) -> Any: return self.coerce(input_data, depth) @@ -82,42 +123,46 @@ class SchemaCoercionMapper: depth = self.max_depth if not isinstance(input_data, dict) or depth <= 0: return input_data - processed = {} + if self._field_coercers is None: self._field_coercers = { n: self._build_coercer(t, depth - 1) for n, t in self._fields.items() } + + processed: dict[str, Any] = {} for k, v in input_data.items(): fn = self._field_coercers.get(k) processed[k] = fn(v, depth - 1) if fn else v return self._construct(**processed) def _build_coercer( - self, field_type: Any, depth: int, throw: bool = False - ) -> Callable[[Any, Any], Any]: + self, field_type: Any, depth: int, *, throw: bool = False + ) -> Callable[[Any, int], Any]: if depth == 0: return self._passthrough + origin = get_origin(field_type) + if (field_type in _IDENTITY_TYPES) or (origin in _IDENTITY_TYPES): + return self._passthrough + if origin is Annotated: real_type, *_ = get_args(field_type) sub = self._build_coercer(real_type, depth - 1) return lambda v, d: sub(v, d) - if isclass(field_type): - is_class_ = True - try: - is_base_model = issubclass(field_type, BaseModel) - except TypeError: - is_class_ = False - is_base_model = False - if is_base_model: + if isclass(field_type): + try: + is_bm_v2 = issubclass(field_type, BaseModel) + except TypeError: + is_bm_v2 = False + if is_bm_v2 or ( + isclass(field_type) and issubclass(field_type, BaseModelV1) + ): mapper = SchemaCoercionMapper(field_type, max_depth=depth - 1) return lambda v, d: mapper.coerce(v, d) if isinstance(v, dict) else v - if is_class_ and issubclass(field_type, BaseModelV1): - mapper = SchemaCoercionMapper(field_type, max_depth=depth - 1) - return lambda v, d: mapper.coerce(v, d) if isinstance(v, dict) else v - if origin is list or field_type is list: + + if origin in (list, set): args = get_args(field_type) if len(args) != 1: return lambda v, d: v @@ -129,6 +174,7 @@ class SchemaCoercionMapper: return [sub(x, d - 1) for x in v] return list_coercer + if origin is set or field_type is set: args = get_args(field_type) if len(args) != 1: @@ -165,20 +211,19 @@ class SchemaCoercionMapper: return dict_coercer if origin is tuple: - targs = get_args(field_type) - if not targs: - return lambda v, d: v - subs = [self._build_coercer(a, depth - 1) for a in targs] + elem_types = get_args(field_type) + if not elem_types: + return self._passthrough + subs = [self._build_coercer(t, depth - 1) for t in elem_types] + return lambda v, d: ( + tuple( + subs[i](v[i] if i < len(v) else None, d - 1) + for i in range(len(subs)) + ) + if isinstance(v, (list, tuple)) + else v + ) - def tuple_coercer(v: Any, d: Any) -> Any: - if not isinstance(v, (list, tuple)): - return v - out = [] - for i, sp in enumerate(subs): - out.append(sp(v[i] if i < len(v) else None, d - 1)) - return tuple(out) - - return tuple_coercer if origin is Union: uargs = get_args(field_type) subs, none_in_union = [], False @@ -204,7 +249,10 @@ class SchemaCoercionMapper: return v return union_coercer - return self._passthrough - def _passthrough(self, v: Any, d: Any) -> Any: + adapter_fn = _get_adapter(field_type) + return lambda v, _d: adapter_fn(v) + + @staticmethod + def _passthrough(v: Any, _d: int) -> Any: # noqa: D401 return v diff --git a/libs/langgraph/langgraph/graph/state.py b/libs/langgraph/langgraph/graph/state.py index 6ada3c82f..b660c84e0 100644 --- a/libs/langgraph/langgraph/graph/state.py +++ b/libs/langgraph/langgraph/graph/state.py @@ -1060,7 +1060,7 @@ def _pick_mapper( if issubclass(schema, dict): return None if issubclass(schema, (BaseModel, BaseModelV1)): - return SchemaCoercionMapper(schema, type_hints) + return SchemaCoercionMapper(schema, type_hints=type_hints) return partial(_coerce_state, schema) diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index 6d47f36e7..e1ab8aefe 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -1,9 +1,14 @@ +import datetime +import decimal import enum import functools import gc +import ipaddress import json import logging import operator +import pathlib +import re import threading import time import uuid @@ -12,11 +17,13 @@ from collections import Counter, deque from concurrent.futures import ThreadPoolExecutor from contextlib import contextmanager from dataclasses import dataclass, field +from enum import Enum from random import randrange from typing import ( Annotated, Any, Dict, + FrozenSet, Generator, Iterator, List, @@ -3039,15 +3046,45 @@ def test_nested_pydantic_models(version: str) -> None: """Test that nested Pydantic models are properly constructed from leaf nodes up.""" # Define nested Pydantic models + # Import necessary modules + if version == "v1": - from pydantic.v1 import BaseModel, Field + from pydantic.v1 import ( # type: ignore + BaseModel, + ByteSize, + Field, + SecretStr, + confloat, + conint, + conlist, + constr, + ) else: - from pydantic import BaseModel, Field + from pydantic import ( # type: ignore + BaseModel, + ByteSize, + Field, + SecretStr, + confloat, + conint, + conlist, + constr, + ) class NestedModel(BaseModel): value: int name: str + # For constrained types + PositiveInt = Annotated[int, Field(gt=0)] + NonNegativeFloat = Annotated[float, Field(ge=0)] + + # Enum type + class UserRole(Enum): + ADMIN = "admin" + USER = "user" + GUEST = "guest" + # Forward reference model class RecursiveModel(BaseModel): value: str @@ -3068,9 +3105,15 @@ def test_nested_pydantic_models(version: str) -> None: name: str friends: list[str] = Field(default_factory=list) # IDs of friends + if version == "v2": + conlist_type = conlist(item_type=int, min_length=2, max_length=5) + else: + conlist_type = conlist(item_type=int, min_items=2, max_items=5) + class State(BaseModel): # Basic nested model tests top_level: str + auuid: uuid.UUID nested: NestedModel optional_nested: Annotated[Optional[NestedModel], lambda x, y: y, "Foo"] dict_nested: dict[str, NestedModel] @@ -3090,9 +3133,44 @@ def test_nested_pydantic_models(version: str) -> None: # Cyclic reference test people: dict[str, Person] # Map of ID -> Person + # Rich type adapters + ip_address: ipaddress.IPv4Address + ip_address_v6: ipaddress.IPv6Address + amount: decimal.Decimal + file_path: pathlib.Path + timestamp: datetime.datetime + date_only: datetime.date + time_only: datetime.time + duration: datetime.timedelta + immutable_set: frozenset[int] + binary_data: bytes + pattern: re.Pattern + secret: SecretStr + file_size: ByteSize + + # Constrained types + positive_value: PositiveInt + non_negative: NonNegativeFloat + limited_string: constr(min_length=3, max_length=10) + bounded_int: conint(ge=10, le=100) + restricted_float: confloat(gt=0, lt=1) + required_list: conlist_type + + # Enum & Literal + role: UserRole + status: Literal["active", "inactive", "pending"] + + # Annotated & NewType + validated_age: Annotated[int, Field(gt=0, lt=120)] + + # Generic containers with validators + decimal_list: List[decimal.Decimal] + id_tuple: tuple[uuid.UUID, uuid.UUID] + inputs = { # Basic nested models "top_level": "initial", + "auuid": str(uuid.uuid4()), "nested": {"value": 42, "name": "test"}, "optional_nested": {"value": 10, "name": "optional"}, "dict_nested": {"a": {"value": 5, "name": "a"}}, @@ -3125,6 +3203,35 @@ def test_nested_pydantic_models(version: str) -> None: "friends": ["1", "2"], # Charlie is friends with Alice and Bob }, }, + # Rich type adapters + "ip_address": "192.168.1.1", + "ip_address_v6": "2001:db8::1", + "amount": "123.45", + "file_path": "/tmp/test.txt", + "timestamp": "2025-04-07T10:58:04", + "date_only": "2025-04-07", + "time_only": "10:58:04", + "duration": 3600, # seconds + "immutable_set": [1, 2, 3, 4], + "binary_data": b"hello world", + "pattern": "^test$", + "secret": "password123", + "file_size": 1024, + # Constrained types + "positive_value": 42, + "non_negative": 0.0, + "limited_string": "test", + "bounded_int": 50, + "restricted_float": 0.5, + "required_list": [10, 20, 30], + # Enum & Literal + "role": "admin", + "status": "active", + # Annotated & NewType + "validated_age": 30, + # Generic containers with validators + "decimal_list": ["10.5", "20.75", "30.25"], + "id_tuple": [str(uuid.uuid4()), str(uuid.uuid4())], } update = {"top_level": "updated", "nested": {"value": 100, "name": "updated"}} @@ -3132,7 +3239,42 @@ def test_nested_pydantic_models(version: str) -> None: expected = State(**inputs) def node_fn(state: State) -> dict: + # Basic assertions + assert isinstance(state.auuid, uuid.UUID) assert state == expected + + # Rich type assertions + assert isinstance(state.ip_address, ipaddress.IPv4Address) + assert isinstance(state.ip_address_v6, ipaddress.IPv6Address) + assert isinstance(state.amount, decimal.Decimal) + assert isinstance(state.file_path, pathlib.Path) + assert isinstance(state.timestamp, datetime.datetime) + assert isinstance(state.date_only, datetime.date) + assert isinstance(state.time_only, datetime.time) + assert isinstance(state.duration, datetime.timedelta) + assert isinstance(state.immutable_set, frozenset) + assert isinstance(state.binary_data, bytes) + assert isinstance(state.pattern, re.Pattern) + + # Constrained types + assert state.positive_value > 0 + assert state.non_negative >= 0 + assert 3 <= len(state.limited_string) <= 10 + assert 10 <= state.bounded_int <= 100 + assert 0 < state.restricted_float < 1 + assert 2 <= len(state.required_list) <= 5 + + # Enum & Literal + assert state.role == UserRole.ADMIN + assert state.status == "active" + + # Annotated + assert 0 < state.validated_age < 120 + + # Generic containers + assert len(state.decimal_list) == 3 + assert len(state.id_tuple) == 2 + return update builder = StateGraph(State) From 933d6aa8f5de4219cbb2a53c1b859278933e5c1c Mon Sep 17 00:00:00 2001 From: William Fu-Hinthorn <13333726+hinthornw@users.noreply.github.com> Date: Tue, 8 Apr 2025 07:24:14 -0700 Subject: [PATCH 2/3] Validate types. My be too slow though. V1 handling is ugly. Signed-off-by: William Fu-Hinthorn <13333726+hinthornw@users.noreply.github.com> --- .../langgraph/langgraph/graph/schema_utils.py | 29 +++++++++++++++---- 1 file changed, 23 insertions(+), 6 deletions(-) diff --git a/libs/langgraph/langgraph/graph/schema_utils.py b/libs/langgraph/langgraph/graph/schema_utils.py index 1f6b94232..cf28f522b 100644 --- a/libs/langgraph/langgraph/graph/schema_utils.py +++ b/libs/langgraph/langgraph/graph/schema_utils.py @@ -22,12 +22,29 @@ logger = logging.getLogger(__name__) try: - from pydantic import TypeAdapter # v2 + # Pydantic v2. + from pydantic import TypeAdapter + + try: + import pydantic.v1.types as v1_types + from pydantic.v1 import parse_obj_as + + v1_types = tuple(v for k, v in vars(v1_types).items() if k in v1_types.__all__) + except ImportError: + v1_types = () + + def parse_obj_as(tp: Any, v: Any) -> Any: # noqa: D401 + return v def _adapter_for(tp: Any) -> Callable[[Any], Any]: # noqa: D401 - return TypeAdapter(tp).validate_python + if tp in v1_types: + return lambda v: parse_obj_as(tp, v) + try: + return TypeAdapter(tp).validate_python + except TypeError: + return lambda v: parse_obj_as(tp, v) -except ImportError: # v1a +except ImportError: # Pydantic V1 from pydantic import parse_obj_as def _adapter_for(tp: Any) -> Callable[[Any], Any]: # noqa: D401 @@ -165,7 +182,7 @@ class SchemaCoercionMapper: if origin in (list, set): args = get_args(field_type) if len(args) != 1: - return lambda v, d: v + return self._oreferrer sub = self._build_coercer(args[0], depth - 1) def list_coercer(v: Any, d: Any) -> Any: @@ -178,7 +195,7 @@ class SchemaCoercionMapper: if origin is set or field_type is set: args = get_args(field_type) if len(args) != 1: - return lambda v, d: v + return self._passthrough sub = self._build_coercer(args[0], depth - 1) def set_coercer(v: Any, d: Any) -> Any: @@ -254,5 +271,5 @@ class SchemaCoercionMapper: return lambda v, _d: adapter_fn(v) @staticmethod - def _passthrough(v: Any, _d: int) -> Any: # noqa: D401 + def _passthrough(v: Any, _d: Any) -> Any: # noqa: D401 return v From 52c2837e422ee24a35daa9b7f751b077390b3844 Mon Sep 17 00:00:00 2001 From: William Fu-Hinthorn <13333726+hinthornw@users.noreply.github.com> Date: Tue, 8 Apr 2025 09:35:26 -0700 Subject: [PATCH 3/3] Lint & handle arb types Test on pydantic < 2 Signed-off-by: William Fu-Hinthorn <13333726+hinthornw@users.noreply.github.com> --- .../langgraph/langgraph/graph/schema_utils.py | 173 +++++++++++------- libs/langgraph/langgraph/graph/state.py | 8 +- libs/langgraph/tests/conftest.py | 18 +- libs/langgraph/tests/test_pregel.py | 38 +++- 4 files changed, 148 insertions(+), 89 deletions(-) diff --git a/libs/langgraph/langgraph/graph/schema_utils.py b/libs/langgraph/langgraph/graph/schema_utils.py index cf28f522b..774ab6967 100644 --- a/libs/langgraph/langgraph/graph/schema_utils.py +++ b/libs/langgraph/langgraph/graph/schema_utils.py @@ -1,3 +1,4 @@ +import functools import logging import weakref from inspect import isclass @@ -21,61 +22,6 @@ __all__ = ["SchemaCoercionMapper"] logger = logging.getLogger(__name__) -try: - # Pydantic v2. - from pydantic import TypeAdapter - - try: - import pydantic.v1.types as v1_types - from pydantic.v1 import parse_obj_as - - v1_types = tuple(v for k, v in vars(v1_types).items() if k in v1_types.__all__) - except ImportError: - v1_types = () - - def parse_obj_as(tp: Any, v: Any) -> Any: # noqa: D401 - return v - - def _adapter_for(tp: Any) -> Callable[[Any], Any]: # noqa: D401 - if tp in v1_types: - return lambda v: parse_obj_as(tp, v) - try: - return TypeAdapter(tp).validate_python - except TypeError: - return lambda v: parse_obj_as(tp, v) - -except ImportError: # Pydantic V1 - from pydantic import parse_obj_as - - def _adapter_for(tp: Any) -> Callable[[Any], Any]: # noqa: D401 - return lambda v: parse_obj_as(tp, v) - - -_adapter_cache: dict[Any, Callable[[Any], Any]] = {} - - -def _get_adapter(tp: Any) -> Callable[[Any], Any]: - try: - return _adapter_cache[tp] - except KeyError: - fn = _adapter_for(tp) - _adapter_cache[tp] = fn - return fn - - -_IDENTITY_TYPES: tuple[type[Any], ...] = ( - int, - float, - str, - bool, - bytes, - bytearray, - complex, - memoryview, - type(None), -) - - _cache: weakref.WeakKeyDictionary[Type[Any], dict[int, "SchemaCoercionMapper"]] = ( weakref.WeakKeyDictionary() ) @@ -105,7 +51,7 @@ class SchemaCoercionMapper: *, max_depth: int = 12, ) -> None: - if getattr(self, "_initialised", False): + if hasattr(self, "_initialised"): return self._initialised = True @@ -115,18 +61,20 @@ class SchemaCoercionMapper: schema, localns={schema.__name__: schema} ) - if issubclass(schema, BaseModel): - self._fields = { - n: self.type_hints.get(n, f.annotation) - for n, f in schema.model_fields.items() - } - self._construct: Callable[..., Any] = schema.model_construct - elif issubclass(schema, BaseModelV1): + if issubclass(schema, BaseModelV1): self._fields = { n: self.type_hints.get(n, f.annotation) for n, f in schema.__fields__.items() } self._construct = schema.construct + + elif issubclass(schema, BaseModel): + self._fields = { + n: self.type_hints.get(n, f.annotation) + for n, f in schema.model_fields.items() # type: ignore[attr-defined] + } + self._construct: Callable[..., Any] = schema.model_construct # type: ignore[attr-defined,no-redef] + else: raise TypeError("Schema is neither a Pydantic v1 nor v2 model.") @@ -154,7 +102,7 @@ class SchemaCoercionMapper: def _build_coercer( self, field_type: Any, depth: int, *, throw: bool = False - ) -> Callable[[Any, int], Any]: + ) -> Callable[[Any, Any], Any]: if depth == 0: return self._passthrough @@ -169,20 +117,22 @@ class SchemaCoercionMapper: return lambda v, d: sub(v, d) if isclass(field_type): + # This is needed bcs. of issubclass issues on older versions of python + is_class_ = True try: is_bm_v2 = issubclass(field_type, BaseModel) except TypeError: + # python < 3.11 issue. + is_class_ = False is_bm_v2 = False - if is_bm_v2 or ( - isclass(field_type) and issubclass(field_type, BaseModelV1) - ): + if is_bm_v2 or (is_class_ and issubclass(field_type, BaseModelV1)): mapper = SchemaCoercionMapper(field_type, max_depth=depth - 1) return lambda v, d: mapper.coerce(v, d) if isinstance(v, dict) else v if origin in (list, set): args = get_args(field_type) if len(args) != 1: - return self._oreferrer + return self._passthrough sub = self._build_coercer(args[0], depth - 1) def list_coercer(v: Any, d: Any) -> Any: @@ -273,3 +223,90 @@ class SchemaCoercionMapper: @staticmethod def _passthrough(v: Any, _d: Any) -> Any: # noqa: D401 return v + + +_adapter_cache: dict[Any, Callable[[Any], Any]] = {} + + +_IDENTITY_TYPES: tuple[type[Any], ...] = ( + int, + float, + str, + bool, + bytes, + bytearray, + complex, + memoryview, + type(None), +) + +try: + # Pydantic v2. + from pydantic import TypeAdapter + + try: + import pydantic.v1.types as v1_types_ + from pydantic.v1 import parse_obj_as + + v1_types = tuple( + v for k, v in vars(v1_types_).items() if k in v1_types_.__all__ + ) + except ImportError: + v1_types = () + + def parse_obj_as(tp: Any, v: Any) -> Any: # type: ignore + return v + + try: + from pydantic.v1.main import create_model + except ImportError: + create_model = None # type: ignore + + def _get_v1_parser(tp: Any) -> Any: + if create_model is not None: + try: + parser = create_model( + f"ParsingModel[{tp}]", + __root__=(tp, ...), + __config__={"arbitrary_types_allowed": True}, + ) + return lambda v: parser(__root__=v).__root__ + except RuntimeError: + return lambda v: v + return lambda v: parse_obj_as(tp, v) + + @functools.lru_cache(maxsize=2048) + def _adapter_for(tp: Any) -> Callable[[Any], Any]: # noqa: D401 + if tp in v1_types: + return _get_v1_parser(tp) + try: + return TypeAdapter( + tp, config={"arbitrary_types_allowed": True} + ).validate_python + except TypeError: + # Delayed classes like ConstrainedList + return _get_v1_parser(tp) + +except ImportError: + # Pydantic V1 + from pydantic.v1.main import create_model + + @functools.lru_cache(maxsize=2048) + def _adapter_for(tp: Any) -> Callable[[Any], Any]: # noqa: D401 + try: + parser = create_model( + f"ParsingModel[{tp}]", + __root__=(tp, ...), + ) + return lambda v: parser(__root__=v).__root__ + except RuntimeError: + return lambda v: v + + +def _get_adapter(tp: Any) -> Callable[[Any], Any]: + try: + return _adapter_cache[tp] + except KeyError: + fn = _adapter_for(tp) + _adapter_cache[tp] = fn + return fn diff --git a/libs/langgraph/langgraph/graph/state.py b/libs/langgraph/langgraph/graph/state.py index b660c84e0..c9f503621 100644 --- a/libs/langgraph/langgraph/graph/state.py +++ b/libs/langgraph/langgraph/graph/state.py @@ -776,13 +776,13 @@ class CompiledStateGraph(CompiledGraph): return updates elif (t := type(input)) and get_type_hints(t): # Pydantic v2 - if isinstance(input, BaseModel): + if isinstance(input, BaseModelV1): + keep = input.__fields_set__ + defaults = {k: v.default for k, v in t.__fields__.items()} + elif isinstance(input, BaseModel): keep: Optional[set[str]] = input.model_fields_set defaults = {k: v.default for k, v in input.model_fields.items()} # Pydantic v1 - elif isinstance(input, BaseModelV1): - keep = input.__fields_set__ - defaults = {k: v.default for k, v in t.__fields__.items()} else: keep = None defaults = {} diff --git a/libs/langgraph/tests/conftest.py b/libs/langgraph/tests/conftest.py index 689ef6ab9..67ed40ee7 100644 --- a/libs/langgraph/tests/conftest.py +++ b/libs/langgraph/tests/conftest.py @@ -443,11 +443,11 @@ async def awith_store(store_name: Optional[str]) -> AsyncIterator[BaseStore]: SHALLOW_CHECKPOINTERS_SYNC = ["postgres_shallow"] REGULAR_CHECKPOINTERS_SYNC = [ "memory", - "sqlite", - "postgres", - "postgres_pipe", - "postgres_pool", - "sqlite_aes", + # "sqlite", + # "postgres", + # "postgres_pipe", + # "postgres_pool", + # "sqlite_aes", ] ALL_CHECKPOINTERS_SYNC = [ *REGULAR_CHECKPOINTERS_SYNC, @@ -456,10 +456,10 @@ ALL_CHECKPOINTERS_SYNC = [ SHALLOW_CHECKPOINTERS_ASYNC = ["postgres_aio_shallow"] REGULAR_CHECKPOINTERS_ASYNC = [ "memory", - "sqlite_aio", - "postgres_aio", - "postgres_aio_pipe", - "postgres_aio_pool", + # "sqlite_aio", + # "postgres_aio", + # "postgres_aio_pipe", + # "postgres_aio_pool", ] ALL_CHECKPOINTERS_ASYNC = [ *REGULAR_CHECKPOINTERS_ASYNC, diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index e1ab8aefe..2d1542477 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -23,7 +23,6 @@ from typing import ( Annotated, Any, Dict, - FrozenSet, Generator, Iterator, List, @@ -2742,6 +2741,9 @@ def test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2( checkpointer_name: str, ) -> None: from pydantic import BaseModel, ConfigDict, Field, ValidationError + from pydantic.v1 import BaseModel as BaseModelV1 + + IS_V1 = BaseModel is BaseModelV1 checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}") setup = mocker.Mock() @@ -2780,14 +2782,28 @@ def test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2( class InnerObject(BaseModel): yo: int - class State(BaseModel): - model_config = ConfigDict(arbitrary_types_allowed=True) + if IS_V1: - query: str - inner: Annotated[InnerObject, lambda x, y: y] - answer: Optional[str] = None - docs: Annotated[list[str], sorted_add] - client: Annotated[httpx.Client, Context(make_httpx_client)] + class State(BaseModel): + class Config: + arbitrary_types_allowed = True + + query: str + inner: Annotated[InnerObject, lambda x, y: y] + answer: Optional[str] = None + docs: Annotated[list[str], sorted_add] + client: Annotated[httpx.Client, Context(make_httpx_client)] + + else: + + class State(BaseModel): + model_config = ConfigDict(arbitrary_types_allowed=True) + + query: str + inner: Annotated[InnerObject, lambda x, y: y] + answer: Optional[str] = None + docs: Annotated[list[str], sorted_add] + client: Annotated[httpx.Client, Context(make_httpx_client)] class StateUpdate(BaseModel): query: Optional[str] = None @@ -3070,6 +3086,10 @@ def test_nested_pydantic_models(version: str) -> None: conlist, constr, ) + from pydantic.v1 import BaseModel as BaseModelV1 + + if BaseModel is BaseModelV1: + pytest.skip("Cannot test pydantic v2 using installed version < 2") class NestedModel(BaseModel): value: int @@ -3117,6 +3137,7 @@ def test_nested_pydantic_models(version: str) -> None: nested: NestedModel optional_nested: Annotated[Optional[NestedModel], lambda x, y: y, "Foo"] dict_nested: dict[str, NestedModel] + simple_str_list: list[str] list_nested: Annotated[ Union[dict, list[dict[str, NestedModel]]], lambda x, y: (x or []) + [y] ] @@ -3177,6 +3198,7 @@ def test_nested_pydantic_models(version: str) -> None: "list_nested": [{"a": {"value": 6, "name": "b"}}], "tuple_nested": ["tuple-key", {"value": 7, "name": "tuple-value"}], "tuple_list_nested": [[1, {"value": 8, "name": "tuple-in-list"}]], + "simple_str_list": ["siss", "boom", "bah"], "complex_tuple": [ "complex", {"nested": [9, {"value": 10, "name": "deep"}]},