Schema coercer should never throw (#3871)

- pydantic will do that for us if needed

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

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

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,