_evaluate() forward ref

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.")
@@ -75,6 +98,7 @@ class SchemaCoercionMapper:
         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, depth - 1)
@@ -88,10 +112,10 @@ class SchemaCoercionMapper:
                 is_base_model = False
 
             if is_base_model:
-                mapper = SchemaCoercionMapper(field_type, depth - 1)
+                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, depth - 1)
+                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)
@@ -105,16 +129,27 @@ 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)
+
+            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:
-                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()}
+                        if throw:
+                            raise TypeError("Expected dict, got %s" % type(v))
+                    return v
 
                 return dict_coercer
             k_sub = self._build_coercer(args[0], depth - 1)
@@ -122,6 +157,8 @@ class SchemaCoercionMapper:
 
             def dict_coercer(v: Any, d: Any) -> Any:
                 if not isinstance(v, dict):
+                    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()}
 
@@ -160,7 +197,7 @@ 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

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:
@@ -827,7 +829,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(
@@ -942,12 +948,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
@@ -955,7 +961,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)
 
 
@@ -1017,18 +1023,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

@@ -4578,7 +4578,7 @@ 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_set": [1, 2, 7],
         "my_typed_dict": {"x": 1},
         "dict_nested": {"a": {"value": 5, "name": "a"}},
         "list_nested": [{"a": {"value": 6, "name": "b"}}],