diff --git a/libs/langgraph/langgraph/graph/schema_utils.py b/libs/langgraph/langgraph/graph/schema_utils.py index bf8d5c4a0..065129bf7 100644 --- a/libs/langgraph/langgraph/graph/schema_utils.py +++ b/libs/langgraph/langgraph/graph/schema_utils.py @@ -9,6 +9,7 @@ from typing import ( Union, get_args, get_origin, + get_type_hints, ) from pydantic import BaseModel @@ -18,34 +19,56 @@ from typing_extensions import Annotated logger = logging.getLogger(__name__) -class SchemaCoercionMapper: - _cache: weakref.WeakKeyDictionary[Type[Any], dict[int, "SchemaCoercionMapper"]] = ( - weakref.WeakKeyDictionary() - ) +_cache: weakref.WeakKeyDictionary[Type[Any], dict[int, "SchemaCoercionMapper"]] = ( + weakref.WeakKeyDictionary() +) - def __new__(cls, schema: Type[Any], max_depth: int = 5) -> "SchemaCoercionMapper": - if schema not in cls._cache: - cls._cache[schema] = {} - if max_depth in cls._cache[schema]: - return cls._cache[schema][max_depth] + +class SchemaCoercionMapper: + 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] inst = super().__new__(cls) - cls._cache[schema][max_depth] = inst + _cache[schema][max_depth] = inst return inst - def __init__(self, schema: Type[Any], max_depth: int = 5): + def __init__( + self, + schema: Type[Any], + type_hints: Optional[dict[str, Any]] = None, + max_depth: int = 12, + ): if hasattr(self, "_inited"): return self._inited = 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 - if hasattr(schema, "model_fields") and hasattr(schema, "model_construct"): - self._fields = {n: f.annotation for n, f in schema.model_fields.items()} - self._construct = schema.model_construct - elif hasattr(schema, "__fields__") and callable( - getattr(schema, "construct", None) - ): - self._fields = {n: f.annotation for n, f in schema.__fields__.items()} + + 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): + self._fields = { + n: self.type_hints.get(n, f.annotation) + for n, f in schema.__fields__.items() + } self._construct = schema.construct else: raise TypeError("Schema is neither valid Pydantic v1 nor v2 model.") @@ -62,18 +85,23 @@ class SchemaCoercionMapper: processed = {} if self._field_coercers is None: self._field_coercers = { - n: self._build_coercer(t) for n, t in self._fields.items() + n: self._build_coercer(t, depth - 1) for n, t in self._fields.items() } 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) -> Callable[[Any, Any], Any]: + def _build_coercer( + 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 origin is Annotated: real_type, *_ = get_args(field_type) - sub = self._build_coercer(real_type) + sub = self._build_coercer(real_type, depth - 1) return lambda v, d: sub(v, d) if isclass(field_type): is_class_ = True @@ -84,39 +112,54 @@ class SchemaCoercionMapper: is_base_model = False if is_base_model: - mapper = SchemaCoercionMapper(field_type, self.max_depth) + 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, self.max_depth) + 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: args = get_args(field_type) if len(args) != 1: return lambda v, d: v - sub = self._build_coercer(args[0]) + sub = self._build_coercer(args[0], depth - 1) def list_coercer(v: Any, d: Any) -> Any: if not isinstance(v, (list, tuple)): - raise TypeError(f"Expected list, got {type(v).__name__}") + return v 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) + + def set_coercer(v: Any, d: Any) -> Any: + if not isinstance(v, (list, tuple, set)): + return v + return {sub(x, d - 1) for x in v} + + return set_coercer if origin is dict or field_type is dict: args = get_args(field_type) if len(args) != 2: - def plain_dict_coercer(v: Any, d: Any) -> Any: + def dict_coercer(v: Any, d: Any) -> Any: if not isinstance(v, dict): - raise TypeError(f"Expected dict, got {type(v).__name__}") + if throw: + raise TypeError("Expected dict, got %s" % type(v)) return v - return plain_dict_coercer - k_sub = self._build_coercer(args[0]) - v_sub = self._build_coercer(args[1]) + return dict_coercer + k_sub = self._build_coercer(args[0], depth - 1) + v_sub = self._build_coercer(args[1], depth - 1) def dict_coercer(v: Any, d: Any) -> Any: if not isinstance(v, dict): - raise TypeError(f"Expected dict, got {type(v).__name__}") + if throw: + raise TypeError("Expected dict, got %s" % type(v)) + return v return {k_sub(k, d - 1): v_sub(val, d - 1) for k, val in v.items()} return dict_coercer @@ -125,11 +168,11 @@ class SchemaCoercionMapper: targs = get_args(field_type) if not targs: return lambda v, d: v - subs = [self._build_coercer(a) for a in targs] + subs = [self._build_coercer(a, depth - 1) for a in targs] def tuple_coercer(v: Any, d: Any) -> Any: if not isinstance(v, (list, tuple)): - raise TypeError(f"Expected tuple-like, got {type(v).__name__}") + return v out = [] for i, sp in enumerate(subs): out.append(sp(v[i] if i < len(v) else None, d - 1)) @@ -139,11 +182,13 @@ class SchemaCoercionMapper: if origin is Union: uargs = get_args(field_type) subs, none_in_union = [], False - for arg in uargs: + for ix, arg in enumerate(uargs): if arg is type(None): none_in_union = True else: - subs.append(self._build_coercer(arg)) + subs.append( + self._build_coercer(arg, depth - 1, throw=ix < len(uargs) - 1) + ) def union_coercer(v: Any, d: Any) -> Any: if v is None and none_in_union: @@ -152,11 +197,14 @@ class SchemaCoercionMapper: for sp in subs: try: return sp(v, d - 1) - except Exception as e: + except TypeError as e: err = e if err: raise err return v return union_coercer - return lambda v, d: v + return self._passthrough + + def _passthrough(self, v: Any, d: Any) -> Any: + return v diff --git a/libs/langgraph/langgraph/graph/state.py b/libs/langgraph/langgraph/graph/state.py index d9009d200..684e61c8e 100644 --- a/libs/langgraph/langgraph/graph/state.py +++ b/libs/langgraph/langgraph/graph/state.py @@ -185,6 +185,7 @@ class StateGraph(Graph): self.schemas = {} self.channels = {} self.managed = {} + self.type_hints: dict[Type[Any], dict[str, Any]] = {} self.schema = state_schema self.input = input self.output = output @@ -203,7 +204,7 @@ class StateGraph(Graph): def _add_schema(self, schema: Type[Any], /, allow_managed: bool = True) -> None: if schema not in self.schemas: _warn_invalid_state_schema(schema) - channels, managed = _get_channels(schema) + channels, managed, type_hints = _get_channels(schema) if managed and not allow_managed: names = ", ".join(managed) schema_name = getattr(schema, "__name__", "") @@ -212,6 +213,7 @@ class StateGraph(Graph): " Managed channels are not permitted in Input/Output schema." ) self.schemas[schema] = {**channels, **managed} + self.type_hints[schema] = type_hints for key, channel in channels.items(): if key in self.channels: if self.channels[key] != channel: @@ -829,7 +831,11 @@ class CompiledStateGraph(CompiledGraph): # read state keys and managed values channels=(list(input_values) if is_single_input else input_values), # coerce state dict to schema class (eg. pydantic model) - mapper=_pick_mapper(list(input_values), input_schema), + mapper=_pick_mapper( + list(input_values), + input_schema, + self.builder.type_hints[input_schema], + ), writers=[ # publish to this channel and state keys ChannelWrite( @@ -939,12 +945,12 @@ def _get_state_reader( select=select[0] if select == ["__root__"] else select, fresh=True, # coerce state dict to schema class (eg. pydantic model) - mapper=_pick_mapper(state_keys, schema), + mapper=_pick_mapper(state_keys, schema, builder.type_hints[schema]), ) def _pick_mapper( - state_keys: Sequence[str], schema: Type[Any] + state_keys: Sequence[str], schema: Type[Any], type_hints: Optional[dict[str, Any]] ) -> Optional[Callable[[Any], Any]]: if state_keys == ["__root__"]: return None @@ -952,7 +958,7 @@ def _pick_mapper( if issubclass(schema, dict): return None if issubclass(schema, (BaseModel, BaseModelV1)): - return SchemaCoercionMapper(schema) + return SchemaCoercionMapper(schema, type_hints) return partial(_coerce_state, schema) @@ -1014,18 +1020,24 @@ CONTROL_BRANCH = Branch(CONTROL_BRANCH_PATH, None) def _get_channels( schema: Type[dict], -) -> tuple[dict[str, BaseChannel], dict[str, ManagedValueSpec]]: +) -> tuple[dict[str, BaseChannel], dict[str, ManagedValueSpec], dict[str, Any]]: if not hasattr(schema, "__annotations__"): - return {"__root__": _get_channel("__root__", schema, allow_managed=False)}, {} + return ( + {"__root__": _get_channel("__root__", schema, allow_managed=False)}, + {}, + {}, + ) + type_hints = get_type_hints(schema, include_extras=True) all_keys = { name: _get_channel(name, typ) - for name, typ in get_type_hints(schema, include_extras=True).items() + for name, typ in type_hints.items() if name != "__slots__" } return ( {k: v for k, v in all_keys.items() if isinstance(v, BaseChannel)}, {k: v for k, v in all_keys.items() if is_managed_value(v)}, + type_hints, ) diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index f1bb4d8ab..60cbd2547 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/libs/langgraph/tests/test_pregel_async.py @@ -27,11 +27,7 @@ from uuid import UUID import httpx import pytest from langchain_core.language_models import GenericFakeChatModel -from langchain_core.runnables import ( - RunnableConfig, - RunnableLambda, - RunnablePassthrough, -) +from langchain_core.runnables import RunnableConfig, RunnableLambda, RunnablePassthrough from langchain_core.utils.aiter import aclosing from pytest_mock import MockerFixture from syrupy import SnapshotAssertion @@ -4524,6 +4520,7 @@ async def test_nested_pydantic_models(version: str) -> None: class NestedModel(BaseModel): value: int name: str + something: Optional[str] = None # Forward reference model class RecursiveModel(BaseModel): @@ -4545,18 +4542,27 @@ async def test_nested_pydantic_models(version: str) -> None: name: str friends: list[str] = Field(default_factory=list) # IDs of friends + class MyTypedDict(TypedDict): + x: int + class State(BaseModel): # Basic nested model tests top_level: str nested: NestedModel optional_nested: Optional[NestedModel] = None dict_nested: dict[str, NestedModel] + my_set: set[int] list_nested: Annotated[ Union[dict, list[dict[str, NestedModel]]], lambda x, y: (x or []) + [y] ] + list_nested_reversed: Annotated[ + Union[list[dict[str, NestedModel]], NestedModel, dict, list], + lambda x, y: (x or []) + [y], + ] tuple_nested: tuple[str, NestedModel] tuple_list_nested: list[tuple[int, NestedModel]] complex_tuple: tuple[str, dict[str, tuple[int, NestedModel]]] + my_typed_dict: MyTypedDict # Forward reference test recursive: RecursiveModel @@ -4572,8 +4578,11 @@ async def test_nested_pydantic_models(version: str) -> None: "top_level": "initial", "nested": {"value": 42, "name": "test"}, "optional_nested": {"value": 10, "name": "optional"}, + "my_set": [1, 2, 7], + "my_typed_dict": {"x": 1}, "dict_nested": {"a": {"value": 5, "name": "a"}}, "list_nested": [{"a": {"value": 6, "name": "b"}}], + "list_nested_reversed": ["foo", "bar"], "tuple_nested": ["tuple-key", {"value": 7, "name": "tuple-value"}], "tuple_list_nested": [[1, {"value": 8, "name": "tuple-in-list"}]], "complex_tuple": [ @@ -5882,10 +5891,12 @@ async def test_store_injected_async(checkpointer_name: str, store_name: str) -> ): assert isinstance(store, BaseStore) await store.aput( - namespace - if self.i is not None - and config["configurable"]["thread_id"] in (thread_1, thread_2) - else (f"foo_{self.i}", "bar"), + ( + namespace + if self.i is not None + and config["configurable"]["thread_id"] in (thread_1, thread_2) + else (f"foo_{self.i}", "bar") + ), doc_id, { **doc,