Update tests

Signed-off-by: William Fu-Hinthorn <13333726+hinthornw@users.noreply.github.com>

libs/langgraph/langgraph/graph/schema_utils.py

@@ -5,23 +5,25 @@ from inspect import isclass
 from typing import (
     Any,
     Callable,
+    Hashable,
     Optional,
     Type,
+    TypeVar,
     Union,
     get_args,
     get_origin,
     get_type_hints,
 )
 
-from pydantic import BaseModel
+from pydantic import BaseModel, Discriminator
+from pydantic.fields import FieldInfo
 from pydantic.v1 import BaseModel as BaseModelV1
-from typing_extensions import Annotated
+from typing_extensions import Annotated, Literal
 
 __all__ = ["SchemaCoercionMapper"]
 
 logger = logging.getLogger(__name__)
 
-
 _cache: weakref.WeakKeyDictionary[Type[Any], dict[int, "SchemaCoercionMapper"]] = (
     weakref.WeakKeyDictionary()
 )
@@ -61,7 +63,9 @@ class SchemaCoercionMapper:
         self.type_hints = (
             type_hints
             if type_hints is not None
-            else get_type_hints(schema, localns={schema.__name__: schema})
+            else get_type_hints(
+                schema, localns={schema.__name__: schema}, include_extras=True
+            )
         )
 
         if issubclass(schema, BaseModelV1):
@@ -127,15 +131,38 @@ class SchemaCoercionMapper:
         if depth == 0:
             return self._passthrough
 
+        field_type, metadata = self._unwrap_annotated(field_type)
         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 isinstance(field_type, TypeVar):
+            concrete = self.type_hints.get(field_type)  # type: ignore
+            if concrete is not None:
+                return self._build_coercer(concrete, depth - 1)
+            return self._passthrough
+
+        if hasattr(field_type, "__parameters__") and hasattr(
+            field_type, "model_fields"
+        ):
+            try:
+                type_hints = self.resolve_concrete_type_hints(field_type)
+
+                def generic_model_coercer(v: Any, d: int) -> Any:
+                    if not isinstance(v, dict):
+                        if throw:
+                            raise TypeError(
+                                f"Expected dict for {field_type}, got {type(v)}"
+                            )
+                        return v
+                    mapper = SchemaCoercionMapper(field_type, type_hints, max_depth=d)
+                    return mapper.coerce(v, d)
+
+                return generic_model_coercer
+            except Exception as e:
+                logger.debug(f"Generic type resolution failed: {e}")
+                return self._passthrough
 
         if isclass(field_type):
             # This is needed bcs. of issubclass issues on older versions of python
@@ -180,6 +207,7 @@ class SchemaCoercionMapper:
                 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:
@@ -218,27 +246,75 @@ class SchemaCoercionMapper:
             )
 
         if origin is Union:
-            uargs = get_args(field_type)
-            subs, none_in_union = [], False
-            for ix, arg in enumerate(uargs):
+            args = get_args(field_type)
+            discriminator_key = self._extract_discriminator_key(metadata)
+            none_in_union = False
+            discriminator_map = {}
+
+            for arg in args:
                 if arg is type(None):
                     none_in_union = True
-                else:
-                    subs.append(
-                        self._build_coercer(arg, depth - 1, throw=ix < len(uargs) - 1)
-                    )
+                    continue
+                base_type = arg
+                if get_origin(arg) is Annotated:
+                    base_type, _ = get_args(arg)[0], get_args(arg)[1:]
+                try:
+                    hint = get_type_hints(base_type)
+                    lit = hint.get(discriminator_key)
+                    if get_origin(lit) is Literal:
+                        for val in get_args(lit):
+                            discriminator_map[val] = base_type
+                except Exception as e:
+                    if throw:
+                        raise e
+                    else:
+                        logger.debug(f"Failed to extract discriminator: {e}")
 
             def union_coercer(v: Any, d: Any) -> Any:
                 if v is None and none_in_union:
                     return None
-                err = None
-                for sp in subs:
+
+                tag = None
+                if callable(discriminator_key):
                     try:
-                        return sp(v, d - 1)
-                    except TypeError as e:
-                        err = e
-                if err:
-                    raise err
+                        tag = discriminator_key(v)
+                    except Exception as e:
+                        logger.debug(f"Failed to call discriminator func: {e}")
+                elif (
+                    isinstance(v, dict)
+                    and isinstance(discriminator_key, str)
+                    and discriminator_key in v
+                ):
+                    tag = v[discriminator_key]
+
+                if tag is not None:
+                    for arg in args:
+                        base_type = arg
+                        if get_origin(arg) is Annotated:
+                            base_type, _ = get_args(arg)[0], get_args(arg)[1:]
+
+                        try:
+                            if issubclass(base_type, (BaseModel, BaseModelV1)):
+                                return SchemaCoercionMapper(
+                                    base_type, max_depth=d
+                                ).coerce(v, d)
+                        except Exception as e:
+                            logger.debug(
+                                f"Coercion with {base_type} failed for tag={tag}: {e}"
+                            )
+                            continue
+
+                # fallback: try coercing each branch
+                for arg in args:
+                    try:
+                        sub = self._build_coercer(arg, d - 1)
+                        return sub(v, d - 1)
+                    except Exception as e:
+                        if throw:
+                            raise e
+                        else:
+                            logger.debug(f"Fallback coercion failed for arg={arg}: {e}")
+
                 return v
 
             return union_coercer
