From ee650ab85fd34c171b8554a2e0fd23e33eeb3d50 Mon Sep 17 00:00:00 2001 From: William Fu-Hinthorn <13333726+hinthornw@users.noreply.github.com> Date: Tue, 18 Mar 2025 09:23:26 -0700 Subject: [PATCH] Only throw if in union --- .../langgraph/langgraph/graph/schema_utils.py | 38 +++++++++++++------ libs/langgraph/tests/test_pregel_async.py | 29 +++++++++----- 2 files changed, 46 insertions(+), 21 deletions(-) diff --git a/libs/langgraph/langgraph/graph/schema_utils.py b/libs/langgraph/langgraph/graph/schema_utils.py index 4596331ad..1cbec6dd3 100644 --- a/libs/langgraph/langgraph/graph/schema_utils.py +++ b/libs/langgraph/langgraph/graph/schema_utils.py @@ -62,18 +62,22 @@ 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,16 +88,16 @@ class SchemaCoercionMapper: is_base_model = False if is_base_model: - mapper = SchemaCoercionMapper(field_type, self.max_depth) + mapper = SchemaCoercionMapper(field_type, 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, 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)): @@ -104,9 +108,17 @@ class SchemaCoercionMapper: if origin is dict or field_type is dict: args = get_args(field_type) if len(args) != 2: - return self._passthrough - k_sub = self._build_coercer(args[0]) - v_sub = self._build_coercer(args[1]) + if not throw: + return self._passthrough + + def dict_coercer(v: Any, d: Any) -> Any: + if not isinstance(v, dict): + raise TypeError("Expected dict, got %s" % type(v)) + return {k_sub(k, d - 1): v_sub(val, d - 1) for k, val in v.items()} + + 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): @@ -119,7 +131,7 @@ 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)): @@ -133,11 +145,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: diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index 39d3d34a6..a7159133b 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, 4.5], + "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,