Update tests

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

libs/langgraph/langgraph/graph/schema_utils.py

@@ -5,13 +5,14 @@ from inspect import isclass
 from typing import (
     Any,
     Callable,
+    Hashable,
     Optional,
     Type,
-    get_args,
     TypeVar,
+    Union,
+    get_args,
+    get_origin,
     get_type_hints,
-    Hashable,
-    get_origin, Union,
 )
 
 from pydantic import BaseModel, Discriminator
@@ -32,11 +33,11 @@ class SchemaCoercionMapper:
     """Lightweight coercion of *dict* → *BaseModel* instances."""
 
     def __new__(
-            cls,
-            schema: Type[Any],
-            type_hints: Optional[dict[str, Any]] = None,
-            *,
-            max_depth: int = 12,
+        cls,
+        schema: Type[Any],
+        type_hints: Optional[dict[str, Any]] = None,
+        *,
+        max_depth: int = 12,
     ) -> "SchemaCoercionMapper":
         by_depth = _cache.setdefault(schema, {})
         if max_depth in by_depth:
@@ -46,11 +47,11 @@ class SchemaCoercionMapper:
         return inst
 
     def __init__(
-            self,
-            schema: Type[Any],
-            type_hints: Optional[dict[str, Any]] = None,
-            *,
-            max_depth: int = 12,
+        self,
+        schema: Type[Any],
+        type_hints: Optional[dict[str, Any]] = None,
+        *,
+        max_depth: int = 12,
     ) -> None:
         if hasattr(self, "_initialised"):
             return
@@ -62,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}, include_extras=True)
+            else get_type_hints(
+                schema, localns={schema.__name__: schema}, include_extras=True
+            )
         )
 
         if issubclass(schema, BaseModelV1):
@@ -123,34 +126,35 @@ class SchemaCoercionMapper:
         return self._construct(**processed)
 
     def _build_coercer(
-            self, field_type: Any, depth: int, *, throw: bool = False
+        self, field_type: Any, depth: int, *, throw: bool = False
     ) -> Callable[[Any, Any], Any]:
         if depth == 0:
             return self._passthrough
 
-        # unwrap Annotated
         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
 
-        # support TypeVar
         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
 
-        # support generics like Wrapper[int] Wrapper[AnyMessage]
-        if hasattr(field_type, "__parameters__") and hasattr(field_type, "model_fields"):
+        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)}")
+                            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)
@@ -276,7 +280,11 @@ class SchemaCoercionMapper:
                         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:
+                elif (
+                    isinstance(v, dict)
+                    and isinstance(discriminator_key, str)
+                    and discriminator_key in v
+                ):
                     tag = v[discriminator_key]
 
                 if tag is not None:
@@ -287,9 +295,13 @@ class SchemaCoercionMapper:
 
                         try:
                             if issubclass(base_type, (BaseModel, BaseModelV1)):
-                                return SchemaCoercionMapper(base_type, max_depth=d).coerce(v, d)
+                                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}")
+                            logger.debug(
+                                f"Coercion with {base_type} failed for tag={tag}: {e}"
+                            )
                             continue
 
                 # fallback: try coercing each branch

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

@@ -1,9 +1,8 @@
-import pytest
-from typing import List, Dict, Set, Tuple, Optional, Union, TypeVar, Generic, Literal
+from typing import Dict, Generic, List, Literal, Optional, Set, Tuple, TypeVar, Union
 
-from langchain_core.messages import HumanMessage, AIMessage, AnyMessage
+from langchain_core.messages import AIMessage, AnyMessage, HumanMessage
+from pydantic import BaseModel, Discriminator, Field, Tag
 from typing_extensions import Annotated
-from pydantic import BaseModel, Field, Discriminator, Tag
 
 from langgraph.graph.schema_utils import SchemaCoercionMapper
 
@@ -12,17 +11,10 @@ def test_any_message():
     class MyMessage(BaseModel):
         msg: List[AnyMessage]
 
-
     data = {
         "msg": [
-            {
-                "type": "human",
-                "content": "Hello"
-            },
-            {
-                "type": "ai",
-                "content": "Hi there!"
-            }
+            {"type": "human", "content": "Hello"},
+            {"type": "ai", "content": "Hi there!"},
         ]
     }
 
@@ -36,6 +28,7 @@ def test_any_message():
     assert isinstance(result.msg[0], (HumanMessage))
     assert isinstance(result.msg[1], (AIMessage))
 
+
 # ==== 基础模型 ====
 class SimpleModel(BaseModel):
     name: str
@@ -182,29 +175,13 @@ class Warehouse(BaseModel):
 
 def test_nested_optional_generic_union():
     # Box[TaggedPet]
-    data1 = {
-        "animal": {
-            "content": {
-                "type": "cat",
-                "name": "Kitty"
-            }
-        }
-    }
+    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
-                }
-            }
-        }
-    }
+    data2 = {"cage": {"payload": {"content": {"type": "dog", "age": 8}}}}
     mapper2 = SchemaCoercionMapper(Warehouse)
     result2 = mapper2(data2)
     assert isinstance(result2.cage.payload.content, Dog)
@@ -217,4 +194,4 @@ def test_nested_optional_generic_union():
     # deeply nested Optional
     data4 = {"cage": {"payload": {"content": None}}}
     result4 = mapper2(data4)
-    assert result4.cage.payload.content is None
+    assert result4.cage.payload.content is None