diff --git a/libs/langgraph/langgraph/graph/schema_utils.py b/libs/langgraph/langgraph/graph/schema_utils.py index 065129bf7..35ccada7b 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 @@ -16,6 +17,8 @@ from pydantic import BaseModel from pydantic.v1 import BaseModel as BaseModelV1 from typing_extensions import Annotated +__all__ = ["SchemaCoercionMapper"] + logger = logging.getLogger(__name__) @@ -25,54 +28,60 @@ _cache: weakref.WeakKeyDictionary[Type[Any], dict[int, "SchemaCoercionMapper"]] 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 hasattr(self, "_initialised"): return - self._inited = True + self._initialised = True + self.schema = schema + self.max_depth = max_depth + 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 - 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() + } + self._construct: Callable[..., Any] = schema.model_construct # type: ignore + 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,45 +91,51 @@ 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 + self, field_type: Any, depth: int, *, throw: bool = False ) -> Callable[[Any, Any], 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): + # This is needed bcs. of issubclass issues on older versions of python is_class_ = True try: - is_base_model = issubclass(field_type, BaseModel) + is_bm_v2 = issubclass(field_type, BaseModel) except TypeError: + # python < 3.11 issue. is_class_ = False - is_base_model = False + is_bm_v2 = False + 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 is_base_model: - 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 is list: 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 list_coercer(v: Any, d: Any) -> Any: @@ -129,15 +144,21 @@ 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: - return lambda v, d: v - sub = self._build_coercer(args[0], depth - 1) + if len(args) > 1: + return self._passthrough + elif len(args) == 1: + sub = self._build_coercer(args[0], depth - 1) + else: + sub = None # type: ignore def set_coercer(v: Any, d: Any) -> Any: if not isinstance(v, (list, tuple, set)): return v + if sub is None: + return set(v) return {sub(x, d - 1) for x in v} return set_coercer @@ -165,20 +186,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 +224,97 @@ 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: 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 import parse_obj_as + 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, ...), + ) + return lambda v: parser(__root__=v).__root__ # type: ignore + 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__ # type: ignore + 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 6ada3c82f..093faf5d4 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): - keep: Optional[set[str]] = input.model_fields_set + if isinstance(input, BaseModelV1): + keep: Optional[set[str]] = input.__fields_set__ + defaults = {k: v.default for k, v in t.__fields__.items()} + elif isinstance(input, BaseModel): + keep = 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 = {} @@ -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 486b4cdf9..be3f94ec7 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,6 +17,7 @@ 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, @@ -2775,6 +2781,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() @@ -2813,14 +2822,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 @@ -3079,15 +3102,49 @@ 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, + ) + 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 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 @@ -3108,12 +3165,19 @@ 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] + simple_str_list: list[str] list_nested: Annotated[ Union[dict, list[dict[str, NestedModel]]], lambda x, y: (x or []) + [y] ] @@ -3130,15 +3194,51 @@ 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"}}, "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"}]}, @@ -3165,6 +3265,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"}} @@ -3172,7 +3301,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) diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index dd146dbb7..0e74f0b74 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/libs/langgraph/tests/test_pregel_async.py @@ -4638,6 +4638,7 @@ async def test_nested_pydantic_models(version: str) -> None: optional_nested: Optional[NestedModel] = None dict_nested: dict[str, NestedModel] my_set: set[int] + another_set: set my_enum: MyEnum list_nested: Annotated[ Union[dict, list[dict[str, NestedModel]]], lambda x, y: (x or []) + [y] @@ -4666,6 +4667,7 @@ async def test_nested_pydantic_models(version: str) -> None: "nested": {"value": 42, "name": "test"}, "optional_nested": {"value": 10, "name": "optional"}, "my_set": [1, 2, 7], + "another_set": ["foo", 3], "my_enum": MyEnum.B, "my_typed_dict": {"x": 1, "my_enum": MyEnum.A}, "dict_nested": {"a": {"value": 5, "name": "a"}},