From 654c5f21e43357f04f2de2a5ad686c982771da4b Mon Sep 17 00:00:00 2001 From: baii <33405932+littlebai3618@users.noreply.github.com> Date: Tue, 22 Apr 2025 03:50:50 +0800 Subject: [PATCH] =?UTF-8?q?fix(langgraph):=20Fix=20deserialization=20issue?= =?UTF-8?q?=20for=20AnyMessage=20objects,=20and=E2=80=A6=20(#4317)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Fix deserialization issue for AnyMessage objects, and add support for deserializing Pydantic generic and polymorphic models. bug detail: https://github.com/langchain-ai/langgraph/issues/4316 --- .../langgraph/langgraph/graph/schema_utils.py | 184 ++++++++++++--- .../tests/test_schema_coercion_mapper.py | 220 ++++++++++++++++++ 2 files changed, 368 insertions(+), 36 deletions(-) create mode 100644 libs/langgraph/tests/test_schema_coercion_mapper.py diff --git a/libs/langgraph/langgraph/graph/schema_utils.py b/libs/langgraph/langgraph/graph/schema_utils.py index c1e6eae5e..7eebb5915 100644 --- a/libs/langgraph/langgraph/graph/schema_utils.py +++ b/libs/langgraph/langgraph/graph/schema_utils.py @@ -7,21 +7,22 @@ from typing import ( Callable, Optional, Type, - Union, get_args, - get_origin, + TypeVar, get_type_hints, + Hashable, + get_origin, Union, ) -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() ) @@ -31,11 +32,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: @@ -45,11 +46,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 @@ -61,7 +62,7 @@ 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): @@ -122,20 +123,42 @@ 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 - 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) + # 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"): + 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 +203,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 +242,67 @@ 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 +314,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, diff --git a/libs/langgraph/tests/test_schema_coercion_mapper.py b/libs/langgraph/tests/test_schema_coercion_mapper.py new file mode 100644 index 000000000..e23ecde59 --- /dev/null +++ b/libs/langgraph/tests/test_schema_coercion_mapper.py @@ -0,0 +1,220 @@ +import pytest +from typing import List, Dict, Set, Tuple, Optional, Union, TypeVar, Generic, Literal + +from langchain_core.messages import HumanMessage, AIMessage, AnyMessage +from typing_extensions import Annotated +from pydantic import BaseModel, Field, Discriminator, Tag + +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 \ No newline at end of file