@@ -250,10 +326,58 @@ class SchemaCoercionMapper:
     def _passthrough(v: Any, _d: Any) -> Any:  # noqa: D401
         return v
 
+    @staticmethod
+    def _extract_discriminator_key(meta: list[Any]) -> str | Callable[[Any], Hashable]:
+        """Extract discriminator field name or function from Annotated metadata"""
+        for m in meta:
+            if isinstance(m, FieldInfo):
+                disc = getattr(m, "discriminator", None)
+                if isinstance(disc, Discriminator):
+                    return disc.discriminator
+                elif isinstance(disc, str):
+                    return disc
+        return "type"
+
+    @staticmethod
+    def _unwrap_annotated(tp: Any) -> tuple[Any, list[Any]]:
+        """Unwrap nested Annotated types, extracting the base type and all metadata"""
+        metadata = []
+        while get_origin(tp) is Annotated:
+            tp, *meta = get_args(tp)
+            metadata.extend(meta)
+        return tp, metadata
+
+    @staticmethod
+    def resolve_concrete_type_hints(generic_model_type: Any) -> dict[Any, Any]:
+        """Resolve concrete type hints in a generic model"""
+        origin = get_origin(generic_model_type)
+        args = get_args(generic_model_type)
+        param_names = getattr(origin, "__parameters__", [])
+
+        if not args or not param_names:
+            return {}
+
+        type_map = dict(zip(param_names, args))
+        result = {}
+
+        for field_name, model_field in origin.model_fields.items():
+            anno = model_field.annotation
+            if get_origin(anno) is Annotated:
+                base, *meta = get_args(anno)
+                if isinstance(base, TypeVar) and base in type_map:
+                    result[field_name] = Annotated[type_map[base], *meta]
+                else:
+                    result[field_name] = anno
+            elif isinstance(anno, TypeVar) and anno in type_map:
+                result[field_name] = type_map[anno]
+            else:
+                result[field_name] = anno
+
+        return result
+
 
 _adapter_cache: dict[Any, Callable[[Any], Any]] = {}
 
