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] 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)