mirror of
https://github.com/langchain-ai/langgraph.git
synced 2026-09-08 02:37:52 +02:00
fix(langgraph): Fix deserialization issue for AnyMessage objects, and… (#4317)
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
This commit is contained in:
baii
committed by
William Fu-Hinthorn
William Fu-Hinthorn
parent
38d806733d
commit
654c5f21e4
@@ -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,
|
||||
|
||||
@@ -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
|
||||
Reference in New Issue
Block a user