diff --git a/libs/langgraph/langgraph/graph/schema_utils.py b/libs/langgraph/langgraph/graph/schema_utils.py index 7eebb5915..a58549e8b 100644 --- a/libs/langgraph/langgraph/graph/schema_utils.py +++ b/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 diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index 5bac66252..65cc758fb 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/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, diff --git a/libs/langgraph/tests/test_schema_coercion_mapper.py b/libs/langgraph/tests/test_schema_coercion_mapper.py index e23ecde59..a1414d5e4 100644 --- a/libs/langgraph/tests/test_schema_coercion_mapper.py +++ b/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 \ No newline at end of file + assert result4.cage.payload.content is None