-
 _IDENTITY_TYPES: tuple[type[Any], ...] = (
     int,
     float,

libs/langgraph/tests/test_pregel.py

@@ -3123,8 +3123,10 @@ def test_nested_pydantic_models(version: str) -> None:
         from pydantic import (  # type: ignore
             BaseModel,
             ByteSize,
+            Discriminator,
             Field,
             SecretStr,
+            Tag,
             confloat,
             conint,
             conlist,
@@ -3173,6 +3175,16 @@ def test_nested_pydantic_models(version: str) -> None:
         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)
+    if version == "v2":
+        FuncDiscriminatorPet = Annotated[
+            Union[
+                Annotated[Dog, Tag(tag="dog")],
+                Annotated[Cat, Tag(tag="cat")],
+            ],
+            Field(discriminator=Discriminator(lambda obj: obj.get("pet_type"))),
+        ]
+    else:
+        FuncDiscriminatorPet = Union[Dog, Cat]
 
     class State(BaseModel):
         # Basic nested model tests
@@ -3212,6 +3224,7 @@ def test_nested_pydantic_models(version: str) -> None:
         pattern: re.Pattern
         secret: SecretStr
         file_size: ByteSize
+        discriminated_pet: FuncDiscriminatorPet
 
         # Constrained types
         positive_value: PositiveInt
@@ -3283,6 +3296,7 @@ def test_nested_pydantic_models(version: str) -> None:
         "pattern": "^test$",
         "secret": "password123",
         "file_size": 1024,
+        "discriminated_pet": {"pet_type": "cat", "meow": "indubitably"},
         # Constrained types
         "positive_value": 42,
         "non_negative": 0.0,

libs/langgraph/tests/test_schema_coercion_mapper.py

@@ -0,0 +1,197 @@
+from typing import Dict, Generic, List, Literal, Optional, Set, Tuple, TypeVar, Union
+
+from langchain_core.messages import AIMessage, AnyMessage, HumanMessage
+from pydantic import BaseModel, Discriminator, Field, Tag
+from typing_extensions import Annotated
+
+from langgraph.graph.schema_utils import SchemaCoercionMapper
+
+
+def test_any_message():
+    class MyMessage(BaseModel):
+        msg: List[AnyMessage]
+
+    data = {
+        "msg": [
+            {"type": "human", "content": "Hello"},
+            {"type": "ai", "content": "Hi there!"},
+        ]
+    }
+
+    MyMessage.model_validate(data)
+
+    mapper = SchemaCoercionMapper(MyMessage)
+    result = mapper(data)
+    assert isinstance(result, MyMessage)
+    assert isinstance(result.msg, list)
+    assert len(result.msg) == 2
+    assert isinstance(result.msg[0], (HumanMessage))
+    assert isinstance(result.msg[1], (AIMessage))
+
+
+# ==== 基础模型 ====
+class SimpleModel(BaseModel):
+    name: str
+    age: int
+
+
+def test_simple_model():
+    data = {"name": "Alice", "age": 30}
+    mapper = SchemaCoercionMapper(SimpleModel)
+    result = mapper(data)
+    assert isinstance(result, SimpleModel)
+    assert result.name == "Alice"
+    assert result.age == 30
+
+
+# ==== 容器类型 ====
+class ContainerModel(BaseModel):
+    items: List[int]
+    mapping: Dict[str, float]
+    tags: Set[str]
+    coords: Tuple[int, int]
+
+
+def test_container_model():
+    data = {
+        "items": [1, 2, 3],
+        "mapping": {"a": 1.1},
+        "tags": ["x", "y"],
+        "coords": [10, 20],
+    }
+    mapper = SchemaCoercionMapper(ContainerModel)
+    result = mapper(data)
+    assert isinstance(result.items, list)
+    assert isinstance(result.mapping, dict)
+    assert isinstance(result.tags, set)
+    assert isinstance(result.coords, tuple)
+
+
+# ==== 泛型 ====
+T = TypeVar("T")
+
+
+class Wrapper(BaseModel, Generic[T]):
+    value: T
+
+
+def test_generic_model():
+    class IntWrapper(Wrapper[int]):
+        pass
+
+    data = {"value": 123}
+    mapper = SchemaCoercionMapper(IntWrapper)
+    result = mapper(data)
+    assert result.value == 123
+
+
+# ==== Union 类型 ====
+class Dog(BaseModel):
+    type: Literal["dog"]
+    age: int
+
+
+class Cat(BaseModel):
+    type: Literal["cat"]
+    name: str
+
+
+Pet = Union[Dog, Cat]
+
+
+class Owner(BaseModel):
+    pet: Pet
+
+
+def test_union_type():
+    data = {"pet": {"type": "dog", "age": 5}}
+    mapper = SchemaCoercionMapper(Owner)
+    result = mapper(data)
+    assert isinstance(result.pet, Dog)
+
+
+# ==== Annotated + Tag + discriminator ====
+TaggedPet = Annotated[
+    Union[
+        Annotated[Dog, Tag(tag="dog")],
+        Annotated[Cat, Tag(tag="cat")],
+    ],
+    Field(discriminator="type"),
+]
+
+
+class TaggedOwner(BaseModel):
+    pet: TaggedPet
+
+
+def test_tagged_union():
+    data = {"pet": {"type": "cat", "name": "Mimi"}}
+    mapper = SchemaCoercionMapper(TaggedOwner)
+    result = mapper(data)
+    assert isinstance(result.pet, Cat)
+
+
+# ==== Annotated + Field(discriminator=Discriminator(func)) ====
+def _get_type(obj):
+    return obj.get("type")
+
+
+FuncDiscriminatorPet = Annotated[
+    Union[
+        Annotated[Dog, Tag(tag="dog")],
+        Annotated[Cat, Tag(tag="cat")],
+    ],
+    Field(discriminator=Discriminator(_get_type)),
+]
+
+
+class FuncOwner(BaseModel):
+    pet: FuncDiscriminatorPet
+
+
+def test_func_discriminator():
+    data = {"pet": {"type": "dog", "age": 9}}
+    mapper = SchemaCoercionMapper(FuncOwner)
+    result = mapper(data)
+    assert isinstance(result.pet, Dog)
+
+
+# ==== Optional + 泛型 + 多态嵌套 ====
+class Box(BaseModel, Generic[T]):
+    content: Optional[T]
+
+
+class Crate(BaseModel, Generic[T]):
+    payload: Box[T]
+
+
+class Zoo(BaseModel):
+    animal: Box[TaggedPet]
+
+
+class Warehouse(BaseModel):
+    cage: Crate[TaggedPet]
+
+
+def test_nested_optional_generic_union():
+    # Box[TaggedPet]
+    data1 = {"animal": {"content": {"type": "cat", "name": "Kitty"}}}
+    mapper1 = SchemaCoercionMapper(Zoo)
+    result1 = mapper1(data1)
+    assert isinstance(result1.animal.content, Cat)
+
+    # Crate[TaggedPet]
+    data2 = {"cage": {"payload": {"content": {"type": "dog", "age": 8}}}}
+    mapper2 = SchemaCoercionMapper(Warehouse)
+    result2 = mapper2(data2)
+    assert isinstance(result2.cage.payload.content, Dog)
+
+    # Optional None
+    data3 = {"animal": {"content": None}}
+    result3 = mapper1(data3)
+    assert result3.animal.content is None
+
+    # deeply nested Optional
+    data4 = {"cage": {"payload": {"content": None}}}
+    result4 = mapper2(data4)
+    assert result4.cage.payload.content is None