mirror of
https://github.com/langchain-ai/langgraph.git
synced 2026-08-25 09:02:25 +02:00
Remove SchemaCoercionMapper (#4855)
This commit is contained in:
Nuno Campos
@@ -1,268 +0,0 @@
|
||||
import functools
|
||||
import logging
|
||||
import weakref
|
||||
from dataclasses import is_dataclass
|
||||
from inspect import isclass
|
||||
from typing import (
|
||||
Annotated,
|
||||
Any,
|
||||
Callable,
|
||||
Optional,
|
||||
Union,
|
||||
get_args,
|
||||
get_origin,
|
||||
get_type_hints,
|
||||
)
|
||||
|
||||
from pydantic import BaseModel, ConfigDict, TypeAdapter
|
||||
from typing_extensions import is_typeddict
|
||||
|
||||
__all__ = ["SchemaCoercionMapper"]
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
_cache: weakref.WeakKeyDictionary[type[Any], dict[int, "SchemaCoercionMapper"]] = (
|
||||
weakref.WeakKeyDictionary()
|
||||
)
|
||||
|
||||
|
||||
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,
|
||||
) -> "SchemaCoercionMapper":
|
||||
by_depth = _cache.setdefault(schema, {})
|
||||
if max_depth in by_depth:
|
||||
return by_depth[max_depth]
|
||||
inst = super().__new__(cls)
|
||||
by_depth[max_depth] = inst
|
||||
return inst
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
schema: type[BaseModel],
|
||||
type_hints: Optional[dict[str, Any]] = None,
|
||||
*,
|
||||
max_depth: int = 12,
|
||||
) -> None:
|
||||
if hasattr(self, "_initialised"):
|
||||
return
|
||||
self._initialised = True
|
||||
|
||||
self.schema = schema
|
||||
self.max_depth = max_depth
|
||||
|
||||
self.type_hints = (
|
||||
type_hints
|
||||
if type_hints is not None
|
||||
else get_type_hints(schema, localns={schema.__name__: schema})
|
||||
)
|
||||
|
||||
if issubclass(schema, BaseModel):
|
||||
self._fields = {
|
||||
n: self.type_hints.get(n, f.annotation)
|
||||
for n, f in schema.model_fields.items()
|
||||
}
|
||||
self._construct: Callable[..., Any] = schema.model_construct
|
||||
unhandled_attrs = ("validators", "field_validators", "root_validators")
|
||||
if (decorators := getattr(schema, "__pydantic_decorators__", None)) and any(
|
||||
getattr(decorators, attr, None) for attr in unhandled_attrs
|
||||
):
|
||||
self.coerce = lambda v, _: schema.model_validate(v)
|
||||
else:
|
||||
self.coerce = self._coerce
|
||||
else:
|
||||
raise TypeError("Schema must be a Pydantic V2 model.")
|
||||
|
||||
self._field_coercers: Optional[dict[str, Callable[[Any, int], Any]]] = None
|
||||
|
||||
def __call__(self, input_data: Any, depth: Optional[int] = None) -> Any:
|
||||
return self.coerce(input_data, depth)
|
||||
|
||||
def _coerce(self, input_data: Any, depth: Optional[int] = None) -> Any:
|
||||
if depth is None:
|
||||
depth = self.max_depth
|
||||
if not isinstance(input_data, dict) or depth <= 0:
|
||||
return input_data
|
||||
|
||||
if self._field_coercers is None:
|
||||
self._field_coercers = {
|
||||
n: self._build_coercer(t, depth - 1) for n, t in self._fields.items()
|
||||
}
|
||||
|
||||
processed: dict[str, Any] = {}
|
||||
for k, v in input_data.items():
|
||||
fn = self._field_coercers.get(k)
|
||||
processed[k] = fn(v, depth - 1) if fn else v
|
||||
return self._construct(**processed)
|
||||
|
||||
def _build_coercer(
|
||||
self, field_type: Any, depth: int, *, throw: bool = False
|
||||
) -> Callable[[Any, Any], Any]:
|
||||
if depth == 0:
|
||||
return self._passthrough
|
||||
|
||||
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 isclass(field_type):
|
||||
# This is needed bcs. of issubclass issues on older versions of python
|
||||
try:
|
||||
is_bm_subclass = issubclass(field_type, BaseModel)
|
||||
except TypeError:
|
||||
# python < 3.11 issue.
|
||||
is_bm_subclass = False
|
||||
if is_bm_subclass:
|
||||
mapper = SchemaCoercionMapper(field_type, max_depth=depth - 1)
|
||||
return lambda v, d: mapper.coerce(v, d) if isinstance(v, dict) else v
|
||||
|
||||
if origin is list:
|
||||
args = get_args(field_type)
|
||||
if len(args) != 1:
|
||||
return self._passthrough
|
||||
sub = self._build_coercer(args[0], depth - 1)
|
||||
|
||||
def list_coercer(v: Any, d: Any) -> Any:
|
||||
if not isinstance(v, (list, tuple)):
|
||||
return v
|
||||
return [sub(x, d - 1) for x in v]
|
||||
|
||||
return list_coercer
|
||||
|
||||
if origin is set or field_type is set:
|
||||
args = get_args(field_type)
|
||||
if len(args) > 1:
|
||||
return self._passthrough
|
||||
elif len(args) == 1:
|
||||
sub = self._build_coercer(args[0], depth - 1)
|
||||
else:
|
||||
sub = None # type: ignore
|
||||
|
||||
def set_coercer(v: Any, d: Any) -> Any:
|
||||
if not isinstance(v, (list, tuple, set)):
|
||||
return v
|
||||
if sub is None:
|
||||
return set(v)
|
||||
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:
|
||||
|
||||
def dict_coercer(v: Any, d: Any) -> Any:
|
||||
if not isinstance(v, dict):
|
||||
if throw:
|
||||
raise TypeError(f"Expected dict, got {type(v)}")
|
||||
return v
|
||||
|
||||
return dict_coercer
|
||||
k_sub = self._build_coercer(args[0], depth - 1)
|
||||
v_sub = self._build_coercer(args[1], depth - 1)
|
||||
|
||||
def dict_coercer(v: Any, d: Any) -> Any:
|
||||
if not isinstance(v, dict):
|
||||
if throw:
|
||||
raise TypeError(f"Expected dict, got {type(v)}")
|
||||
return v
|
||||
return {k_sub(k, d - 1): v_sub(val, d - 1) for k, val in v.items()}
|
||||
|
||||
return dict_coercer
|
||||
|
||||
if origin is tuple:
|
||||
elem_types = get_args(field_type)
|
||||
if not elem_types:
|
||||
return self._passthrough
|
||||
subs = [self._build_coercer(t, depth - 1) for t in elem_types]
|
||||
return lambda v, d: (
|
||||
tuple(
|
||||
subs[i](v[i] if i < len(v) else None, d - 1)
|
||||
for i in range(len(subs))
|
||||
)
|
||||
if isinstance(v, (list, tuple))
|
||||
else v
|
||||
)
|
||||
|
||||
if origin is Union:
|
||||
uargs = get_args(field_type)
|
||||
subs, none_in_union = [], False
|
||||
for ix, arg in enumerate(uargs):
|
||||
if arg is type(None):
|
||||
none_in_union = True
|
||||
else:
|
||||
subs.append(
|
||||
self._build_coercer(arg, depth - 1, throw=ix < len(uargs) - 1)
|
||||
)
|
||||
|
||||
def union_coercer(v: Any, d: Any) -> Any:
|
||||
if v is None and none_in_union:
|
||||
return None
|
||||
err = None
|
||||
for sp in subs:
|
||||
try:
|
||||
return sp(v, d - 1)
|
||||
except TypeError as e:
|
||||
err = e
|
||||
if err:
|
||||
raise err
|
||||
return v
|
||||
|
||||
return union_coercer
|
||||
|
||||
adapter_fn = _get_adapter(field_type)
|
||||
return lambda v, _d: adapter_fn(v)
|
||||
|
||||
@staticmethod
|
||||
def _passthrough(v: Any, _d: Any) -> Any: # noqa: D401
|
||||
return v
|
||||
|
||||
|
||||
_adapter_cache: dict[Any, Callable[[Any], Any]] = {}
|
||||
|
||||
|
||||
_IDENTITY_TYPES: tuple[type[Any], ...] = (
|
||||
int,
|
||||
float,
|
||||
str,
|
||||
bool,
|
||||
bytes,
|
||||
bytearray,
|
||||
complex,
|
||||
memoryview,
|
||||
type(None),
|
||||
)
|
||||
|
||||
|
||||
@functools.lru_cache(maxsize=2048)
|
||||
def _adapter_for(tp: Any) -> Callable[[Any], Any]: # noqa: D401
|
||||
try:
|
||||
config = (
|
||||
None
|
||||
if (issubclass(tp, BaseModel) or is_dataclass(tp) or is_typeddict(tp))
|
||||
else ConfigDict(arbitrary_types_allowed=True)
|
||||
)
|
||||
except TypeError:
|
||||
config = None
|
||||
return TypeAdapter(tp, config=config).validate_python
|
||||
|
||||
|
||||
def _get_adapter(tp: Any) -> Callable[[Any], Any]:
|
||||
try:
|
||||
return _adapter_cache[tp]
|
||||
except KeyError:
|
||||
fn = _adapter_for(tp)
|
||||
_adapter_cache[tp] = fn
|
||||
return fn
|
||||
@@ -64,7 +64,6 @@ from langgraph.graph.graph import (
|
||||
Graph,
|
||||
Send,
|
||||
)
|
||||
from langgraph.graph.schema_utils import SchemaCoercionMapper
|
||||
from langgraph.managed.base import (
|
||||
ChannelKeyPlaceholder,
|
||||
ChannelTypePlaceholder,
|
||||
@@ -1042,11 +1041,8 @@ def _pick_mapper(
|
||||
) -> Optional[Callable[[Any], Any]]:
|
||||
if state_keys == ["__root__"]:
|
||||
return None
|
||||
if isclass(schema):
|
||||
if issubclass(schema, dict):
|
||||
return None
|
||||
if issubclass(schema, BaseModel):
|
||||
return SchemaCoercionMapper(schema, type_hints=type_hints)
|
||||
if isclass(schema) and issubclass(schema, dict):
|
||||
return None
|
||||
return partial(_coerce_state, schema)
|
||||
|
||||
|
||||
|
||||
Reference in New Issue
Block a user