mirror of
https://github.com/langchain-ai/langgraph.git
synced 2026-09-09 11:17:53 +02:00
Merge commit from fork
* Patch * Add more tests * update idempotency tests --------- Co-authored-by: William Fu-Hinthorn <13333726+hinthornw@users.noreply.github.com>
This commit is contained in:
Eugene Yurtsev
co-authored by
William Fu-Hinthorn
William Fu-Hinthorn
parent
c4a4a46473
commit
50df7d423a
Generated
+1
@@ -280,6 +280,7 @@ dev = [
|
||||
{ name = "numpy" },
|
||||
{ name = "pandas" },
|
||||
{ name = "pandas-stubs", specifier = ">=2.2.2.240807" },
|
||||
{ name = "pycryptodome", specifier = ">=3.23.0" },
|
||||
{ name = "pytest" },
|
||||
{ name = "pytest-asyncio" },
|
||||
{ name = "pytest-mock" },
|
||||
|
||||
Generated
+1
@@ -289,6 +289,7 @@ dev = [
|
||||
{ name = "numpy" },
|
||||
{ name = "pandas" },
|
||||
{ name = "pandas-stubs", specifier = ">=2.2.2.240807" },
|
||||
{ name = "pycryptodome", specifier = ">=3.23.0" },
|
||||
{ name = "pytest" },
|
||||
{ name = "pytest-asyncio" },
|
||||
{ name = "pytest-mock" },
|
||||
|
||||
@@ -37,4 +37,4 @@ type:
|
||||
|
||||
format format_diff:
|
||||
uv run ruff format $(PYTHON_FILES)
|
||||
uv run ruff check --select I --fix $(PYTHON_FILES)
|
||||
uv run ruff check --fix $(PYTHON_FILES)
|
||||
|
||||
@@ -1,6 +1,8 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import AsyncIterator, Iterator, Mapping, Sequence
|
||||
import copy
|
||||
import logging
|
||||
from collections.abc import AsyncIterator, Collection, Iterator, Mapping, Sequence
|
||||
from typing import ( # noqa: UP035
|
||||
Any,
|
||||
Generic,
|
||||
@@ -14,6 +16,7 @@ from langchain_core.runnables import RunnableConfig
|
||||
|
||||
from langgraph.checkpoint.base.id import uuid6
|
||||
from langgraph.checkpoint.serde.base import SerializerProtocol, maybe_add_typed_methods
|
||||
from langgraph.checkpoint.serde.encrypted import EncryptedSerializer
|
||||
from langgraph.checkpoint.serde.jsonplus import JsonPlusSerializer
|
||||
from langgraph.checkpoint.serde.types import (
|
||||
ERROR,
|
||||
@@ -25,6 +28,7 @@ from langgraph.checkpoint.serde.types import (
|
||||
|
||||
V = TypeVar("V", int, float, str)
|
||||
PendingWrite = tuple[str, str, Any]
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
# Marked as total=False to allow for future expansion.
|
||||
@@ -474,6 +478,37 @@ class BaseCheckpointSaver(Generic[V]):
|
||||
else:
|
||||
return current + 1
|
||||
|
||||
def with_allowlist(
|
||||
self, extra_allowlist: Collection[tuple[str, ...]]
|
||||
) -> BaseCheckpointSaver[V]:
|
||||
"""Return a shallow clone with a derived msgpack allowlist."""
|
||||
serde = _with_msgpack_allowlist(self.serde, extra_allowlist)
|
||||
if serde is self.serde:
|
||||
return self
|
||||
clone = copy.copy(self)
|
||||
clone.serde = maybe_add_typed_methods(serde)
|
||||
return clone
|
||||
|
||||
|
||||
def _with_msgpack_allowlist(
|
||||
serde: SerializerProtocol, extra_allowlist: Collection[tuple[str, ...]]
|
||||
) -> SerializerProtocol:
|
||||
if isinstance(serde, JsonPlusSerializer):
|
||||
return serde.with_msgpack_allowlist(extra_allowlist)
|
||||
if isinstance(serde, EncryptedSerializer):
|
||||
inner = serde.serde
|
||||
if isinstance(inner, JsonPlusSerializer):
|
||||
updated_inner = inner.with_msgpack_allowlist(extra_allowlist)
|
||||
if updated_inner is inner:
|
||||
return serde
|
||||
return EncryptedSerializer(serde.cipher, updated_inner)
|
||||
logger.warning(
|
||||
"Serializer %s does not support msgpack allowlist. "
|
||||
"Strict msgpack deserialization will not be enforced.",
|
||||
type(serde).__name__,
|
||||
)
|
||||
return serde
|
||||
|
||||
|
||||
class EmptyChannelError(Exception):
|
||||
"""Raised when attempting to get the value of a channel that hasn't been updated
|
||||
|
||||
@@ -0,0 +1,71 @@
|
||||
import os
|
||||
from collections.abc import Iterable
|
||||
from typing import cast
|
||||
|
||||
STRICT_MSGPACK_ENABLED = os.getenv("LANGGRAPH_STRICT_MSGPACK", "false").lower() in (
|
||||
"1",
|
||||
"true",
|
||||
"yes",
|
||||
)
|
||||
|
||||
|
||||
_SENTINEL = cast(None, object())
|
||||
|
||||
SAFE_MSGPACK_TYPES: frozenset[tuple[str, ...]] = frozenset(
|
||||
{
|
||||
# datetime types
|
||||
("datetime", "datetime"),
|
||||
("datetime", "date"),
|
||||
("datetime", "time"),
|
||||
("datetime", "timedelta"),
|
||||
("datetime", "timezone"),
|
||||
# uuid
|
||||
("uuid", "UUID"),
|
||||
# numeric
|
||||
("decimal", "Decimal"),
|
||||
# collections
|
||||
("builtins", "set"),
|
||||
("builtins", "frozenset"),
|
||||
("collections", "deque"),
|
||||
# ip addresses
|
||||
("ipaddress", "IPv4Address"),
|
||||
("ipaddress", "IPv4Interface"),
|
||||
("ipaddress", "IPv4Network"),
|
||||
("ipaddress", "IPv6Address"),
|
||||
("ipaddress", "IPv6Interface"),
|
||||
("ipaddress", "IPv6Network"),
|
||||
# pathlib
|
||||
("pathlib", "Path"),
|
||||
("pathlib", "PosixPath"),
|
||||
("pathlib", "WindowsPath"),
|
||||
# pathlib in Python 3.13+
|
||||
("pathlib._local", "Path"),
|
||||
("pathlib._local", "PosixPath"),
|
||||
("pathlib._local", "WindowsPath"),
|
||||
# zoneinfo
|
||||
("zoneinfo", "ZoneInfo"),
|
||||
# regex
|
||||
("re", "compile"),
|
||||
# langgraph
|
||||
("langgraph.types", "Send"),
|
||||
("langgraph.types", "Interrupt"),
|
||||
("langgraph.types", "Command"),
|
||||
("langgraph.types", "StateSnapshot"),
|
||||
("langgraph.types", "PregelTask"),
|
||||
("langgraph.types", "Overwrite"),
|
||||
("langgraph.store.base", "Item"),
|
||||
("langgraph.store.base", "GetOp"),
|
||||
}
|
||||
)
|
||||
|
||||
# Allowed (module, name, method) triples for EXT_METHOD_SINGLE_ARG.
|
||||
# Only these specific method invocations are permitted during deserialization.
|
||||
# This is separate from SAFE_MSGPACK_TYPES which only governs construction.
|
||||
SAFE_MSGPACK_METHODS: frozenset[tuple[str, str, str]] = frozenset(
|
||||
{
|
||||
("datetime", "datetime", "fromisoformat"),
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
AllowedMsgpackModules = Iterable[tuple[str, ...] | type]
|
||||
@@ -41,7 +41,7 @@ class EncryptedSerializer(SerializerProtocol):
|
||||
) -> "EncryptedSerializer":
|
||||
"""Create an `EncryptedSerializer` using AES encryption."""
|
||||
try:
|
||||
from Crypto.Cipher import AES # type: ignore
|
||||
from Crypto.Cipher import AES
|
||||
except ImportError:
|
||||
raise ImportError(
|
||||
"Pycryptodome is not installed. Please install it with `pip install pycryptodome`."
|
||||
|
||||
@@ -10,7 +10,7 @@ import pickle
|
||||
import re
|
||||
import sys
|
||||
from collections import deque
|
||||
from collections.abc import Callable, Sequence
|
||||
from collections.abc import Callable, Iterable, Sequence
|
||||
from datetime import date, datetime, time, timedelta, timezone
|
||||
from enum import Enum
|
||||
from inspect import isclass
|
||||
@@ -22,17 +22,24 @@ from ipaddress import (
|
||||
IPv6Interface,
|
||||
IPv6Network,
|
||||
)
|
||||
from typing import Any, Literal
|
||||
from typing import TYPE_CHECKING, Any, Literal, cast
|
||||
from uuid import UUID
|
||||
from zoneinfo import ZoneInfo
|
||||
|
||||
import ormsgpack
|
||||
from langchain_core.load.load import Reviver
|
||||
|
||||
from langgraph.checkpoint.serde import _msgpack as _lg_msgpack
|
||||
from langgraph.checkpoint.serde.base import SerializerProtocol
|
||||
from langgraph.checkpoint.serde.types import SendProtocol
|
||||
from langgraph.store.base import Item
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from langgraph.checkpoint.serde._msgpack import (
|
||||
AllowedMsgpackModules,
|
||||
)
|
||||
from langgraph.checkpoint.serde.types import SendProtocol
|
||||
|
||||
LC_REVIVER = Reviver()
|
||||
EMPTY_BYTES = b""
|
||||
logger = logging.getLogger(__name__)
|
||||
@@ -53,19 +60,59 @@ class JsonPlusSerializer(SerializerProtocol):
|
||||
self,
|
||||
*,
|
||||
pickle_fallback: bool = False,
|
||||
allowed_json_modules: Sequence[tuple[str, ...]] | Literal[True] | None = None,
|
||||
allowed_json_modules: Iterable[tuple[str, ...]] | Literal[True] | None = None,
|
||||
allowed_msgpack_modules: (
|
||||
AllowedMsgpackModules | Literal[True] | None
|
||||
) = _lg_msgpack._SENTINEL,
|
||||
__unpack_ext_hook__: Callable[[int, bytes], Any] | None = None,
|
||||
) -> None:
|
||||
if allowed_msgpack_modules is _lg_msgpack._SENTINEL:
|
||||
if _lg_msgpack.STRICT_MSGPACK_ENABLED:
|
||||
allowed_msgpack_modules = None
|
||||
else:
|
||||
allowed_msgpack_modules = True
|
||||
self.pickle_fallback = pickle_fallback
|
||||
self._allowed_modules = (
|
||||
{mod_and_name for mod_and_name in allowed_json_modules}
|
||||
if allowed_json_modules and allowed_json_modules is not True
|
||||
else (allowed_json_modules if allowed_json_modules is True else None)
|
||||
self._allowed_json_modules: set[tuple[str, ...]] | Literal[True] | None = (
|
||||
_normalize_allowlist(allowed_json_modules)
|
||||
)
|
||||
self._allowed_msgpack_modules = _normalize_allowlist(allowed_msgpack_modules)
|
||||
|
||||
self._custom_unpack_ext_hook = __unpack_ext_hook__ is not None
|
||||
self._unpack_ext_hook = (
|
||||
__unpack_ext_hook__
|
||||
if __unpack_ext_hook__ is not None
|
||||
else _msgpack_ext_hook
|
||||
else _create_msgpack_ext_hook(self._allowed_msgpack_modules)
|
||||
)
|
||||
|
||||
def with_msgpack_allowlist(
|
||||
self, extra_allowlist: Iterable[tuple[str, ...] | type]
|
||||
) -> JsonPlusSerializer:
|
||||
"""Return a new serializer with a merged msgpack allowlist."""
|
||||
base_allowlist = self._allowed_msgpack_modules
|
||||
if base_allowlist is True or base_allowlist is False:
|
||||
return self
|
||||
elif base_allowlist:
|
||||
base_allowlist = set(base_allowlist)
|
||||
else:
|
||||
base_allowlist = set()
|
||||
extra = _normalize_module_keys(tuple(extra_allowlist))
|
||||
merged = base_allowlist | extra
|
||||
if merged == base_allowlist:
|
||||
return self
|
||||
allowed_msgpack_modules: AllowedMsgpackModules | Literal[True] | None
|
||||
if merged:
|
||||
allowed_msgpack_modules = tuple(merged)
|
||||
elif isinstance(self._allowed_msgpack_modules, set):
|
||||
allowed_msgpack_modules = tuple(self._allowed_msgpack_modules)
|
||||
else:
|
||||
allowed_msgpack_modules = self._allowed_msgpack_modules
|
||||
return self.__class__(
|
||||
pickle_fallback=self.pickle_fallback,
|
||||
allowed_json_modules=self._allowed_json_modules,
|
||||
allowed_msgpack_modules=allowed_msgpack_modules,
|
||||
__unpack_ext_hook__=(
|
||||
self._unpack_ext_hook if self._custom_unpack_ext_hook else None
|
||||
),
|
||||
)
|
||||
|
||||
def _encode_constructor_args(
|
||||
@@ -90,7 +137,7 @@ class JsonPlusSerializer(SerializerProtocol):
|
||||
return out
|
||||
|
||||
def _reviver(self, value: dict[str, Any]) -> Any:
|
||||
if self._allowed_modules and (
|
||||
if self._allowed_json_modules and (
|
||||
value.get("lc", None) == 2
|
||||
and value.get("type", None) == "constructor"
|
||||
and value.get("id", None) is not None
|
||||
@@ -107,7 +154,7 @@ class JsonPlusSerializer(SerializerProtocol):
|
||||
return LC_REVIVER(value)
|
||||
|
||||
def _revive_lc2(self, value: dict[str, Any]) -> Any:
|
||||
self._check_allowed_modules(value)
|
||||
self._check_allowed_json_modules(value)
|
||||
|
||||
[*module, name] = value["id"]
|
||||
try:
|
||||
@@ -139,7 +186,7 @@ class JsonPlusSerializer(SerializerProtocol):
|
||||
except Exception:
|
||||
return None
|
||||
|
||||
def _check_allowed_modules(self, value: dict[str, Any]) -> None:
|
||||
def _check_allowed_json_modules(self, value: dict[str, Any]) -> None:
|
||||
needed = tuple(value["id"])
|
||||
method = value.get("method")
|
||||
if isinstance(method, list):
|
||||
@@ -150,7 +197,7 @@ class JsonPlusSerializer(SerializerProtocol):
|
||||
method_display = "<init>"
|
||||
|
||||
dotted = ".".join(needed)
|
||||
if not self._allowed_modules:
|
||||
if not self._allowed_json_modules:
|
||||
raise InvalidModuleError(
|
||||
f"Refused to deserialize JSON constructor: {dotted} (method: {method_display}). "
|
||||
"No allowed_json_modules configured.\n\n"
|
||||
@@ -161,9 +208,9 @@ class JsonPlusSerializer(SerializerProtocol):
|
||||
"or plain-JSON representations revived without import-time side effects."
|
||||
)
|
||||
|
||||
if self._allowed_modules is True:
|
||||
if self._allowed_json_modules is True:
|
||||
return
|
||||
if needed in self._allowed_modules:
|
||||
if needed in self._allowed_json_modules:
|
||||
return
|
||||
|
||||
raise InvalidModuleError(
|
||||
@@ -448,92 +495,174 @@ def _msgpack_default(obj: Any) -> str | ormsgpack.Ext:
|
||||
raise TypeError(f"Object of type {obj.__class__.__name__} is not serializable")
|
||||
|
||||
|
||||
def _msgpack_ext_hook(code: int, data: bytes) -> Any:
|
||||
if code == EXT_CONSTRUCTOR_SINGLE_ARG:
|
||||
try:
|
||||
tup = ormsgpack.unpackb(
|
||||
data, ext_hook=_msgpack_ext_hook, option=ormsgpack.OPT_NON_STR_KEYS
|
||||
)
|
||||
# module, name, arg
|
||||
return getattr(importlib.import_module(tup[0]), tup[1])(tup[2])
|
||||
except Exception:
|
||||
return
|
||||
elif code == EXT_CONSTRUCTOR_POS_ARGS:
|
||||
try:
|
||||
tup = ormsgpack.unpackb(
|
||||
data, ext_hook=_msgpack_ext_hook, option=ormsgpack.OPT_NON_STR_KEYS
|
||||
)
|
||||
# module, name, args
|
||||
return getattr(importlib.import_module(tup[0]), tup[1])(*tup[2])
|
||||
except Exception:
|
||||
return
|
||||
elif code == EXT_CONSTRUCTOR_KW_ARGS:
|
||||
try:
|
||||
tup = ormsgpack.unpackb(
|
||||
data, ext_hook=_msgpack_ext_hook, option=ormsgpack.OPT_NON_STR_KEYS
|
||||
)
|
||||
# module, name, args
|
||||
return getattr(importlib.import_module(tup[0]), tup[1])(**tup[2])
|
||||
except Exception:
|
||||
return
|
||||
elif code == EXT_METHOD_SINGLE_ARG:
|
||||
try:
|
||||
tup = ormsgpack.unpackb(
|
||||
data, ext_hook=_msgpack_ext_hook, option=ormsgpack.OPT_NON_STR_KEYS
|
||||
)
|
||||
# module, name, arg, method
|
||||
return getattr(getattr(importlib.import_module(tup[0]), tup[1]), tup[3])(
|
||||
tup[2]
|
||||
)
|
||||
except Exception:
|
||||
return
|
||||
elif code == EXT_PYDANTIC_V1:
|
||||
try:
|
||||
tup = ormsgpack.unpackb(
|
||||
data, ext_hook=_msgpack_ext_hook, option=ormsgpack.OPT_NON_STR_KEYS
|
||||
)
|
||||
# module, name, kwargs
|
||||
cls = getattr(importlib.import_module(tup[0]), tup[1])
|
||||
try:
|
||||
return cls(**tup[2])
|
||||
except Exception:
|
||||
return cls.construct(**tup[2])
|
||||
except Exception:
|
||||
# for pydantic objects we can't find/reconstruct
|
||||
# let's return the kwargs dict instead
|
||||
try:
|
||||
return tup[2]
|
||||
except NameError:
|
||||
return
|
||||
elif code == EXT_PYDANTIC_V2:
|
||||
try:
|
||||
tup = ormsgpack.unpackb(
|
||||
data, ext_hook=_msgpack_ext_hook, option=ormsgpack.OPT_NON_STR_KEYS
|
||||
)
|
||||
# module, name, kwargs, method
|
||||
cls = getattr(importlib.import_module(tup[0]), tup[1])
|
||||
try:
|
||||
return cls(**tup[2])
|
||||
except Exception:
|
||||
return cls.model_construct(**tup[2])
|
||||
except Exception:
|
||||
# for pydantic objects we can't find/reconstruct
|
||||
# let's return the kwargs dict instead
|
||||
try:
|
||||
return tup[2]
|
||||
except NameError:
|
||||
return
|
||||
elif code == EXT_NUMPY_ARRAY:
|
||||
try:
|
||||
import numpy as _np
|
||||
def _create_msgpack_ext_hook(
|
||||
allowed_modules: set[tuple[str, ...]] | Literal[True] | None,
|
||||
) -> Callable[[int, bytes], Any]:
|
||||
"""Create msgpack ext hook with allowlist.
|
||||
|
||||
dtype_str, shape, order, buf = ormsgpack.unpackb(
|
||||
data, ext_hook=_msgpack_ext_hook, option=ormsgpack.OPT_NON_STR_KEYS
|
||||
Args:
|
||||
allowed_modules: Set of (module, name) tuples that are allowed to be
|
||||
deserialized, or True to allow all with warnings for unregistered types, or None to only allow safe types.
|
||||
|
||||
Returns:
|
||||
An ext_hook function for use with ormsgpack.unpackb.
|
||||
"""
|
||||
|
||||
def _check_allowed(module: str, name: str) -> bool:
|
||||
"""Check if type is allowed. Returns True if allowed, False if blocked."""
|
||||
key = (module, name)
|
||||
|
||||
if key in _lg_msgpack.SAFE_MSGPACK_TYPES:
|
||||
return True
|
||||
|
||||
if allowed_modules is True:
|
||||
# default is to warn but allow unregistered types
|
||||
logger.warning(
|
||||
"Deserializing unregistered type %s.%s from checkpoint. "
|
||||
"This will be blocked in a future version. "
|
||||
"Add to allowed_msgpack_modules to silence: [(%r, %r)]",
|
||||
module,
|
||||
name,
|
||||
module,
|
||||
name,
|
||||
)
|
||||
arr = _np.frombuffer(buf, dtype=_np.dtype(dtype_str))
|
||||
return arr.reshape(shape, order=order)
|
||||
except Exception:
|
||||
return
|
||||
return True
|
||||
if allowed_modules is not None:
|
||||
if key in allowed_modules:
|
||||
return True
|
||||
# strict mode blocks unregistered types
|
||||
logger.warning(
|
||||
"Blocked deserialization of %s.%s - not in allowed_msgpack_modules. "
|
||||
"Add to allowed_msgpack_modules to allow: [(%r, %r)]",
|
||||
module,
|
||||
name,
|
||||
module,
|
||||
name,
|
||||
)
|
||||
return False
|
||||
|
||||
def _check_allowed_method(module: str, name: str, method: str) -> bool:
|
||||
"""Check if a method invocation is allowed."""
|
||||
key = (module, name, method)
|
||||
if key in _lg_msgpack.SAFE_MSGPACK_METHODS:
|
||||
return True
|
||||
logger.warning(
|
||||
"Blocked deserialization of method call %s.%s.%s - "
|
||||
"not in allowed methods set.",
|
||||
module,
|
||||
name,
|
||||
method,
|
||||
)
|
||||
return False
|
||||
|
||||
def ext_hook(code: int, data: bytes) -> Any:
|
||||
if code == EXT_CONSTRUCTOR_SINGLE_ARG:
|
||||
try:
|
||||
tup = ormsgpack.unpackb(
|
||||
data, ext_hook=ext_hook, option=ormsgpack.OPT_NON_STR_KEYS
|
||||
)
|
||||
if not _check_allowed(tup[0], tup[1]):
|
||||
# We default to returning the raw data. If the user
|
||||
# is using this in the context of a pydantic state, etc., then
|
||||
# it would be validated upon construction.
|
||||
return tup[2]
|
||||
# module, name, arg
|
||||
return getattr(importlib.import_module(tup[0]), tup[1])(tup[2])
|
||||
except Exception:
|
||||
return None
|
||||
elif code == EXT_CONSTRUCTOR_POS_ARGS:
|
||||
try:
|
||||
tup = ormsgpack.unpackb(
|
||||
data, ext_hook=ext_hook, option=ormsgpack.OPT_NON_STR_KEYS
|
||||
)
|
||||
if not _check_allowed(tup[0], tup[1]):
|
||||
return tup[2]
|
||||
# module, name, args
|
||||
return getattr(importlib.import_module(tup[0]), tup[1])(*tup[2])
|
||||
except Exception:
|
||||
return None
|
||||
elif code == EXT_CONSTRUCTOR_KW_ARGS:
|
||||
try:
|
||||
tup = ormsgpack.unpackb(
|
||||
data, ext_hook=ext_hook, option=ormsgpack.OPT_NON_STR_KEYS
|
||||
)
|
||||
if not _check_allowed(tup[0], tup[1]):
|
||||
return tup[2]
|
||||
# module, name, kwargs
|
||||
return getattr(importlib.import_module(tup[0]), tup[1])(**tup[2])
|
||||
except Exception:
|
||||
return None
|
||||
elif code == EXT_METHOD_SINGLE_ARG:
|
||||
try:
|
||||
tup = ormsgpack.unpackb(
|
||||
data, ext_hook=ext_hook, option=ormsgpack.OPT_NON_STR_KEYS
|
||||
)
|
||||
if not _check_allowed_method(tup[0], tup[1], tup[3]):
|
||||
return tup[2]
|
||||
# module, name, arg, method
|
||||
return getattr(
|
||||
getattr(importlib.import_module(tup[0]), tup[1]), tup[3]
|
||||
)(tup[2])
|
||||
except Exception:
|
||||
return None
|
||||
elif code == EXT_PYDANTIC_V1:
|
||||
try:
|
||||
tup = ormsgpack.unpackb(
|
||||
data, ext_hook=ext_hook, option=ormsgpack.OPT_NON_STR_KEYS
|
||||
)
|
||||
if not _check_allowed(tup[0], tup[1]):
|
||||
return tup[2]
|
||||
# module, name, kwargs
|
||||
cls = getattr(importlib.import_module(tup[0]), tup[1])
|
||||
try:
|
||||
return cls(**tup[2])
|
||||
except Exception:
|
||||
return cls.construct(**tup[2])
|
||||
except Exception:
|
||||
# for pydantic objects we can't find/reconstruct
|
||||
# let's return the kwargs dict instead
|
||||
try:
|
||||
return tup[2]
|
||||
except NameError:
|
||||
return None
|
||||
elif code == EXT_PYDANTIC_V2:
|
||||
try:
|
||||
tup = ormsgpack.unpackb(
|
||||
data, ext_hook=ext_hook, option=ormsgpack.OPT_NON_STR_KEYS
|
||||
)
|
||||
if not _check_allowed(tup[0], tup[1]):
|
||||
return tup[2]
|
||||
# module, name, kwargs, method
|
||||
cls = getattr(importlib.import_module(tup[0]), tup[1])
|
||||
try:
|
||||
return cls(**tup[2])
|
||||
except Exception:
|
||||
return cls.model_construct(**tup[2])
|
||||
except Exception:
|
||||
# for pydantic objects we can't find/reconstruct
|
||||
# let's return the kwargs dict instead
|
||||
try:
|
||||
return tup[2]
|
||||
except NameError:
|
||||
return None
|
||||
elif code == EXT_NUMPY_ARRAY:
|
||||
try:
|
||||
import numpy as _np
|
||||
|
||||
dtype_str, shape, order, buf = ormsgpack.unpackb(
|
||||
data, ext_hook=ext_hook, option=ormsgpack.OPT_NON_STR_KEYS
|
||||
)
|
||||
arr = _np.frombuffer(buf, dtype=_np.dtype(dtype_str))
|
||||
return arr.reshape(shape, order=order)
|
||||
except Exception:
|
||||
return None
|
||||
return None
|
||||
|
||||
return ext_hook
|
||||
|
||||
|
||||
# Aliasing in case anyone imported it directly
|
||||
_msgpack_ext_hook = _create_msgpack_ext_hook(allowed_modules=None)
|
||||
|
||||
|
||||
def _msgpack_ext_hook_to_json(code: int, data: bytes) -> Any:
|
||||
@@ -648,3 +777,26 @@ _option = (
|
||||
|
||||
def _msgpack_enc(data: Any) -> bytes:
|
||||
return ormsgpack.packb(data, default=_msgpack_default, option=_option)
|
||||
|
||||
|
||||
def _normalize_allowlist(
|
||||
allowlist: AllowedMsgpackModules | Literal[True] | None,
|
||||
) -> set[tuple[str, ...]] | Literal[True] | None:
|
||||
if allowlist is True:
|
||||
return allowlist
|
||||
elif allowlist:
|
||||
return _normalize_module_keys(allowlist)
|
||||
else:
|
||||
return None
|
||||
|
||||
|
||||
def _normalize_module_keys(
|
||||
modules: AllowedMsgpackModules,
|
||||
) -> set[tuple[str, ...]]:
|
||||
normalized: set[tuple[str, ...]] = set()
|
||||
for module in modules:
|
||||
if isclass(module):
|
||||
normalized.add((module.__module__, module.__name__))
|
||||
else:
|
||||
normalized.add(cast(tuple[str, ...], module))
|
||||
return normalized
|
||||
|
||||
@@ -42,6 +42,7 @@ lint = [
|
||||
dev = [
|
||||
{include-group = "test"},
|
||||
{include-group = "lint"},
|
||||
"pycryptodome>=3.23.0",
|
||||
]
|
||||
|
||||
[tool.hatch.build.targets.wheel]
|
||||
|
||||
@@ -0,0 +1,437 @@
|
||||
"""Tests for EncryptedSerializer with msgpack allowlist functionality.
|
||||
|
||||
These tests mirror the msgpack allowlist tests in test_jsonplus.py but run them
|
||||
through the EncryptedSerializer to ensure the allowlist behavior is preserved
|
||||
when encryption is enabled.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import pathlib
|
||||
import re
|
||||
import uuid
|
||||
from collections import deque
|
||||
from datetime import date, datetime, time, timezone
|
||||
from decimal import Decimal
|
||||
from ipaddress import IPv4Address
|
||||
from typing import Literal, cast
|
||||
|
||||
import ormsgpack
|
||||
import pytest
|
||||
from pydantic import BaseModel
|
||||
|
||||
from langgraph.checkpoint.base import BaseCheckpointSaver, _with_msgpack_allowlist
|
||||
from langgraph.checkpoint.serde import _msgpack as _lg_msgpack
|
||||
from langgraph.checkpoint.serde.base import CipherProtocol
|
||||
from langgraph.checkpoint.serde.encrypted import EncryptedSerializer
|
||||
from langgraph.checkpoint.serde.jsonplus import (
|
||||
EXT_METHOD_SINGLE_ARG,
|
||||
JsonPlusSerializer,
|
||||
_msgpack_enc,
|
||||
)
|
||||
|
||||
|
||||
class InnerPydantic(BaseModel):
|
||||
hello: str
|
||||
|
||||
|
||||
class MyPydantic(BaseModel):
|
||||
foo: str
|
||||
bar: int
|
||||
inner: InnerPydantic
|
||||
|
||||
|
||||
class AnotherPydantic(BaseModel):
|
||||
foo: str
|
||||
|
||||
|
||||
class _PassthroughCipher(CipherProtocol):
|
||||
def encrypt(self, plaintext: bytes) -> tuple[str, bytes]:
|
||||
return "passthrough", plaintext
|
||||
|
||||
def decrypt(self, ciphername: str, ciphertext: bytes) -> bytes:
|
||||
assert ciphername == "passthrough"
|
||||
return ciphertext
|
||||
|
||||
|
||||
def _make_encrypted_serde(
|
||||
allowed_msgpack_modules: (
|
||||
_lg_msgpack.AllowedMsgpackModules | Literal[True] | None | object
|
||||
) = _lg_msgpack._SENTINEL,
|
||||
) -> EncryptedSerializer:
|
||||
"""Create an EncryptedSerializer with AES encryption for testing."""
|
||||
inner = JsonPlusSerializer(
|
||||
allowed_msgpack_modules=cast(
|
||||
_lg_msgpack.AllowedMsgpackModules | Literal[True] | None,
|
||||
allowed_msgpack_modules,
|
||||
)
|
||||
)
|
||||
return EncryptedSerializer.from_pycryptodome_aes(
|
||||
serde=inner, key=b"1234567890123456"
|
||||
)
|
||||
|
||||
|
||||
def test_msgpack_method_pathlib_blocked_encrypted_strict(
|
||||
tmp_path: pathlib.Path, caplog: pytest.LogCaptureFixture
|
||||
) -> None:
|
||||
target = tmp_path / "secret.txt"
|
||||
target.write_text("secret")
|
||||
payload = ormsgpack.packb(
|
||||
ormsgpack.Ext(
|
||||
EXT_METHOD_SINGLE_ARG,
|
||||
_msgpack_enc(("pathlib", "Path", target, "read_text")),
|
||||
),
|
||||
option=ormsgpack.OPT_NON_STR_KEYS,
|
||||
)
|
||||
serde = EncryptedSerializer(
|
||||
_PassthroughCipher(),
|
||||
JsonPlusSerializer(allowed_msgpack_modules=None),
|
||||
)
|
||||
|
||||
caplog.set_level(logging.WARNING, logger="langgraph.checkpoint.serde.jsonplus")
|
||||
caplog.clear()
|
||||
result = serde.loads_typed(("msgpack+passthrough", payload))
|
||||
|
||||
assert result == target
|
||||
assert "blocked deserialization of method call pathlib.path.read_text" in (
|
||||
caplog.text.lower()
|
||||
)
|
||||
|
||||
|
||||
class TestEncryptedSerializerMsgpackAllowlist:
|
||||
"""Test msgpack allowlist behavior through EncryptedSerializer."""
|
||||
|
||||
def test_safe_types_no_warning(self, caplog: pytest.LogCaptureFixture) -> None:
|
||||
"""Test safe types deserialize without warnings through encryption."""
|
||||
serde = _make_encrypted_serde()
|
||||
|
||||
safe_objects = [
|
||||
datetime.now(),
|
||||
date.today(),
|
||||
time(12, 30),
|
||||
timezone.utc,
|
||||
uuid.uuid4(),
|
||||
Decimal("123.45"),
|
||||
{1, 2, 3},
|
||||
frozenset([1, 2, 3]),
|
||||
deque([1, 2, 3]),
|
||||
IPv4Address("192.168.1.1"),
|
||||
pathlib.Path("/tmp/test"),
|
||||
]
|
||||
|
||||
for obj in safe_objects:
|
||||
caplog.clear()
|
||||
dumped = serde.dumps_typed(obj)
|
||||
# Verify encryption is happening
|
||||
assert "+aes" in dumped[0], f"Expected encryption for {type(obj)}"
|
||||
result = serde.loads_typed(dumped)
|
||||
assert "unregistered type" not in caplog.text.lower(), (
|
||||
f"Unexpected warning for {type(obj)}"
|
||||
)
|
||||
assert result is not None
|
||||
|
||||
def test_pydantic_warns_by_default(self, caplog: pytest.LogCaptureFixture) -> None:
|
||||
"""Pydantic models not in allowlist should log warning but still deserialize."""
|
||||
current = _lg_msgpack.STRICT_MSGPACK_ENABLED
|
||||
_lg_msgpack.STRICT_MSGPACK_ENABLED = False
|
||||
serde = _make_encrypted_serde()
|
||||
|
||||
obj = MyPydantic(foo="test", bar=42, inner=InnerPydantic(hello="world"))
|
||||
|
||||
caplog.clear()
|
||||
dumped = serde.dumps_typed(obj)
|
||||
assert "+aes" in dumped[0]
|
||||
result = serde.loads_typed(dumped)
|
||||
|
||||
assert "unregistered type" in caplog.text.lower()
|
||||
assert "allowed_msgpack_modules" in caplog.text
|
||||
assert result == obj
|
||||
_lg_msgpack.STRICT_MSGPACK_ENABLED = current
|
||||
|
||||
def test_strict_mode_blocks_unregistered(
|
||||
self, caplog: pytest.LogCaptureFixture
|
||||
) -> None:
|
||||
"""Strict mode should block unregistered types through encryption."""
|
||||
serde = _make_encrypted_serde(allowed_msgpack_modules=None)
|
||||
|
||||
obj = MyPydantic(foo="test", bar=42, inner=InnerPydantic(hello="world"))
|
||||
|
||||
caplog.clear()
|
||||
dumped = serde.dumps_typed(obj)
|
||||
assert "+aes" in dumped[0]
|
||||
result = serde.loads_typed(dumped)
|
||||
|
||||
assert "blocked" in caplog.text.lower()
|
||||
expected = obj.model_dump()
|
||||
assert result == expected
|
||||
|
||||
def test_allowlist_silences_warning(self, caplog: pytest.LogCaptureFixture) -> None:
|
||||
"""Types in allowed_msgpack_modules should deserialize without warnings."""
|
||||
serde = _make_encrypted_serde(
|
||||
allowed_msgpack_modules=[
|
||||
("tests.test_encrypted", "MyPydantic"),
|
||||
("tests.test_encrypted", "InnerPydantic"),
|
||||
]
|
||||
)
|
||||
|
||||
obj = MyPydantic(foo="test", bar=42, inner=InnerPydantic(hello="world"))
|
||||
|
||||
caplog.clear()
|
||||
dumped = serde.dumps_typed(obj)
|
||||
assert "+aes" in dumped[0]
|
||||
result = serde.loads_typed(dumped)
|
||||
|
||||
assert "unregistered type" not in caplog.text.lower()
|
||||
assert "blocked" not in caplog.text.lower()
|
||||
assert result == obj
|
||||
|
||||
def test_allowlist_blocks_non_listed(
|
||||
self, caplog: pytest.LogCaptureFixture
|
||||
) -> None:
|
||||
"""Allowlists should block unregistered types even through encryption."""
|
||||
serde = _make_encrypted_serde(
|
||||
allowed_msgpack_modules=[("tests.test_encrypted", "MyPydantic")]
|
||||
)
|
||||
|
||||
obj = AnotherPydantic(foo="nope")
|
||||
|
||||
caplog.clear()
|
||||
dumped = serde.dumps_typed(obj)
|
||||
assert "+aes" in dumped[0]
|
||||
result = serde.loads_typed(dumped)
|
||||
|
||||
assert "blocked" in caplog.text.lower()
|
||||
expected = obj.model_dump()
|
||||
assert result == expected
|
||||
|
||||
def test_safe_types_value_equality(self, caplog: pytest.LogCaptureFixture) -> None:
|
||||
"""Verify safe types are correctly restored with proper values through encryption."""
|
||||
serde = _make_encrypted_serde(allowed_msgpack_modules=None)
|
||||
|
||||
test_cases = [
|
||||
datetime(2024, 1, 15, 12, 30, 45, 123456),
|
||||
date(2024, 6, 15),
|
||||
time(14, 30, 0),
|
||||
uuid.UUID("12345678-1234-5678-1234-567812345678"),
|
||||
Decimal("123.456789"),
|
||||
{1, 2, 3, 4, 5},
|
||||
frozenset(["a", "b", "c"]),
|
||||
deque([1, 2, 3]),
|
||||
IPv4Address("10.0.0.1"),
|
||||
pathlib.Path("/some/test/path"),
|
||||
re.compile(r"\d+", re.MULTILINE),
|
||||
]
|
||||
|
||||
for obj in test_cases:
|
||||
caplog.clear()
|
||||
dumped = serde.dumps_typed(obj)
|
||||
assert "+aes" in dumped[0], f"Expected encryption for {type(obj)}"
|
||||
result = serde.loads_typed(dumped)
|
||||
|
||||
assert "blocked" not in caplog.text.lower(), f"Blocked for {type(obj)}"
|
||||
if isinstance(obj, re.Pattern):
|
||||
assert result.pattern == obj.pattern
|
||||
assert result.flags == obj.flags
|
||||
else:
|
||||
assert result == obj, (
|
||||
f"Value mismatch for {type(obj)}: {result} != {obj}"
|
||||
)
|
||||
|
||||
def test_regex_safe_type(self, caplog: pytest.LogCaptureFixture) -> None:
|
||||
"""re.compile patterns should deserialize without warnings as a safe type."""
|
||||
serde = _make_encrypted_serde(allowed_msgpack_modules=None)
|
||||
pattern = re.compile(r"foo.*bar", re.IGNORECASE | re.DOTALL)
|
||||
|
||||
caplog.clear()
|
||||
dumped = serde.dumps_typed(pattern)
|
||||
assert "+aes" in dumped[0]
|
||||
result = serde.loads_typed(dumped)
|
||||
|
||||
assert "blocked" not in caplog.text.lower()
|
||||
assert "unregistered" not in caplog.text.lower()
|
||||
assert result.pattern == pattern.pattern
|
||||
assert result.flags == pattern.flags
|
||||
|
||||
|
||||
class TestWithMsgpackAllowlistEncrypted:
|
||||
"""Test _with_msgpack_allowlist function with EncryptedSerializer."""
|
||||
|
||||
def test_propagates_allowlist_to_inner_serde(self) -> None:
|
||||
"""_with_msgpack_allowlist should propagate allowlist to inner JsonPlusSerializer."""
|
||||
inner = JsonPlusSerializer(allowed_msgpack_modules=None)
|
||||
encrypted = EncryptedSerializer.from_pycryptodome_aes(
|
||||
serde=inner, key=b"1234567890123456"
|
||||
)
|
||||
|
||||
extra = [("my.module", "MyClass")]
|
||||
result = _with_msgpack_allowlist(encrypted, extra)
|
||||
|
||||
# Should return a new EncryptedSerializer
|
||||
assert isinstance(result, EncryptedSerializer)
|
||||
assert result is not encrypted
|
||||
# Inner serde should have the allowlist
|
||||
assert isinstance(result.serde, JsonPlusSerializer)
|
||||
assert isinstance(result.serde._allowed_msgpack_modules, set)
|
||||
assert ("my.module", "MyClass") in result.serde._allowed_msgpack_modules
|
||||
|
||||
def test_preserves_cipher(self) -> None:
|
||||
"""_with_msgpack_allowlist should preserve the cipher from the original."""
|
||||
inner = JsonPlusSerializer(allowed_msgpack_modules=None)
|
||||
encrypted = EncryptedSerializer.from_pycryptodome_aes(
|
||||
serde=inner, key=b"1234567890123456"
|
||||
)
|
||||
|
||||
result = _with_msgpack_allowlist(encrypted, [("my.module", "MyClass")])
|
||||
|
||||
assert isinstance(result, EncryptedSerializer)
|
||||
# Should use the same cipher
|
||||
assert result.cipher is encrypted.cipher
|
||||
|
||||
def test_returns_same_if_not_jsonplus_inner(self) -> None:
|
||||
"""_with_msgpack_allowlist should return same serde if inner is not JsonPlusSerializer."""
|
||||
|
||||
class DummyInnerSerde:
|
||||
def dumps_typed(self, obj: object) -> tuple[str, bytes]:
|
||||
return ("dummy", b"")
|
||||
|
||||
def loads_typed(self, data: tuple[str, bytes]) -> None:
|
||||
return None
|
||||
|
||||
from langgraph.checkpoint.serde.base import CipherProtocol
|
||||
|
||||
class DummyCipher(CipherProtocol):
|
||||
def encrypt(self, plaintext: bytes) -> tuple[str, bytes]:
|
||||
return "dummy", plaintext
|
||||
|
||||
def decrypt(self, ciphername: str, ciphertext: bytes) -> bytes:
|
||||
return ciphertext
|
||||
|
||||
encrypted = EncryptedSerializer(DummyCipher(), DummyInnerSerde())
|
||||
result = _with_msgpack_allowlist(encrypted, [("my.module", "MyClass")])
|
||||
|
||||
assert result is encrypted
|
||||
|
||||
def test_warns_if_allowlist_unsupported(
|
||||
self, caplog: pytest.LogCaptureFixture
|
||||
) -> None:
|
||||
class DummySerde:
|
||||
def dumps_typed(self, obj: object) -> tuple[str, bytes]:
|
||||
return ("dummy", b"")
|
||||
|
||||
def loads_typed(self, data: tuple[str, bytes]) -> object:
|
||||
return data
|
||||
|
||||
serde = DummySerde()
|
||||
caplog.set_level(logging.WARNING, logger="langgraph.checkpoint.base")
|
||||
caplog.clear()
|
||||
|
||||
result = _with_msgpack_allowlist(serde, [("my.module", "MyClass")])
|
||||
|
||||
assert result is serde
|
||||
assert "does not support msgpack allowlist" in caplog.text.lower()
|
||||
|
||||
def test_noop_allowlist_returns_same_encrypted_instance(self) -> None:
|
||||
inner = JsonPlusSerializer(allowed_msgpack_modules=None)
|
||||
encrypted = EncryptedSerializer.from_pycryptodome_aes(
|
||||
serde=inner, key=b"1234567890123456"
|
||||
)
|
||||
|
||||
result = _with_msgpack_allowlist(encrypted, ())
|
||||
|
||||
assert result is encrypted
|
||||
|
||||
def test_functional_roundtrip_with_allowlist(
|
||||
self, caplog: pytest.LogCaptureFixture
|
||||
) -> None:
|
||||
"""End-to-end test: allowlist applied via _with_msgpack_allowlist works."""
|
||||
inner = JsonPlusSerializer(allowed_msgpack_modules=None)
|
||||
encrypted = EncryptedSerializer.from_pycryptodome_aes(
|
||||
serde=inner, key=b"1234567890123456"
|
||||
)
|
||||
|
||||
# Apply allowlist for MyPydantic
|
||||
updated = _with_msgpack_allowlist(
|
||||
encrypted,
|
||||
[
|
||||
("tests.test_encrypted", "MyPydantic"),
|
||||
("tests.test_encrypted", "InnerPydantic"),
|
||||
],
|
||||
)
|
||||
|
||||
obj = MyPydantic(foo="test", bar=42, inner=InnerPydantic(hello="world"))
|
||||
|
||||
caplog.clear()
|
||||
dumped = updated.dumps_typed(obj)
|
||||
assert "+aes" in dumped[0]
|
||||
result = updated.loads_typed(dumped)
|
||||
|
||||
# Should deserialize without blocking
|
||||
assert "blocked" not in caplog.text.lower()
|
||||
assert result == obj
|
||||
|
||||
def test_original_still_blocks_after_with_allowlist(
|
||||
self, caplog: pytest.LogCaptureFixture
|
||||
) -> None:
|
||||
"""Original serde should still block after _with_msgpack_allowlist creates a new one."""
|
||||
inner = JsonPlusSerializer(allowed_msgpack_modules=None)
|
||||
encrypted = EncryptedSerializer.from_pycryptodome_aes(
|
||||
serde=inner, key=b"1234567890123456"
|
||||
)
|
||||
|
||||
# Apply allowlist - this should create a NEW serde
|
||||
_with_msgpack_allowlist(
|
||||
encrypted,
|
||||
[("tests.test_encrypted", "MyPydantic")],
|
||||
)
|
||||
|
||||
# Original should still block
|
||||
obj = MyPydantic(foo="test", bar=42, inner=InnerPydantic(hello="world"))
|
||||
|
||||
caplog.clear()
|
||||
dumped = encrypted.dumps_typed(obj)
|
||||
result = encrypted.loads_typed(dumped)
|
||||
|
||||
assert "blocked" in caplog.text.lower()
|
||||
assert result == obj.model_dump()
|
||||
|
||||
|
||||
class TestEncryptedSerializerUnencryptedFallback:
|
||||
"""Test that EncryptedSerializer handles unencrypted data correctly."""
|
||||
|
||||
def test_loads_unencrypted_data(self) -> None:
|
||||
"""EncryptedSerializer should handle unencrypted data for backwards compat."""
|
||||
plain = JsonPlusSerializer(allowed_msgpack_modules=None)
|
||||
encrypted = _make_encrypted_serde(allowed_msgpack_modules=None)
|
||||
|
||||
obj = {"key": "value", "number": 42}
|
||||
|
||||
# Serialize with plain serde
|
||||
dumped = plain.dumps_typed(obj)
|
||||
assert "+aes" not in dumped[0]
|
||||
|
||||
# Should still deserialize with encrypted serde
|
||||
result = encrypted.loads_typed(dumped)
|
||||
assert result == obj
|
||||
|
||||
|
||||
def test_with_allowlist_uses_copy_protocol() -> None:
|
||||
class CopyAwareSaver(BaseCheckpointSaver[str]):
|
||||
def __init__(self) -> None:
|
||||
super().__init__(serde=JsonPlusSerializer(allowed_msgpack_modules=None))
|
||||
self.copy_was_used = False
|
||||
|
||||
def __copy__(self) -> object:
|
||||
clone = object.__new__(self.__class__)
|
||||
clone.__dict__ = self.__dict__.copy()
|
||||
clone.copy_was_used = True
|
||||
return clone
|
||||
|
||||
saver = CopyAwareSaver()
|
||||
|
||||
updated = saver.with_allowlist([("tests.test_encrypted", "MyPydantic")])
|
||||
|
||||
assert isinstance(updated, CopyAwareSaver)
|
||||
assert updated is not saver
|
||||
assert updated.copy_was_used is True
|
||||
assert saver.copy_was_used is False
|
||||
@@ -1,5 +1,6 @@
|
||||
import dataclasses
|
||||
import json
|
||||
import logging
|
||||
import pathlib
|
||||
import re
|
||||
import sys
|
||||
@@ -13,15 +14,20 @@ from zoneinfo import ZoneInfo
|
||||
|
||||
import dataclasses_json
|
||||
import numpy as np
|
||||
import ormsgpack
|
||||
import pandas as pd
|
||||
import pytest
|
||||
from pydantic import BaseModel, SecretStr
|
||||
from pydantic.v1 import BaseModel as BaseModelV1
|
||||
from pydantic.v1 import SecretStr as SecretStrV1
|
||||
|
||||
from langgraph.checkpoint.serde import _msgpack as _lg_msgpack
|
||||
from langgraph.checkpoint.serde._msgpack import AllowedMsgpackModules
|
||||
from langgraph.checkpoint.serde.jsonplus import (
|
||||
EXT_METHOD_SINGLE_ARG,
|
||||
InvalidModuleError,
|
||||
JsonPlusSerializer,
|
||||
_msgpack_enc,
|
||||
_msgpack_ext_hook_to_json,
|
||||
)
|
||||
from langgraph.store.base import Item
|
||||
@@ -37,6 +43,10 @@ class MyPydantic(BaseModel):
|
||||
inner: InnerPydantic
|
||||
|
||||
|
||||
class AnotherPydantic(BaseModel):
|
||||
foo: str
|
||||
|
||||
|
||||
class InnerPydanticV1(BaseModelV1):
|
||||
hello: str
|
||||
|
||||
@@ -138,7 +148,27 @@ def test_serde_jsonplus() -> None:
|
||||
)
|
||||
to_serialize["my_secret_str_v1"] = SecretStrV1("meow")
|
||||
|
||||
serde = JsonPlusSerializer()
|
||||
allowed_msgpack_modules: AllowedMsgpackModules = [
|
||||
InnerDataclass,
|
||||
MyDataclass,
|
||||
MyDataclassWSlots,
|
||||
MyEnum,
|
||||
InnerPydantic,
|
||||
MyPydantic,
|
||||
# Testing that it supports both.
|
||||
(Person.__module__, Person.__name__),
|
||||
(SecretStr.__module__, SecretStr.__name__),
|
||||
]
|
||||
if sys.version_info < (3, 14):
|
||||
allowed_msgpack_modules.extend( # type: ignore
|
||||
[
|
||||
(InnerPydanticV1.__module__, InnerPydanticV1.__name__),
|
||||
(MyPydanticV1.__module__, MyPydanticV1.__name__),
|
||||
(SecretStrV1.__module__, SecretStrV1.__name__),
|
||||
]
|
||||
)
|
||||
|
||||
serde = JsonPlusSerializer(allowed_msgpack_modules=allowed_msgpack_modules)
|
||||
|
||||
dumped = serde.dumps_typed(to_serialize)
|
||||
|
||||
@@ -512,5 +542,337 @@ def test_serde_jsonplus_pandas_series(series: pd.Series) -> None:
|
||||
|
||||
assert dumped[0] == "pickle"
|
||||
result = serde.loads_typed(dumped)
|
||||
|
||||
assert result.equals(series)
|
||||
|
||||
|
||||
def test_msgpack_safe_types_no_warning(caplog: pytest.LogCaptureFixture) -> None:
|
||||
"""Test safe types deserialize without warnings."""
|
||||
|
||||
serde = JsonPlusSerializer()
|
||||
|
||||
safe_objects = [
|
||||
datetime.now(),
|
||||
date.today(),
|
||||
time(12, 30),
|
||||
timezone.utc,
|
||||
uuid.uuid4(),
|
||||
Decimal("123.45"),
|
||||
{1, 2, 3},
|
||||
frozenset([1, 2, 3]),
|
||||
deque([1, 2, 3]),
|
||||
IPv4Address("192.168.1.1"),
|
||||
pathlib.Path("/tmp/test"),
|
||||
]
|
||||
|
||||
for obj in safe_objects:
|
||||
caplog.clear()
|
||||
dumped = serde.dumps_typed(obj)
|
||||
result = serde.loads_typed(dumped)
|
||||
assert "unregistered type" not in caplog.text.lower(), (
|
||||
f"Unexpected warning for {type(obj)}"
|
||||
)
|
||||
assert result is not None
|
||||
|
||||
|
||||
def test_msgpack_pydantic_warns_by_default(caplog: pytest.LogCaptureFixture) -> None:
|
||||
"""Pydantic models not in allowlist should log warning but still deserialize."""
|
||||
current = _lg_msgpack.STRICT_MSGPACK_ENABLED
|
||||
_lg_msgpack.STRICT_MSGPACK_ENABLED = False
|
||||
serde = JsonPlusSerializer()
|
||||
|
||||
obj = MyPydantic(foo="test", bar=42, inner=InnerPydantic(hello="world"))
|
||||
|
||||
caplog.clear()
|
||||
dumped = serde.dumps_typed(obj)
|
||||
result = serde.loads_typed(dumped)
|
||||
|
||||
assert "unregistered type" in caplog.text.lower()
|
||||
assert "allowed_msgpack_modules" in caplog.text
|
||||
assert result == obj
|
||||
_lg_msgpack.STRICT_MSGPACK_ENABLED = current
|
||||
|
||||
|
||||
def test_msgpack_env_strict_default(
|
||||
caplog: pytest.LogCaptureFixture,
|
||||
) -> None:
|
||||
"""Strict msgpack env should default to blocking unregistered types."""
|
||||
current = _lg_msgpack.STRICT_MSGPACK_ENABLED
|
||||
_lg_msgpack.STRICT_MSGPACK_ENABLED = True
|
||||
serde = JsonPlusSerializer()
|
||||
|
||||
obj = MyPydantic(foo="test", bar=42, inner=InnerPydantic(hello="world"))
|
||||
|
||||
caplog.clear()
|
||||
dumped = serde.dumps_typed(obj)
|
||||
result = serde.loads_typed(dumped)
|
||||
|
||||
assert "blocked" in caplog.text.lower()
|
||||
assert result == obj.model_dump()
|
||||
_lg_msgpack.STRICT_MSGPACK_ENABLED = current
|
||||
|
||||
|
||||
def test_msgpack_allowlist_silences_warning(caplog: pytest.LogCaptureFixture) -> None:
|
||||
"""Types in allowed_msgpack_modules should deserialize without warnings."""
|
||||
|
||||
serde = JsonPlusSerializer(
|
||||
allowed_msgpack_modules=[
|
||||
("tests.test_jsonplus", "MyPydantic"),
|
||||
("tests.test_jsonplus", "InnerPydantic"),
|
||||
]
|
||||
)
|
||||
|
||||
obj = MyPydantic(foo="test", bar=42, inner=InnerPydantic(hello="world"))
|
||||
|
||||
caplog.clear()
|
||||
dumped = serde.dumps_typed(obj)
|
||||
result = serde.loads_typed(dumped)
|
||||
|
||||
assert "unregistered type" not in caplog.text.lower()
|
||||
assert result == obj
|
||||
|
||||
|
||||
def test_msgpack_none_blocks_unregistered(caplog: pytest.LogCaptureFixture) -> None:
|
||||
"""allowed_msgpack_modules=None should block unregistered types."""
|
||||
|
||||
serde = JsonPlusSerializer(allowed_msgpack_modules=None)
|
||||
|
||||
obj = MyPydantic(foo="test", bar=42, inner=InnerPydantic(hello="world"))
|
||||
|
||||
caplog.clear()
|
||||
dumped = serde.dumps_typed(obj)
|
||||
result = serde.loads_typed(dumped)
|
||||
|
||||
assert "blocked" in caplog.text.lower()
|
||||
expected = obj.model_dump()
|
||||
assert result == expected
|
||||
|
||||
|
||||
def test_msgpack_allowlist_blocks_non_listed(
|
||||
caplog: pytest.LogCaptureFixture,
|
||||
) -> None:
|
||||
"""Allowlists should block unregistered types even if msgpack is enabled."""
|
||||
|
||||
serde = JsonPlusSerializer(
|
||||
allowed_msgpack_modules=[("tests.test_jsonplus", "MyPydantic")]
|
||||
)
|
||||
|
||||
obj = AnotherPydantic(foo="nope")
|
||||
|
||||
caplog.clear()
|
||||
dumped = serde.dumps_typed(obj)
|
||||
result = serde.loads_typed(dumped)
|
||||
|
||||
assert "blocked" in caplog.text.lower()
|
||||
expected = obj.model_dump()
|
||||
# It's not allowed, so we just leave it as a dict
|
||||
assert result == expected
|
||||
|
||||
|
||||
def test_msgpack_strict_allows_safe_types(
|
||||
caplog: pytest.LogCaptureFixture,
|
||||
) -> None:
|
||||
"""Safe types should still deserialize in strict mode without warnings."""
|
||||
|
||||
serde = JsonPlusSerializer(allowed_msgpack_modules=None)
|
||||
safe = uuid.uuid4()
|
||||
|
||||
caplog.clear()
|
||||
dumped = serde.dumps_typed(safe)
|
||||
result = serde.loads_typed(dumped)
|
||||
|
||||
assert "blocked" not in caplog.text.lower()
|
||||
assert result == safe
|
||||
|
||||
|
||||
def test_msgpack_regex_safe_type(caplog: pytest.LogCaptureFixture) -> None:
|
||||
"""re.compile patterns should deserialize without warnings as a safe type."""
|
||||
|
||||
serde = JsonPlusSerializer(allowed_msgpack_modules=None)
|
||||
pattern = re.compile(r"foo.*bar", re.IGNORECASE | re.DOTALL)
|
||||
|
||||
caplog.clear()
|
||||
dumped = serde.dumps_typed(pattern)
|
||||
result = serde.loads_typed(dumped)
|
||||
|
||||
assert "blocked" not in caplog.text.lower()
|
||||
assert "unregistered" not in caplog.text.lower()
|
||||
assert result.pattern == pattern.pattern
|
||||
assert result.flags == pattern.flags
|
||||
|
||||
|
||||
def test_msgpack_method_pathlib_blocked_in_strict(
|
||||
tmp_path: pathlib.Path, caplog: pytest.LogCaptureFixture
|
||||
) -> None:
|
||||
target = tmp_path / "secret.txt"
|
||||
target.write_text("secret")
|
||||
serde = JsonPlusSerializer(allowed_msgpack_modules=None)
|
||||
payload = ormsgpack.packb(
|
||||
ormsgpack.Ext(
|
||||
EXT_METHOD_SINGLE_ARG,
|
||||
_msgpack_enc(("pathlib", "Path", target, "read_text")),
|
||||
),
|
||||
option=ormsgpack.OPT_NON_STR_KEYS,
|
||||
)
|
||||
|
||||
caplog.set_level(logging.WARNING, logger="langgraph.checkpoint.serde.jsonplus")
|
||||
caplog.clear()
|
||||
result = serde.loads_typed(("msgpack", payload))
|
||||
|
||||
assert result == target
|
||||
assert "blocked deserialization of method call pathlib.path.read_text" in (
|
||||
caplog.text.lower()
|
||||
)
|
||||
|
||||
|
||||
def test_msgpack_method_pathlib_blocked_default_mode(
|
||||
tmp_path: pathlib.Path, caplog: pytest.LogCaptureFixture
|
||||
) -> None:
|
||||
target = tmp_path / "secret.txt"
|
||||
target.write_text("secret")
|
||||
serde = JsonPlusSerializer(allowed_msgpack_modules=True)
|
||||
payload = ormsgpack.packb(
|
||||
ormsgpack.Ext(
|
||||
EXT_METHOD_SINGLE_ARG,
|
||||
_msgpack_enc(("pathlib", "Path", target, "read_text")),
|
||||
),
|
||||
option=ormsgpack.OPT_NON_STR_KEYS,
|
||||
)
|
||||
|
||||
caplog.set_level(logging.WARNING, logger="langgraph.checkpoint.serde.jsonplus")
|
||||
caplog.clear()
|
||||
result = serde.loads_typed(("msgpack", payload))
|
||||
|
||||
assert result == target
|
||||
assert "blocked deserialization of method call pathlib.path.read_text" in (
|
||||
caplog.text.lower()
|
||||
)
|
||||
|
||||
|
||||
def test_msgpack_regex_still_works_strict(caplog: pytest.LogCaptureFixture) -> None:
|
||||
serde = JsonPlusSerializer(allowed_msgpack_modules=None)
|
||||
pattern = re.compile(r"pattern", re.IGNORECASE | re.MULTILINE)
|
||||
|
||||
caplog.clear()
|
||||
result = serde.loads_typed(serde.dumps_typed(pattern))
|
||||
|
||||
assert "blocked" not in caplog.text.lower()
|
||||
assert result.pattern == pattern.pattern
|
||||
assert result.flags == pattern.flags
|
||||
|
||||
|
||||
def test_msgpack_path_constructor_still_works() -> None:
|
||||
serde = JsonPlusSerializer(allowed_msgpack_modules=None)
|
||||
path_obj = pathlib.Path("/tmp/foo")
|
||||
|
||||
result = serde.loads_typed(serde.dumps_typed(path_obj))
|
||||
|
||||
assert result == path_obj
|
||||
|
||||
|
||||
def test_with_msgpack_allowlist_noop_returns_same_instance() -> None:
|
||||
serde = JsonPlusSerializer(allowed_msgpack_modules=None)
|
||||
|
||||
result = serde.with_msgpack_allowlist(())
|
||||
|
||||
assert result is serde
|
||||
|
||||
|
||||
@pytest.mark.skipif(sys.version_info >= (3, 14), reason="pydantic v1 not on 3.14+")
|
||||
def test_msgpack_pydantic_v1_allowlist(caplog: pytest.LogCaptureFixture) -> None:
|
||||
"""Pydantic v1 models in allowlist should deserialize without warnings."""
|
||||
|
||||
serde = JsonPlusSerializer(
|
||||
allowed_msgpack_modules=[
|
||||
("tests.test_jsonplus", "MyPydanticV1"),
|
||||
("tests.test_jsonplus", "InnerPydanticV1"),
|
||||
]
|
||||
)
|
||||
|
||||
obj = MyPydanticV1(foo="test", bar=42, inner=InnerPydanticV1(hello="world"))
|
||||
|
||||
caplog.clear()
|
||||
dumped = serde.dumps_typed(obj)
|
||||
result = serde.loads_typed(dumped)
|
||||
|
||||
assert "unregistered type" not in caplog.text.lower()
|
||||
assert "blocked" not in caplog.text.lower()
|
||||
assert result == obj
|
||||
|
||||
|
||||
def test_msgpack_dataclass_allowlist(caplog: pytest.LogCaptureFixture) -> None:
|
||||
"""Dataclasses in allowlist should deserialize without warnings."""
|
||||
|
||||
serde = JsonPlusSerializer(
|
||||
allowed_msgpack_modules=[
|
||||
("tests.test_jsonplus", "MyDataclass"),
|
||||
("tests.test_jsonplus", "InnerDataclass"),
|
||||
]
|
||||
)
|
||||
|
||||
obj = MyDataclass(foo="test", bar=42, inner=InnerDataclass(hello="world"))
|
||||
|
||||
caplog.clear()
|
||||
dumped = serde.dumps_typed(obj)
|
||||
result = serde.loads_typed(dumped)
|
||||
|
||||
assert "unregistered type" not in caplog.text.lower()
|
||||
assert "blocked" not in caplog.text.lower()
|
||||
assert result == obj
|
||||
|
||||
|
||||
def test_msgpack_safe_types_value_equality(caplog: pytest.LogCaptureFixture) -> None:
|
||||
"""Verify safe types are correctly restored with proper values."""
|
||||
|
||||
serde = JsonPlusSerializer(allowed_msgpack_modules=None)
|
||||
|
||||
test_cases = [
|
||||
datetime(2024, 1, 15, 12, 30, 45, 123456),
|
||||
date(2024, 6, 15),
|
||||
time(14, 30, 0),
|
||||
uuid.UUID("12345678-1234-5678-1234-567812345678"),
|
||||
Decimal("123.456789"),
|
||||
{1, 2, 3, 4, 5},
|
||||
frozenset(["a", "b", "c"]),
|
||||
deque([1, 2, 3]),
|
||||
IPv4Address("10.0.0.1"),
|
||||
pathlib.Path("/some/test/path"),
|
||||
re.compile(r"\d+", re.MULTILINE),
|
||||
]
|
||||
|
||||
for obj in test_cases:
|
||||
caplog.clear()
|
||||
dumped = serde.dumps_typed(obj)
|
||||
result = serde.loads_typed(dumped)
|
||||
|
||||
assert "blocked" not in caplog.text.lower(), f"Blocked for {type(obj)}"
|
||||
# For regex patterns, compare pattern and flags
|
||||
if isinstance(obj, re.Pattern):
|
||||
assert result.pattern == obj.pattern
|
||||
assert result.flags == obj.flags
|
||||
else:
|
||||
assert result == obj, f"Value mismatch for {type(obj)}: {result} != {obj}"
|
||||
|
||||
|
||||
def test_msgpack_nested_pydantic_serializes_as_dict(
|
||||
caplog: pytest.LogCaptureFixture,
|
||||
) -> None:
|
||||
"""Nested Pydantic models are serialized via model_dump() as dicts.
|
||||
|
||||
This means nested models don't go through the ext hook and don't need
|
||||
to be in the allowlist - only the outer type does.
|
||||
"""
|
||||
|
||||
# Only allow outer type - inner is serialized as dict via model_dump()
|
||||
serde = JsonPlusSerializer(
|
||||
allowed_msgpack_modules=[("tests.test_jsonplus", "MyPydantic")]
|
||||
)
|
||||
|
||||
obj = MyPydantic(foo="test", bar=42, inner=InnerPydantic(hello="world"))
|
||||
|
||||
caplog.clear()
|
||||
dumped = serde.dumps_typed(obj)
|
||||
result = serde.loads_typed(dumped)
|
||||
|
||||
# No blocking should occur - inner is serialized as dict, not ext
|
||||
assert "blocked" not in caplog.text.lower()
|
||||
assert result == obj
|
||||
|
||||
@@ -1,7 +1,9 @@
|
||||
import logging
|
||||
from typing import Any
|
||||
|
||||
import pytest
|
||||
from langchain_core.runnables import RunnableConfig
|
||||
from pydantic import BaseModel
|
||||
|
||||
from langgraph.checkpoint.base import (
|
||||
Checkpoint,
|
||||
@@ -10,6 +12,11 @@ from langgraph.checkpoint.base import (
|
||||
empty_checkpoint,
|
||||
)
|
||||
from langgraph.checkpoint.memory import InMemorySaver
|
||||
from langgraph.checkpoint.serde.jsonplus import JsonPlusSerializer
|
||||
|
||||
|
||||
class MemoryPydantic(BaseModel):
|
||||
foo: str
|
||||
|
||||
|
||||
class TestMemorySaver:
|
||||
@@ -199,3 +206,105 @@ async def test_memory_saver() -> None:
|
||||
|
||||
with memory_saver as sync_memory_saver:
|
||||
assert sync_memory_saver is memory_saver
|
||||
|
||||
|
||||
def test_memory_saver_warns_on_unregistered_msgpack(
|
||||
caplog: pytest.LogCaptureFixture,
|
||||
) -> None:
|
||||
serde = JsonPlusSerializer()
|
||||
memory_saver = InMemorySaver(serde=serde)
|
||||
obj = MemoryPydantic(foo="bar")
|
||||
|
||||
checkpoint = empty_checkpoint()
|
||||
checkpoint["channel_values"] = {"foo": obj}
|
||||
checkpoint["channel_versions"] = {"foo": 1}
|
||||
|
||||
config: RunnableConfig = {
|
||||
"configurable": {"thread_id": "thread-1", "checkpoint_ns": ""}
|
||||
}
|
||||
|
||||
caplog.set_level(logging.WARNING, logger="langgraph.checkpoint.serde.jsonplus")
|
||||
new_config = memory_saver.put(config, checkpoint, {}, {"foo": 1})
|
||||
result = memory_saver.get_tuple(new_config)
|
||||
|
||||
assert result is not None
|
||||
assert "unregistered type" in caplog.text.lower()
|
||||
assert result.checkpoint["channel_values"]["foo"] == obj
|
||||
|
||||
|
||||
def test_memory_saver_allowlist_silences_warning(
|
||||
caplog: pytest.LogCaptureFixture,
|
||||
) -> None:
|
||||
serde = JsonPlusSerializer(
|
||||
allowed_msgpack_modules=[("tests.test_memory", "MemoryPydantic")]
|
||||
)
|
||||
memory_saver = InMemorySaver(serde=serde)
|
||||
obj = MemoryPydantic(foo="bar")
|
||||
|
||||
checkpoint = empty_checkpoint()
|
||||
checkpoint["channel_values"] = {"foo": obj}
|
||||
checkpoint["channel_versions"] = {"foo": 1}
|
||||
|
||||
config: RunnableConfig = {
|
||||
"configurable": {"thread_id": "thread-1", "checkpoint_ns": ""}
|
||||
}
|
||||
|
||||
caplog.set_level(logging.WARNING, logger="langgraph.checkpoint.serde.jsonplus")
|
||||
new_config = memory_saver.put(config, checkpoint, {}, {"foo": 1})
|
||||
result = memory_saver.get_tuple(new_config)
|
||||
|
||||
assert result is not None
|
||||
assert "unregistered type" not in caplog.text.lower()
|
||||
assert result.checkpoint["channel_values"]["foo"] == obj
|
||||
|
||||
|
||||
def test_memory_saver_strict_blocks_unregistered(
|
||||
caplog: pytest.LogCaptureFixture,
|
||||
) -> None:
|
||||
serde = JsonPlusSerializer(allowed_msgpack_modules=None)
|
||||
memory_saver = InMemorySaver(serde=serde)
|
||||
obj = MemoryPydantic(foo="bar")
|
||||
|
||||
checkpoint = empty_checkpoint()
|
||||
checkpoint["channel_values"] = {"foo": obj}
|
||||
checkpoint["channel_versions"] = {"foo": 1}
|
||||
|
||||
config: RunnableConfig = {
|
||||
"configurable": {"thread_id": "thread-1", "checkpoint_ns": ""}
|
||||
}
|
||||
|
||||
caplog.set_level(logging.WARNING, logger="langgraph.checkpoint.serde.jsonplus")
|
||||
new_config = memory_saver.put(config, checkpoint, {}, {"foo": 1})
|
||||
result = memory_saver.get_tuple(new_config)
|
||||
|
||||
assert result is not None
|
||||
assert "blocked" in caplog.text.lower()
|
||||
expected = obj.model_dump() if hasattr(obj, "model_dump") else obj.dict()
|
||||
assert result.checkpoint["channel_values"]["foo"] == expected
|
||||
|
||||
|
||||
def test_memory_saver_with_allowlist_proxy_isolated() -> None:
|
||||
serde = JsonPlusSerializer(allowed_msgpack_modules=None)
|
||||
memory_saver = InMemorySaver(serde=serde)
|
||||
proxy = memory_saver.with_allowlist([("tests.test_memory", "MemoryPydantic")])
|
||||
|
||||
obj = MemoryPydantic(foo="bar")
|
||||
|
||||
checkpoint = empty_checkpoint()
|
||||
checkpoint["channel_values"] = {"foo": obj}
|
||||
checkpoint["channel_versions"] = {"foo": 1}
|
||||
|
||||
config: RunnableConfig = {
|
||||
"configurable": {"thread_id": "thread-1", "checkpoint_ns": ""}
|
||||
}
|
||||
|
||||
new_config = proxy.put(config, checkpoint, {}, {"foo": 1})
|
||||
|
||||
proxied = proxy.get_tuple(new_config)
|
||||
assert proxied is not None
|
||||
assert proxied.checkpoint["channel_values"]["foo"] == obj
|
||||
|
||||
direct = memory_saver.get_tuple(new_config)
|
||||
assert direct is not None
|
||||
expected = obj.model_dump() if hasattr(obj, "model_dump") else obj.dict()
|
||||
assert direct.checkpoint["channel_values"]["foo"] == expected
|
||||
|
||||
Generated
+37
@@ -302,6 +302,7 @@ dev = [
|
||||
{ name = "numpy", version = "2.4.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" },
|
||||
{ name = "pandas" },
|
||||
{ name = "pandas-stubs" },
|
||||
{ name = "pycryptodome" },
|
||||
{ name = "pytest" },
|
||||
{ name = "pytest-asyncio" },
|
||||
{ name = "pytest-mock" },
|
||||
@@ -341,6 +342,7 @@ dev = [
|
||||
{ name = "numpy" },
|
||||
{ name = "pandas" },
|
||||
{ name = "pandas-stubs", specifier = ">=2.2.2.240807" },
|
||||
{ name = "pycryptodome", specifier = ">=3.23.0" },
|
||||
{ name = "pytest" },
|
||||
{ name = "pytest-asyncio" },
|
||||
{ name = "pytest-mock" },
|
||||
@@ -912,6 +914,41 @@ wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/54/20/4d324d65cc6d9205fabedc306948156824eb9f0ee1633355a8f7ec5c66bf/pluggy-1.6.0-py3-none-any.whl", hash = "sha256:e920276dd6813095e9377c0bc5566d94c932c33b27a3e3945d8389c374dd4746", size = 20538, upload-time = "2025-05-15T12:30:06.134Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "pycryptodome"
|
||||
version = "3.23.0"
|
||||
source = { registry = "https://pypi.org/simple" }
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/8e/a6/8452177684d5e906854776276ddd34eca30d1b1e15aa1ee9cefc289a33f5/pycryptodome-3.23.0.tar.gz", hash = "sha256:447700a657182d60338bab09fdb27518f8856aecd80ae4c6bdddb67ff5da44ef", size = 4921276, upload-time = "2025-05-17T17:21:45.242Z" }
|
||||
wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/04/5d/bdb09489b63cd34a976cc9e2a8d938114f7a53a74d3dd4f125ffa49dce82/pycryptodome-3.23.0-cp313-cp313t-macosx_10_13_universal2.whl", hash = "sha256:0011f7f00cdb74879142011f95133274741778abba114ceca229adbf8e62c3e4", size = 2495152, upload-time = "2025-05-17T17:20:20.833Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/a7/ce/7840250ed4cc0039c433cd41715536f926d6e86ce84e904068eb3244b6a6/pycryptodome-3.23.0-cp313-cp313t-macosx_10_13_x86_64.whl", hash = "sha256:90460fc9e088ce095f9ee8356722d4f10f86e5be06e2354230a9880b9c549aae", size = 1639348, upload-time = "2025-05-17T17:20:23.171Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/ee/f0/991da24c55c1f688d6a3b5a11940567353f74590734ee4a64294834ae472/pycryptodome-3.23.0-cp313-cp313t-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:4764e64b269fc83b00f682c47443c2e6e85b18273712b98aa43bcb77f8570477", size = 2184033, upload-time = "2025-05-17T17:20:25.424Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/54/16/0e11882deddf00f68b68dd4e8e442ddc30641f31afeb2bc25588124ac8de/pycryptodome-3.23.0-cp313-cp313t-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:eb8f24adb74984aa0e5d07a2368ad95276cf38051fe2dc6605cbcf482e04f2a7", size = 2270142, upload-time = "2025-05-17T17:20:27.808Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/d5/fc/4347fea23a3f95ffb931f383ff28b3f7b1fe868739182cb76718c0da86a1/pycryptodome-3.23.0-cp313-cp313t-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:d97618c9c6684a97ef7637ba43bdf6663a2e2e77efe0f863cce97a76af396446", size = 2309384, upload-time = "2025-05-17T17:20:30.765Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/6e/d9/c5261780b69ce66d8cfab25d2797bd6e82ba0241804694cd48be41add5eb/pycryptodome-3.23.0-cp313-cp313t-musllinux_1_2_aarch64.whl", hash = "sha256:9a53a4fe5cb075075d515797d6ce2f56772ea7e6a1e5e4b96cf78a14bac3d265", size = 2183237, upload-time = "2025-05-17T17:20:33.736Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/5a/6f/3af2ffedd5cfa08c631f89452c6648c4d779e7772dfc388c77c920ca6bbf/pycryptodome-3.23.0-cp313-cp313t-musllinux_1_2_i686.whl", hash = "sha256:763d1d74f56f031788e5d307029caef067febf890cd1f8bf61183ae142f1a77b", size = 2343898, upload-time = "2025-05-17T17:20:36.086Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/9a/dc/9060d807039ee5de6e2f260f72f3d70ac213993a804f5e67e0a73a56dd2f/pycryptodome-3.23.0-cp313-cp313t-musllinux_1_2_x86_64.whl", hash = "sha256:954af0e2bd7cea83ce72243b14e4fb518b18f0c1649b576d114973e2073b273d", size = 2269197, upload-time = "2025-05-17T17:20:38.414Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/f9/34/e6c8ca177cb29dcc4967fef73f5de445912f93bd0343c9c33c8e5bf8cde8/pycryptodome-3.23.0-cp313-cp313t-win32.whl", hash = "sha256:257bb3572c63ad8ba40b89f6fc9d63a2a628e9f9708d31ee26560925ebe0210a", size = 1768600, upload-time = "2025-05-17T17:20:40.688Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/e4/1d/89756b8d7ff623ad0160f4539da571d1f594d21ee6d68be130a6eccb39a4/pycryptodome-3.23.0-cp313-cp313t-win_amd64.whl", hash = "sha256:6501790c5b62a29fcb227bd6b62012181d886a767ce9ed03b303d1f22eb5c625", size = 1799740, upload-time = "2025-05-17T17:20:42.413Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/5d/61/35a64f0feaea9fd07f0d91209e7be91726eb48c0f1bfc6720647194071e4/pycryptodome-3.23.0-cp313-cp313t-win_arm64.whl", hash = "sha256:9a77627a330ab23ca43b48b130e202582e91cc69619947840ea4d2d1be21eb39", size = 1703685, upload-time = "2025-05-17T17:20:44.388Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/db/6c/a1f71542c969912bb0e106f64f60a56cc1f0fabecf9396f45accbe63fa68/pycryptodome-3.23.0-cp37-abi3-macosx_10_9_universal2.whl", hash = "sha256:187058ab80b3281b1de11c2e6842a357a1f71b42cb1e15bce373f3d238135c27", size = 2495627, upload-time = "2025-05-17T17:20:47.139Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/6e/4e/a066527e079fc5002390c8acdd3aca431e6ea0a50ffd7201551175b47323/pycryptodome-3.23.0-cp37-abi3-macosx_10_9_x86_64.whl", hash = "sha256:cfb5cd445280c5b0a4e6187a7ce8de5a07b5f3f897f235caa11f1f435f182843", size = 1640362, upload-time = "2025-05-17T17:20:50.392Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/50/52/adaf4c8c100a8c49d2bd058e5b551f73dfd8cb89eb4911e25a0c469b6b4e/pycryptodome-3.23.0-cp37-abi3-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:67bd81fcbe34f43ad9422ee8fd4843c8e7198dd88dd3d40e6de42ee65fbe1490", size = 2182625, upload-time = "2025-05-17T17:20:52.866Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/5f/e9/a09476d436d0ff1402ac3867d933c61805ec2326c6ea557aeeac3825604e/pycryptodome-3.23.0-cp37-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:c8987bd3307a39bc03df5c8e0e3d8be0c4c3518b7f044b0f4c15d1aa78f52575", size = 2268954, upload-time = "2025-05-17T17:20:55.027Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/f9/c5/ffe6474e0c551d54cab931918127c46d70cab8f114e0c2b5a3c071c2f484/pycryptodome-3.23.0-cp37-abi3-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:aa0698f65e5b570426fc31b8162ed4603b0c2841cbb9088e2b01641e3065915b", size = 2308534, upload-time = "2025-05-17T17:20:57.279Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/18/28/e199677fc15ecf43010f2463fde4c1a53015d1fe95fb03bca2890836603a/pycryptodome-3.23.0-cp37-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:53ecbafc2b55353edcebd64bf5da94a2a2cdf5090a6915bcca6eca6cc452585a", size = 2181853, upload-time = "2025-05-17T17:20:59.322Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/ce/ea/4fdb09f2165ce1365c9eaefef36625583371ee514db58dc9b65d3a255c4c/pycryptodome-3.23.0-cp37-abi3-musllinux_1_2_i686.whl", hash = "sha256:156df9667ad9f2ad26255926524e1c136d6664b741547deb0a86a9acf5ea631f", size = 2342465, upload-time = "2025-05-17T17:21:03.83Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/22/82/6edc3fc42fe9284aead511394bac167693fb2b0e0395b28b8bedaa07ef04/pycryptodome-3.23.0-cp37-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:dea827b4d55ee390dc89b2afe5927d4308a8b538ae91d9c6f7a5090f397af1aa", size = 2267414, upload-time = "2025-05-17T17:21:06.72Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/59/fe/aae679b64363eb78326c7fdc9d06ec3de18bac68be4b612fc1fe8902693c/pycryptodome-3.23.0-cp37-abi3-win32.whl", hash = "sha256:507dbead45474b62b2bbe318eb1c4c8ee641077532067fec9c1aa82c31f84886", size = 1768484, upload-time = "2025-05-17T17:21:08.535Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/54/2f/e97a1b8294db0daaa87012c24a7bb714147c7ade7656973fd6c736b484ff/pycryptodome-3.23.0-cp37-abi3-win_amd64.whl", hash = "sha256:c75b52aacc6c0c260f204cbdd834f76edc9fb0d8e0da9fbf8352ef58202564e2", size = 1799636, upload-time = "2025-05-17T17:21:10.393Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/18/3d/f9441a0d798bf2b1e645adc3265e55706aead1255ccdad3856dbdcffec14/pycryptodome-3.23.0-cp37-abi3-win_arm64.whl", hash = "sha256:11eeeb6917903876f134b56ba11abe95c0b0fd5e3330def218083c7d98bbcb3c", size = 1703675, upload-time = "2025-05-17T17:21:13.146Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/d9/12/e33935a0709c07de084d7d58d330ec3f4daf7910a18e77937affdb728452/pycryptodome-3.23.0-pp310-pypy310_pp73-macosx_10_15_x86_64.whl", hash = "sha256:ddb95b49df036ddd264a0ad246d1be5b672000f12d6961ea2c267083a5e19379", size = 1623886, upload-time = "2025-05-17T17:21:20.614Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/22/0b/aa8f9419f25870889bebf0b26b223c6986652bdf071f000623df11212c90/pycryptodome-3.23.0-pp310-pypy310_pp73-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:d8e95564beb8782abfd9e431c974e14563a794a4944c29d6d3b7b5ea042110b4", size = 1672151, upload-time = "2025-05-17T17:21:22.666Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/d4/5e/63f5cbde2342b7f70a39e591dbe75d9809d6338ce0b07c10406f1a140cdc/pycryptodome-3.23.0-pp310-pypy310_pp73-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:14e15c081e912c4b0d75632acd8382dfce45b258667aa3c67caf7a4d4c13f630", size = 1664461, upload-time = "2025-05-17T17:21:25.225Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/d6/92/608fbdad566ebe499297a86aae5f2a5263818ceeecd16733006f1600403c/pycryptodome-3.23.0-pp310-pypy310_pp73-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:a7fc76bf273353dc7e5207d172b83f569540fc9a28d63171061c42e361d22353", size = 1702440, upload-time = "2025-05-17T17:21:27.991Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/d1/92/2eadd1341abd2989cce2e2740b4423608ee2014acb8110438244ee97d7ff/pycryptodome-3.23.0-pp310-pypy310_pp73-win_amd64.whl", hash = "sha256:45c69ad715ca1a94f778215a11e66b7ff989d792a4d63b68dc586a1da1392ff5", size = 1803005, upload-time = "2025-05-17T17:21:31.37Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "pydantic"
|
||||
version = "2.12.5"
|
||||
|
||||
@@ -128,7 +128,7 @@ class SerdeConfig(TypedDict, total=False):
|
||||
If omitted, no serde is set up (the object store will still be present, however)."""
|
||||
|
||||
allowed_json_modules: list[list[str]] | bool | None
|
||||
"""Optional. List of allowed python modules to de-serialize custom objects from.
|
||||
"""Optional. List of allowed python modules to de-serialize custom objects from JSON.
|
||||
|
||||
If provided, only the specified modules will be allowed to be deserialized.
|
||||
If omitted, no modules are allowed, and the object returned will simply be a json object OR
|
||||
@@ -148,7 +148,34 @@ class SerdeConfig(TypedDict, total=False):
|
||||
Example:
|
||||
{...
|
||||
"serde": {
|
||||
"allowed_json_modules": true
|
||||
"allowed_json_modules": True
|
||||
}
|
||||
}
|
||||
|
||||
"""
|
||||
allowed_msgpack_modules: list[list[str]] | bool | None
|
||||
"""Optional. List of allowed python modules to de-serialize custom objects from msgpack.
|
||||
|
||||
Known safe types (langgraph.checkpoint.serde.jsonplus.SAFE_MSGPACK_TYPES) are always
|
||||
allowed regardless of this setting. Use this to allowlist your custom Pydantic models,
|
||||
dataclasses, and other user-defined types.
|
||||
|
||||
If True (default), unregistered types will log a warning but still be deserialized.
|
||||
If None, only known safe types will be deserialized; unregistered types will be blocked.
|
||||
|
||||
Example - allowlist specific types (no warnings for these):
|
||||
{...
|
||||
"serde": {
|
||||
"allowed_msgpack_modules": [
|
||||
["my_agent.models", "MyState"],
|
||||
]
|
||||
}
|
||||
}
|
||||
|
||||
Example - strict mode (only safe types allowed):
|
||||
{...
|
||||
"serde": {
|
||||
"allowed_msgpack_modules": null
|
||||
}
|
||||
}
|
||||
|
||||
@@ -328,8 +355,7 @@ class EncryptionConfig(TypedDict, total=False):
|
||||
"""Configuration for custom at-rest encryption logic.
|
||||
|
||||
Allows you to implement custom encryption for sensitive data stored in the database,
|
||||
including metadata fields and checkpoint blobs.
|
||||
"""
|
||||
including metadata fields and checkpoint blobs."""
|
||||
|
||||
path: str
|
||||
"""Required. Path to an instance of the Encryption() class that implements custom encryption handlers.
|
||||
|
||||
@@ -608,7 +608,27 @@
|
||||
"type": "null"
|
||||
}
|
||||
],
|
||||
"description": "Optional. List of allowed python modules to de-serialize custom objects from.\n\nIf provided, only the specified modules will be allowed to be deserialized.\nIf omitted, no modules are allowed, and the object returned will simply be a json object OR\na deserialized langchain object.\n"
|
||||
"description": "Optional. List of allowed python modules to de-serialize custom objects from JSON.\n\nIf provided, only the specified modules will be allowed to be deserialized.\nIf omitted, no modules are allowed, and the object returned will simply be a json object OR\na deserialized langchain object.\n"
|
||||
},
|
||||
"allowed_msgpack_modules": {
|
||||
"anyOf": [
|
||||
{
|
||||
"type": "array",
|
||||
"items": {
|
||||
"type": "array",
|
||||
"items": {
|
||||
"type": "string"
|
||||
}
|
||||
}
|
||||
},
|
||||
{
|
||||
"type": "boolean"
|
||||
},
|
||||
{
|
||||
"type": "null"
|
||||
}
|
||||
],
|
||||
"description": "Optional. List of allowed python modules to de-serialize custom objects from msgpack.\n\nKnown safe types (langgraph.checkpoint.serde.jsonplus.SAFE_MSGPACK_TYPES) are always\nallowed regardless of this setting. Use this to allowlist your custom Pydantic models,\ndataclasses, and other user-defined types.\n\nIf True (default), unregistered types will log a warning but still be deserialized.\nIf None, only known safe types will be deserialized; unregistered types will be blocked.\n\n{...\n[\"my_agent.models\", \"MyState\"],\n]\n}\n}\n\n{...\n}\n}\n\n"
|
||||
},
|
||||
"pickle_fallback": {
|
||||
"type": "boolean",
|
||||
|
||||
@@ -608,7 +608,27 @@
|
||||
"type": "null"
|
||||
}
|
||||
],
|
||||
"description": "Optional. List of allowed python modules to de-serialize custom objects from.\n\nIf provided, only the specified modules will be allowed to be deserialized.\nIf omitted, no modules are allowed, and the object returned will simply be a json object OR\na deserialized langchain object.\n"
|
||||
"description": "Optional. List of allowed python modules to de-serialize custom objects from JSON.\n\nIf provided, only the specified modules will be allowed to be deserialized.\nIf omitted, no modules are allowed, and the object returned will simply be a json object OR\na deserialized langchain object.\n"
|
||||
},
|
||||
"allowed_msgpack_modules": {
|
||||
"anyOf": [
|
||||
{
|
||||
"type": "array",
|
||||
"items": {
|
||||
"type": "array",
|
||||
"items": {
|
||||
"type": "string"
|
||||
}
|
||||
}
|
||||
},
|
||||
{
|
||||
"type": "boolean"
|
||||
},
|
||||
{
|
||||
"type": "null"
|
||||
}
|
||||
],
|
||||
"description": "Optional. List of allowed python modules to de-serialize custom objects from msgpack.\n\nKnown safe types (langgraph.checkpoint.serde.jsonplus.SAFE_MSGPACK_TYPES) are always\nallowed regardless of this setting. Use this to allowlist your custom Pydantic models,\ndataclasses, and other user-defined types.\n\nIf True (default), unregistered types will log a warning but still be deserialized.\nIf None, only known safe types will be deserialized; unregistered types will be blocked.\n\n{...\n[\"my_agent.models\", \"MyState\"],\n]\n}\n}\n\n{...\n}\n}\n\n"
|
||||
},
|
||||
"pickle_fallback": {
|
||||
"type": "boolean",
|
||||
|
||||
@@ -87,15 +87,15 @@ integration_tests:
|
||||
|
||||
WORKERS ?= auto
|
||||
XDIST_ARGS := $(if $(WORKERS),-n $(WORKERS) --dist worksteal,)
|
||||
MAXFAIL ?=
|
||||
MAXFAIL_ARGS := $(if $(MAXFAIL),--maxfail $(MAXFAIL),)
|
||||
MAXFAIL ?= 1
|
||||
MAXFAIL_ARGS = $(if $(MAXFAIL),--maxfail $(MAXFAIL),)
|
||||
# Add an '-x' if xdist is enabled
|
||||
XDIST_ARGS := $(if $(WORKERS),-x $(XDIST_ARGS),)
|
||||
|
||||
test_watch:
|
||||
make start-services &&\
|
||||
make start-dev-server &&\
|
||||
uv run ptw . -- --ff -vv $(XDIST_ARGS) $(MAXFAIL_ARGS) $(TEST); \
|
||||
uv run ptw -- --ff -vv $(XDIST_ARGS) $(MAXFAIL_ARGS) $(TEST); \
|
||||
EXIT_CODE=$$?; \
|
||||
make stop-services; \
|
||||
make stop-dev-server; \
|
||||
@@ -130,7 +130,7 @@ type:
|
||||
|
||||
format format_diff:
|
||||
uv run ruff format $(PYTHON_FILES)
|
||||
uv run ruff check --select I --fix $(PYTHON_FILES)
|
||||
uv run ruff check --fix $(PYTHON_FILES)
|
||||
|
||||
spell_check:
|
||||
uv run codespell --toml pyproject.toml
|
||||
|
||||
@@ -10,6 +10,7 @@ from bench.fanout_to_subgraph import fanout_to_subgraph, fanout_to_subgraph_sync
|
||||
from bench.pydantic_state import pydantic_state
|
||||
from bench.react_agent import react_agent
|
||||
from bench.sequential import create_sequential
|
||||
from bench.serde_allowlist import collect_allowlist_large, collect_allowlist_small
|
||||
from bench.wide_dict import wide_dict
|
||||
from bench.wide_state import wide_state
|
||||
from langgraph.graph import StateGraph
|
||||
@@ -513,3 +514,7 @@ compilation_benchmarks = (
|
||||
|
||||
for name, graph in compilation_benchmarks:
|
||||
r.bench_func(name + "_compilation", compile_graph, graph)
|
||||
|
||||
# Serde allowlist collection
|
||||
r.bench_func("serde_allowlist_small", collect_allowlist_small)
|
||||
r.bench_func("serde_allowlist_large", collect_allowlist_large)
|
||||
|
||||
@@ -0,0 +1,81 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from collections import deque
|
||||
from dataclasses import dataclass
|
||||
from enum import Enum
|
||||
from typing import Annotated
|
||||
|
||||
from pydantic import BaseModel
|
||||
from typing_extensions import NotRequired, TypedDict
|
||||
|
||||
from langgraph._internal._serde import collect_allowlist_from_schemas
|
||||
|
||||
|
||||
class Color(Enum):
|
||||
RED = "red"
|
||||
BLUE = "blue"
|
||||
|
||||
|
||||
@dataclass
|
||||
class InnerDataclass:
|
||||
value: int
|
||||
|
||||
|
||||
class InnerModel(BaseModel):
|
||||
name: str
|
||||
|
||||
|
||||
class InnerTyped(TypedDict):
|
||||
payload: InnerDataclass
|
||||
optional: NotRequired[InnerModel]
|
||||
|
||||
|
||||
@dataclass
|
||||
class Node:
|
||||
value: int
|
||||
child: Node | None = None
|
||||
|
||||
|
||||
@dataclass
|
||||
class NestedDataclass:
|
||||
inner: InnerDataclass
|
||||
items: list[InnerModel]
|
||||
mapping: dict[str, InnerDataclass]
|
||||
optional: InnerModel | None
|
||||
union: InnerDataclass | InnerModel
|
||||
queue: deque[InnerDataclass]
|
||||
frozen: frozenset[InnerModel]
|
||||
|
||||
|
||||
AnnotatedList = Annotated[list[InnerDataclass], "meta"]
|
||||
|
||||
|
||||
class DummyChannel:
|
||||
@property
|
||||
def ValueType(self) -> type[InnerDataclass]:
|
||||
return InnerDataclass
|
||||
|
||||
@property
|
||||
def UpdateType(self) -> type[InnerModel]:
|
||||
return InnerModel
|
||||
|
||||
|
||||
SCHEMAS_SMALL = [InnerDataclass, InnerModel, Color]
|
||||
SCHEMAS_LARGE = [
|
||||
InnerDataclass,
|
||||
InnerModel,
|
||||
Color,
|
||||
InnerTyped,
|
||||
Node,
|
||||
NestedDataclass,
|
||||
AnnotatedList,
|
||||
]
|
||||
CHANNELS = {"a": DummyChannel(), "b": DummyChannel()}
|
||||
|
||||
|
||||
def collect_allowlist_small() -> None:
|
||||
collect_allowlist_from_schemas(schemas=SCHEMAS_SMALL, channels=CHANNELS)
|
||||
|
||||
|
||||
def collect_allowlist_large() -> None:
|
||||
collect_allowlist_from_schemas(schemas=SCHEMAS_LARGE, channels=CHANNELS)
|
||||
@@ -0,0 +1,253 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import dataclasses
|
||||
import logging
|
||||
import sys
|
||||
import types
|
||||
from collections import deque
|
||||
from enum import Enum
|
||||
from typing import (
|
||||
Annotated,
|
||||
Any,
|
||||
Literal,
|
||||
Union,
|
||||
get_args,
|
||||
get_origin,
|
||||
get_type_hints,
|
||||
)
|
||||
|
||||
from langchain_core import messages as lc_messages
|
||||
from langgraph.checkpoint.base import BaseCheckpointSaver
|
||||
from pydantic import BaseModel
|
||||
from typing_extensions import NotRequired, Required, is_typeddict
|
||||
|
||||
try:
|
||||
from langgraph.checkpoint.serde._msgpack import ( # noqa: F401
|
||||
STRICT_MSGPACK_ENABLED,
|
||||
)
|
||||
except ImportError:
|
||||
STRICT_MSGPACK_ENABLED = False
|
||||
|
||||
_warned_allowlist_unsupported = False
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def _supports_checkpointer_allowlist() -> bool:
|
||||
return hasattr(BaseCheckpointSaver, "with_allowlist")
|
||||
|
||||
|
||||
_SUPPORTS_ALLOWLIST = _supports_checkpointer_allowlist()
|
||||
|
||||
|
||||
def apply_checkpointer_allowlist(
|
||||
checkpointer: Any, allowlist: set[tuple[str, ...]] | None
|
||||
) -> Any:
|
||||
if not checkpointer or allowlist is None or checkpointer in (True, False):
|
||||
return checkpointer
|
||||
if not _SUPPORTS_ALLOWLIST:
|
||||
global _warned_allowlist_unsupported
|
||||
if not _warned_allowlist_unsupported:
|
||||
logger.warning(
|
||||
"Checkpointer does not support with_allowlist; strict msgpack "
|
||||
"allowlist will be skipped."
|
||||
)
|
||||
_warned_allowlist_unsupported = True
|
||||
return checkpointer
|
||||
return checkpointer.with_allowlist(allowlist)
|
||||
|
||||
|
||||
def curated_core_allowlist() -> set[tuple[str, ...]]:
|
||||
allowlist: set[tuple[str, ...]] = set()
|
||||
for name in (
|
||||
"BaseMessage",
|
||||
"BaseMessageChunk",
|
||||
"HumanMessage",
|
||||
"HumanMessageChunk",
|
||||
"AIMessage",
|
||||
"AIMessageChunk",
|
||||
"SystemMessage",
|
||||
"SystemMessageChunk",
|
||||
"ChatMessage",
|
||||
"ChatMessageChunk",
|
||||
"ToolMessage",
|
||||
"ToolMessageChunk",
|
||||
"FunctionMessage",
|
||||
"FunctionMessageChunk",
|
||||
"RemoveMessage",
|
||||
):
|
||||
cls = getattr(lc_messages, name, None)
|
||||
if cls is None:
|
||||
continue
|
||||
allowlist.add((cls.__module__, cls.__name__))
|
||||
|
||||
return allowlist
|
||||
|
||||
|
||||
def build_serde_allowlist(
|
||||
*,
|
||||
schemas: list[type[Any]] | None = None,
|
||||
channels: dict[str, Any] | None = None,
|
||||
) -> set[tuple[str, ...]]:
|
||||
allowlist = curated_core_allowlist()
|
||||
if schemas:
|
||||
schemas = [schema for schema in schemas if schema is not None]
|
||||
return allowlist | collect_allowlist_from_schemas(
|
||||
schemas=schemas,
|
||||
channels=channels,
|
||||
)
|
||||
|
||||
|
||||
def collect_allowlist_from_schemas(
|
||||
*,
|
||||
schemas: list[type[Any]] | None = None,
|
||||
channels: dict[str, Any] | None = None,
|
||||
) -> set[tuple[str, ...]]:
|
||||
allowlist: set[tuple[str, ...]] = set()
|
||||
seen: set[Any] = set()
|
||||
seen_ids: set[int] = set()
|
||||
|
||||
if schemas:
|
||||
for schema in schemas:
|
||||
_collect_from_type(schema, allowlist, seen, seen_ids)
|
||||
|
||||
if channels:
|
||||
for channel in channels.values():
|
||||
value_type = getattr(channel, "ValueType", None)
|
||||
if value_type is not None:
|
||||
_collect_from_type(value_type, allowlist, seen, seen_ids)
|
||||
update_type = getattr(channel, "UpdateType", None)
|
||||
if update_type is not None:
|
||||
_collect_from_type(update_type, allowlist, seen, seen_ids)
|
||||
|
||||
return allowlist
|
||||
|
||||
|
||||
def _collect_from_type(
|
||||
typ: Any,
|
||||
allowlist: set[tuple[str, ...]],
|
||||
seen: set[Any],
|
||||
seen_ids: set[int],
|
||||
) -> None:
|
||||
if _already_seen(typ, seen, seen_ids):
|
||||
return
|
||||
|
||||
if typ is Any or typ is None:
|
||||
return
|
||||
|
||||
if typ is Literal:
|
||||
return
|
||||
|
||||
if isinstance(typ, types.UnionType):
|
||||
for arg in typ.__args__:
|
||||
_collect_from_type(arg, allowlist, seen, seen_ids)
|
||||
return
|
||||
|
||||
origin = get_origin(typ)
|
||||
if origin is Union:
|
||||
for arg in get_args(typ):
|
||||
_collect_from_type(arg, allowlist, seen, seen_ids)
|
||||
return
|
||||
if origin is Annotated or origin in (Required, NotRequired):
|
||||
args = get_args(typ)
|
||||
if args:
|
||||
_collect_from_type(args[0], allowlist, seen, seen_ids)
|
||||
return
|
||||
|
||||
if origin is Literal:
|
||||
return
|
||||
|
||||
if origin in (list, set, tuple, dict, deque, frozenset):
|
||||
for arg in get_args(typ):
|
||||
_collect_from_type(arg, allowlist, seen, seen_ids)
|
||||
return
|
||||
|
||||
if hasattr(typ, "__supertype__"):
|
||||
_collect_from_type(typ.__supertype__, allowlist, seen, seen_ids)
|
||||
return
|
||||
|
||||
if is_typeddict(typ):
|
||||
for field_type in _safe_get_type_hints(typ).values():
|
||||
_collect_from_type(field_type, allowlist, seen, seen_ids)
|
||||
return
|
||||
|
||||
if _is_pydantic_model(typ):
|
||||
allowlist.add((typ.__module__, typ.__name__))
|
||||
field_types = _safe_get_type_hints(typ)
|
||||
if field_types:
|
||||
for field_type in field_types.values():
|
||||
_collect_from_type(field_type, allowlist, seen, seen_ids)
|
||||
else:
|
||||
for field_type in _pydantic_field_types(typ):
|
||||
_collect_from_type(field_type, allowlist, seen, seen_ids)
|
||||
return
|
||||
|
||||
if dataclasses.is_dataclass(typ):
|
||||
if typ_name := getattr(typ, "__name__", None):
|
||||
allowlist.add((typ.__module__, typ_name))
|
||||
field_types = _safe_get_type_hints(typ)
|
||||
if field_types:
|
||||
for field_type in field_types.values():
|
||||
_collect_from_type(field_type, allowlist, seen, seen_ids)
|
||||
else:
|
||||
for field in dataclasses.fields(typ):
|
||||
_collect_from_type(field.type, allowlist, seen, seen_ids)
|
||||
return
|
||||
|
||||
if isinstance(typ, type) and issubclass(typ, Enum):
|
||||
allowlist.add((typ.__module__, typ.__name__))
|
||||
return
|
||||
|
||||
|
||||
def _already_seen(typ: Any, seen: set[Any], seen_ids: set[int]) -> bool:
|
||||
try:
|
||||
if typ in seen:
|
||||
return True
|
||||
seen.add(typ)
|
||||
return False
|
||||
except TypeError:
|
||||
typ_id = id(typ)
|
||||
if typ_id in seen_ids:
|
||||
return True
|
||||
seen_ids.add(typ_id)
|
||||
return False
|
||||
|
||||
|
||||
def _safe_get_type_hints(typ: Any) -> dict[str, Any]:
|
||||
try:
|
||||
module = sys.modules.get(getattr(typ, "__module__", ""))
|
||||
globalns = module.__dict__ if module else None
|
||||
localns = dict(vars(typ)) if hasattr(typ, "__dict__") else None
|
||||
return get_type_hints(
|
||||
typ, globalns=globalns, localns=localns, include_extras=True
|
||||
)
|
||||
except Exception:
|
||||
return {}
|
||||
|
||||
|
||||
def _is_pydantic_model(typ: Any) -> bool:
|
||||
if not isinstance(typ, type):
|
||||
return False
|
||||
if issubclass(typ, BaseModel):
|
||||
return True
|
||||
try:
|
||||
from pydantic.v1 import BaseModel as BaseModelV1
|
||||
except Exception:
|
||||
return False
|
||||
return issubclass(typ, BaseModelV1)
|
||||
|
||||
|
||||
def _pydantic_field_types(typ: type[Any]) -> list[Any]:
|
||||
if hasattr(typ, "model_fields"):
|
||||
return [
|
||||
field.annotation
|
||||
for field in typ.model_fields.values()
|
||||
if getattr(field, "annotation", None) is not None
|
||||
]
|
||||
if hasattr(typ, "__fields__"):
|
||||
return [
|
||||
field.outer_type_
|
||||
for field in typ.__fields__.values()
|
||||
if getattr(field, "outer_type_", None) is not None
|
||||
]
|
||||
return []
|
||||
@@ -20,6 +20,7 @@ from langgraph.checkpoint.base import BaseCheckpointSaver
|
||||
from langgraph.store.base import BaseStore
|
||||
from typing_extensions import Unpack
|
||||
|
||||
from langgraph._internal import _serde
|
||||
from langgraph._internal._constants import CACHE_NS_WRITES, PREVIOUS
|
||||
from langgraph._internal._typing import MISSING, DeprecatedKwargs
|
||||
from langgraph.channels.ephemeral_value import EphemeralValue
|
||||
@@ -528,7 +529,7 @@ class entrypoint(Generic[ContextT]):
|
||||
else:
|
||||
output_type = save_type = sig.return_annotation
|
||||
|
||||
return Pregel(
|
||||
graph: Pregel[Any, ContextT, Any, Any] = Pregel(
|
||||
nodes={
|
||||
func.__name__: PregelNode(
|
||||
bound=bound,
|
||||
@@ -559,5 +560,16 @@ class entrypoint(Generic[ContextT]):
|
||||
cache=self.cache,
|
||||
cache_policy=self.cache_policy,
|
||||
retry_policy=self.retry_policy or (),
|
||||
context_schema=self.context_schema, # type: ignore[arg-type]
|
||||
context_schema=self.context_schema,
|
||||
)
|
||||
if _serde.STRICT_MSGPACK_ENABLED:
|
||||
serde_allowlist = _serde.build_serde_allowlist(
|
||||
schemas=[input_type, output_type, save_type]
|
||||
+ ([self.context_schema] if self.context_schema is not None else []),
|
||||
channels=graph.channels,
|
||||
)
|
||||
graph._serde_allowlist = serde_allowlist
|
||||
graph.checkpointer = _serde.apply_checkpointer_allowlist(
|
||||
graph.checkpointer, serde_allowlist
|
||||
)
|
||||
return graph
|
||||
|
||||
@@ -29,6 +29,7 @@ from langgraph.store.base import BaseStore
|
||||
from pydantic import BaseModel, TypeAdapter
|
||||
from typing_extensions import NotRequired, Required, Self, Unpack, is_typeddict
|
||||
|
||||
from langgraph._internal import _serde
|
||||
from langgraph._internal._constants import (
|
||||
INTERRUPT,
|
||||
NS_END,
|
||||
@@ -1079,6 +1080,28 @@ class StateGraph(Generic[StateT, ContextT, InputT, OutputT]):
|
||||
CompiledStateGraph: The compiled `StateGraph`.
|
||||
"""
|
||||
checkpointer = ensure_valid_checkpointer(checkpointer)
|
||||
serde_allowlist: set[tuple[str, ...]] | None = None
|
||||
if _serde.STRICT_MSGPACK_ENABLED:
|
||||
schema_types: list[type[Any]] = [
|
||||
self.state_schema,
|
||||
self.input_schema,
|
||||
self.output_schema,
|
||||
]
|
||||
if self.context_schema is not None:
|
||||
schema_types.append(self.context_schema)
|
||||
for node in self.nodes.values():
|
||||
schema_types.append(node.input_schema)
|
||||
for branches in self.branches.values():
|
||||
for branch in branches.values():
|
||||
if branch.input_schema is not None:
|
||||
schema_types.append(branch.input_schema)
|
||||
serde_allowlist = _serde.build_serde_allowlist(
|
||||
schemas=schema_types,
|
||||
channels=self.channels,
|
||||
)
|
||||
checkpointer = _serde.apply_checkpointer_allowlist(
|
||||
checkpointer, serde_allowlist
|
||||
)
|
||||
|
||||
# assign default values
|
||||
interrupt_before = interrupt_before or []
|
||||
@@ -1135,6 +1158,7 @@ class StateGraph(Generic[StateT, ContextT, InputT, OutputT]):
|
||||
cache=cache,
|
||||
name=name or "LangGraph",
|
||||
)
|
||||
compiled._serde_allowlist = serde_allowlist
|
||||
|
||||
compiled.attach_node(START, None)
|
||||
for key, node in self.nodes.items():
|
||||
|
||||
@@ -48,6 +48,7 @@ from langgraph.store.base import BaseStore
|
||||
from pydantic import BaseModel, TypeAdapter
|
||||
from typing_extensions import Self, Unpack, deprecated, is_typeddict
|
||||
|
||||
from langgraph._internal import _serde
|
||||
from langgraph._internal._config import (
|
||||
ensure_config,
|
||||
merge_configs,
|
||||
@@ -698,9 +699,17 @@ class Pregel(
|
||||
self.config = config
|
||||
self.trigger_to_nodes = trigger_to_nodes or {}
|
||||
self.name = name
|
||||
self._serde_allowlist: set[tuple[str, ...]] | None = None
|
||||
if auto_validate:
|
||||
self.validate()
|
||||
|
||||
def _apply_checkpointer_allowlist(
|
||||
self, checkpointer: BaseCheckpointSaver | None
|
||||
) -> BaseCheckpointSaver | None:
|
||||
if not _serde.STRICT_MSGPACK_ENABLED:
|
||||
return checkpointer
|
||||
return _serde.apply_checkpointer_allowlist(checkpointer, self._serde_allowlist)
|
||||
|
||||
def get_graph(
|
||||
self, config: RunnableConfig | None = None, *, xray: int | bool = False
|
||||
) -> Graph:
|
||||
@@ -1239,6 +1248,8 @@ class Pregel(
|
||||
checkpointer: BaseCheckpointSaver | None = ensure_config(config)[CONF].get(
|
||||
CONFIG_KEY_CHECKPOINTER, self.checkpointer
|
||||
)
|
||||
if isinstance(checkpointer, BaseCheckpointSaver):
|
||||
checkpointer = self._apply_checkpointer_allowlist(checkpointer)
|
||||
if not checkpointer:
|
||||
raise ValueError("No checkpointer set")
|
||||
|
||||
@@ -1281,6 +1292,8 @@ class Pregel(
|
||||
checkpointer: BaseCheckpointSaver | None = ensure_config(config)[CONF].get(
|
||||
CONFIG_KEY_CHECKPOINTER, self.checkpointer
|
||||
)
|
||||
if isinstance(checkpointer, BaseCheckpointSaver):
|
||||
checkpointer = self._apply_checkpointer_allowlist(checkpointer)
|
||||
if not checkpointer:
|
||||
raise ValueError("No checkpointer set")
|
||||
|
||||
@@ -1329,6 +1342,8 @@ class Pregel(
|
||||
checkpointer: BaseCheckpointSaver | None = config[CONF].get(
|
||||
CONFIG_KEY_CHECKPOINTER, self.checkpointer
|
||||
)
|
||||
if isinstance(checkpointer, BaseCheckpointSaver):
|
||||
checkpointer = self._apply_checkpointer_allowlist(checkpointer)
|
||||
if not checkpointer:
|
||||
raise ValueError("No checkpointer set")
|
||||
|
||||
@@ -1380,6 +1395,8 @@ class Pregel(
|
||||
checkpointer: BaseCheckpointSaver | None = ensure_config(config)[CONF].get(
|
||||
CONFIG_KEY_CHECKPOINTER, self.checkpointer
|
||||
)
|
||||
if isinstance(checkpointer, BaseCheckpointSaver):
|
||||
checkpointer = self._apply_checkpointer_allowlist(checkpointer)
|
||||
if not checkpointer:
|
||||
raise ValueError("No checkpointer set")
|
||||
|
||||
@@ -1446,6 +1463,8 @@ class Pregel(
|
||||
checkpointer: BaseCheckpointSaver | None = ensure_config(config)[CONF].get(
|
||||
CONFIG_KEY_CHECKPOINTER, self.checkpointer
|
||||
)
|
||||
if isinstance(checkpointer, BaseCheckpointSaver):
|
||||
checkpointer = self._apply_checkpointer_allowlist(checkpointer)
|
||||
if not checkpointer:
|
||||
raise ValueError("No checkpointer set")
|
||||
|
||||
@@ -1890,6 +1909,8 @@ class Pregel(
|
||||
checkpointer: BaseCheckpointSaver | None = ensure_config(config)[CONF].get(
|
||||
CONFIG_KEY_CHECKPOINTER, self.checkpointer
|
||||
)
|
||||
if isinstance(checkpointer, BaseCheckpointSaver):
|
||||
checkpointer = self._apply_checkpointer_allowlist(checkpointer)
|
||||
if not checkpointer:
|
||||
raise ValueError("No checkpointer set")
|
||||
|
||||
@@ -2378,6 +2399,8 @@ class Pregel(
|
||||
raise RuntimeError("checkpointer=True cannot be used for root graphs.")
|
||||
else:
|
||||
checkpointer = self.checkpointer
|
||||
if isinstance(checkpointer, BaseCheckpointSaver):
|
||||
checkpointer = self._apply_checkpointer_allowlist(checkpointer)
|
||||
if checkpointer and not config.get(CONF):
|
||||
raise ValueError(
|
||||
"Checkpointer requires one or more of the following 'configurable' "
|
||||
|
||||
@@ -1,3 +1,4 @@
|
||||
import os
|
||||
from contextlib import asynccontextmanager, contextmanager
|
||||
from uuid import uuid4
|
||||
|
||||
@@ -5,6 +6,7 @@ import pytest
|
||||
from langgraph.checkpoint.postgres import PostgresSaver
|
||||
from langgraph.checkpoint.postgres.aio import AsyncPostgresSaver
|
||||
from langgraph.checkpoint.serde.encrypted import EncryptedSerializer
|
||||
from langgraph.checkpoint.serde.jsonplus import JsonPlusSerializer
|
||||
from langgraph.checkpoint.sqlite import SqliteSaver
|
||||
from langgraph.checkpoint.sqlite.aio import AsyncSqliteSaver
|
||||
from psycopg import AsyncConnection, Connection
|
||||
@@ -18,30 +20,60 @@ from tests.memory_assert import ( # noqa: E402
|
||||
)
|
||||
|
||||
DEFAULT_POSTGRES_URI = "postgres://postgres:postgres@localhost:5442/"
|
||||
STRICT_MSGPACK = os.getenv("LANGGRAPH_STRICT_MSGPACK", "false").lower() in (
|
||||
"1",
|
||||
"true",
|
||||
"yes",
|
||||
)
|
||||
|
||||
|
||||
def _strict_msgpack_serde() -> JsonPlusSerializer:
|
||||
return JsonPlusSerializer(allowed_msgpack_modules=None)
|
||||
|
||||
|
||||
def _apply_strict_msgpack(checkpointer) -> None:
|
||||
if not STRICT_MSGPACK:
|
||||
return
|
||||
serde = _strict_msgpack_serde()
|
||||
if hasattr(checkpointer, "serde"):
|
||||
checkpointer.serde = serde
|
||||
if hasattr(checkpointer, "saver") and hasattr(checkpointer.saver, "serde"):
|
||||
checkpointer.saver.serde = serde
|
||||
|
||||
|
||||
@contextmanager
|
||||
def _checkpointer_memory():
|
||||
yield MemorySaverAssertImmutable()
|
||||
if STRICT_MSGPACK:
|
||||
yield MemorySaverAssertImmutable(serde=_strict_msgpack_serde())
|
||||
else:
|
||||
yield MemorySaverAssertImmutable()
|
||||
|
||||
|
||||
@contextmanager
|
||||
def _checkpointer_memory_migrate_sends():
|
||||
yield MemorySaverNeedsPendingSendsMigration()
|
||||
checkpointer = MemorySaverNeedsPendingSendsMigration()
|
||||
_apply_strict_msgpack(checkpointer)
|
||||
yield checkpointer
|
||||
|
||||
|
||||
@contextmanager
|
||||
def _checkpointer_sqlite():
|
||||
with SqliteSaver.from_conn_string(":memory:") as checkpointer:
|
||||
_apply_strict_msgpack(checkpointer)
|
||||
yield checkpointer
|
||||
|
||||
|
||||
@contextmanager
|
||||
def _checkpointer_sqlite_aes():
|
||||
with SqliteSaver.from_conn_string(":memory:") as checkpointer:
|
||||
checkpointer.serde = EncryptedSerializer.from_pycryptodome_aes(
|
||||
key=b"1234567890123456"
|
||||
)
|
||||
if STRICT_MSGPACK:
|
||||
checkpointer.serde = EncryptedSerializer.from_pycryptodome_aes(
|
||||
serde=_strict_msgpack_serde(), key=b"1234567890123456"
|
||||
)
|
||||
else:
|
||||
checkpointer.serde = EncryptedSerializer.from_pycryptodome_aes(
|
||||
key=b"1234567890123456"
|
||||
)
|
||||
yield checkpointer
|
||||
|
||||
|
||||
@@ -57,6 +89,7 @@ def _checkpointer_postgres():
|
||||
DEFAULT_POSTGRES_URI + database
|
||||
) as checkpointer:
|
||||
checkpointer.setup()
|
||||
_apply_strict_msgpack(checkpointer)
|
||||
yield checkpointer
|
||||
finally:
|
||||
# drop unique db
|
||||
@@ -79,6 +112,7 @@ def _checkpointer_postgres_pipe():
|
||||
# setup can't run inside pipeline because of implicit transaction
|
||||
with checkpointer.conn.pipeline() as pipe:
|
||||
checkpointer.pipe = pipe
|
||||
_apply_strict_msgpack(checkpointer)
|
||||
yield checkpointer
|
||||
finally:
|
||||
# drop unique db
|
||||
@@ -99,6 +133,7 @@ def _checkpointer_postgres_pool():
|
||||
) as pool:
|
||||
checkpointer = PostgresSaver(pool)
|
||||
checkpointer.setup()
|
||||
_apply_strict_msgpack(checkpointer)
|
||||
yield checkpointer
|
||||
finally:
|
||||
# drop unique db
|
||||
@@ -109,6 +144,7 @@ def _checkpointer_postgres_pool():
|
||||
@asynccontextmanager
|
||||
async def _checkpointer_sqlite_aio():
|
||||
async with AsyncSqliteSaver.from_conn_string(":memory:") as checkpointer:
|
||||
_apply_strict_msgpack(checkpointer)
|
||||
yield checkpointer
|
||||
|
||||
|
||||
@@ -126,6 +162,7 @@ async def _checkpointer_postgres_aio():
|
||||
DEFAULT_POSTGRES_URI + database
|
||||
) as checkpointer:
|
||||
await checkpointer.setup()
|
||||
_apply_strict_msgpack(checkpointer)
|
||||
yield checkpointer
|
||||
finally:
|
||||
# drop unique db
|
||||
@@ -152,6 +189,7 @@ async def _checkpointer_postgres_aio_pipe():
|
||||
# setup can't run inside pipeline because of implicit transaction
|
||||
async with checkpointer.conn.pipeline() as pipe:
|
||||
checkpointer.pipe = pipe
|
||||
_apply_strict_msgpack(checkpointer)
|
||||
yield checkpointer
|
||||
finally:
|
||||
# drop unique db
|
||||
@@ -176,6 +214,7 @@ async def _checkpointer_postgres_aio_pool():
|
||||
) as pool:
|
||||
checkpointer = AsyncPostgresSaver(pool)
|
||||
await checkpointer.setup()
|
||||
_apply_strict_msgpack(checkpointer)
|
||||
yield checkpointer
|
||||
finally:
|
||||
# drop unique db
|
||||
|
||||
@@ -8,6 +8,7 @@ import uuid
|
||||
from enum import Enum
|
||||
from typing import Annotated, Literal, Optional
|
||||
|
||||
from langgraph.checkpoint.base import BaseCheckpointSaver
|
||||
from pydantic import (
|
||||
BaseModel,
|
||||
ByteSize,
|
||||
@@ -23,7 +24,10 @@ from pydantic import (
|
||||
|
||||
from langgraph._internal._pydantic import is_supported_by_pydantic
|
||||
from langgraph.constants import END, START
|
||||
from langgraph.func import entrypoint, task
|
||||
from langgraph.graph.state import StateGraph
|
||||
from langgraph.types import Command, Interrupt, interrupt
|
||||
from tests.any_str import AnyStr
|
||||
|
||||
|
||||
def test_is_supported_by_pydantic() -> None:
|
||||
@@ -312,3 +316,47 @@ def test_pydantic_state_field_validator():
|
||||
g = builder.compile()
|
||||
res = g.invoke(input_state)
|
||||
assert res["text"] == "Hello, Validated John!"
|
||||
|
||||
|
||||
class FunctionalState(BaseModel):
|
||||
a: str
|
||||
b: str | None = None
|
||||
|
||||
|
||||
def test_interrupt_functional_pydantic(sync_checkpointer: BaseCheckpointSaver) -> None:
|
||||
called_count = 0
|
||||
|
||||
@task
|
||||
def foo(state: FunctionalState) -> FunctionalState:
|
||||
nonlocal called_count
|
||||
called_count += 1
|
||||
return FunctionalState(**{"a": state.a + "foo"})
|
||||
|
||||
@task
|
||||
def bar(state: FunctionalState) -> dict:
|
||||
return {"a": state.a + "bar", "b": state.b}
|
||||
|
||||
@entrypoint(checkpointer=sync_checkpointer)
|
||||
def graph(inputs: FunctionalState) -> FunctionalState:
|
||||
fut_foo = foo(inputs)
|
||||
value = interrupt("Provide value for bar:")
|
||||
foo_res = fut_foo.result()
|
||||
assert isinstance(foo_res, FunctionalState)
|
||||
bar_input = FunctionalState(a=foo_res.a, b=value)
|
||||
fut_bar = bar(bar_input)
|
||||
return fut_bar.result()
|
||||
|
||||
config = {"configurable": {"thread_id": "1"}}
|
||||
# First run, interrupted at bar
|
||||
assert graph.invoke(FunctionalState(a=""), config) == {
|
||||
"__interrupt__": [
|
||||
Interrupt(
|
||||
value="Provide value for bar:",
|
||||
id=AnyStr(),
|
||||
)
|
||||
]
|
||||
}
|
||||
# Resume with an answer
|
||||
res = graph.invoke(Command(resume="bar"), config)
|
||||
assert res == {"a": "foobar", "b": "bar"}
|
||||
assert called_count == 1
|
||||
|
||||
@@ -0,0 +1,159 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from collections import deque
|
||||
from dataclasses import dataclass
|
||||
from enum import Enum
|
||||
from typing import TYPE_CHECKING, Annotated, Any, Literal, NewType, Optional, Union
|
||||
|
||||
import pytest
|
||||
from pydantic import BaseModel
|
||||
from typing_extensions import NotRequired, Required, TypedDict
|
||||
|
||||
from langgraph._internal._serde import (
|
||||
collect_allowlist_from_schemas,
|
||||
curated_core_allowlist,
|
||||
)
|
||||
|
||||
|
||||
class Color(Enum):
|
||||
RED = "red"
|
||||
BLUE = "blue"
|
||||
|
||||
|
||||
@dataclass
|
||||
class InnerDataclass:
|
||||
value: int
|
||||
|
||||
|
||||
class InnerModel(BaseModel):
|
||||
name: str
|
||||
|
||||
|
||||
@dataclass
|
||||
class Node:
|
||||
value: int
|
||||
child: Node | None = None
|
||||
|
||||
|
||||
if TYPE_CHECKING:
|
||||
|
||||
class MissingType:
|
||||
pass
|
||||
|
||||
|
||||
@dataclass
|
||||
class MissingRefDataclass:
|
||||
payload: MissingType
|
||||
|
||||
|
||||
class Payload(TypedDict):
|
||||
item: InnerDataclass
|
||||
maybe: NotRequired[InnerModel]
|
||||
required: Required[str]
|
||||
|
||||
|
||||
@dataclass
|
||||
class NestedDataclass:
|
||||
inner: InnerDataclass
|
||||
items: list[InnerModel]
|
||||
mapping: dict[str, InnerDataclass]
|
||||
optional: InnerModel | None
|
||||
union: InnerDataclass | InnerModel
|
||||
queue: deque[InnerDataclass]
|
||||
frozen: frozenset[InnerModel]
|
||||
|
||||
|
||||
AnnotatedList = Annotated[list[InnerDataclass], "meta"]
|
||||
UserId = NewType("UserId", int)
|
||||
|
||||
|
||||
class DummyChannel:
|
||||
@property
|
||||
def ValueType(self) -> type[InnerDataclass]:
|
||||
return InnerDataclass
|
||||
|
||||
@property
|
||||
def UpdateType(self) -> type[InnerModel]:
|
||||
return InnerModel
|
||||
|
||||
|
||||
def test_curated_core_allowlist_includes_messages() -> None:
|
||||
try:
|
||||
from langchain_core.messages import BaseMessage
|
||||
except Exception:
|
||||
pytest.skip("langchain_core not available")
|
||||
allowlist = curated_core_allowlist()
|
||||
assert (BaseMessage.__module__, BaseMessage.__name__) in allowlist
|
||||
|
||||
|
||||
def test_collect_allowlist_basic_models() -> None:
|
||||
allowlist = collect_allowlist_from_schemas(
|
||||
schemas=[InnerDataclass, InnerModel, Color]
|
||||
)
|
||||
assert (InnerDataclass.__module__, InnerDataclass.__name__) in allowlist
|
||||
assert (InnerModel.__module__, InnerModel.__name__) in allowlist
|
||||
assert (Color.__module__, Color.__name__) in allowlist
|
||||
|
||||
|
||||
def test_collect_allowlist_nested_containers() -> None:
|
||||
allowlist = collect_allowlist_from_schemas(schemas=[NestedDataclass])
|
||||
assert (NestedDataclass.__module__, NestedDataclass.__name__) in allowlist
|
||||
assert (InnerDataclass.__module__, InnerDataclass.__name__) in allowlist
|
||||
assert (InnerModel.__module__, InnerModel.__name__) in allowlist
|
||||
|
||||
|
||||
def test_collect_allowlist_annotated_and_union() -> None:
|
||||
allowlist = collect_allowlist_from_schemas(
|
||||
schemas=[AnnotatedList, InnerModel | None, InnerDataclass | None]
|
||||
)
|
||||
assert (InnerDataclass.__module__, InnerDataclass.__name__) in allowlist
|
||||
assert (InnerModel.__module__, InnerModel.__name__) in allowlist
|
||||
|
||||
|
||||
def test_collect_allowlist_literal_and_any() -> None:
|
||||
allowlist = collect_allowlist_from_schemas(schemas=[Any, Literal["a"]])
|
||||
assert allowlist == set()
|
||||
|
||||
|
||||
def test_collect_allowlist_typeddict_fields_only() -> None:
|
||||
allowlist = collect_allowlist_from_schemas(schemas=[Payload])
|
||||
assert (InnerDataclass.__module__, InnerDataclass.__name__) in allowlist
|
||||
assert (InnerModel.__module__, InnerModel.__name__) in allowlist
|
||||
assert (Payload.__module__, Payload.__name__) not in allowlist
|
||||
|
||||
|
||||
def test_collect_allowlist_forward_refs() -> None:
|
||||
allowlist = collect_allowlist_from_schemas(schemas=[Node])
|
||||
assert (Node.__module__, Node.__name__) in allowlist
|
||||
|
||||
|
||||
def test_collect_allowlist_missing_forward_ref() -> None:
|
||||
allowlist = collect_allowlist_from_schemas(schemas=[MissingRefDataclass])
|
||||
assert allowlist == {(MissingRefDataclass.__module__, MissingRefDataclass.__name__)}
|
||||
|
||||
|
||||
def test_collect_allowlist_newtype_supertype() -> None:
|
||||
allowlist = collect_allowlist_from_schemas(schemas=[UserId])
|
||||
assert allowlist == set()
|
||||
|
||||
|
||||
def test_collect_allowlist_channels() -> None:
|
||||
channels = {"a": DummyChannel(), "b": DummyChannel()}
|
||||
allowlist = collect_allowlist_from_schemas(channels=channels)
|
||||
assert (InnerDataclass.__module__, InnerDataclass.__name__) in allowlist
|
||||
assert (InnerModel.__module__, InnerModel.__name__) in allowlist
|
||||
|
||||
|
||||
def test_collect_allowlist_pep604_union() -> None:
|
||||
schema = InnerDataclass | InnerModel
|
||||
allowlist = collect_allowlist_from_schemas(schemas=[schema])
|
||||
assert (InnerDataclass.__module__, InnerDataclass.__name__) in allowlist
|
||||
assert (InnerModel.__module__, InnerModel.__name__) in allowlist
|
||||
|
||||
|
||||
def test_collect_allowlist_typing_union_optional() -> None:
|
||||
typing_optional = Optional[InnerDataclass] # noqa: UP045
|
||||
typing_union = Union[InnerDataclass, InnerModel] # noqa: UP007
|
||||
allowlist = collect_allowlist_from_schemas(schemas=[typing_optional, typing_union])
|
||||
assert (InnerDataclass.__module__, InnerDataclass.__name__) in allowlist
|
||||
assert (InnerModel.__module__, InnerModel.__name__) in allowlist
|
||||
Generated
+1
@@ -1569,6 +1569,7 @@ dev = [
|
||||
{ name = "numpy" },
|
||||
{ name = "pandas" },
|
||||
{ name = "pandas-stubs", specifier = ">=2.2.2.240807" },
|
||||
{ name = "pycryptodome", specifier = ">=3.23.0" },
|
||||
{ name = "pytest" },
|
||||
{ name = "pytest-asyncio" },
|
||||
{ name = "pytest-mock" },
|
||||
|
||||
Generated
+1
@@ -373,6 +373,7 @@ dev = [
|
||||
{ name = "numpy" },
|
||||
{ name = "pandas" },
|
||||
{ name = "pandas-stubs", specifier = ">=2.2.2.240807" },
|
||||
{ name = "pycryptodome", specifier = ">=3.23.0" },
|
||||
{ name = "pytest" },
|
||||
{ name = "pytest-asyncio" },
|
||||
{ name = "pytest-mock" },
|
||||
|
||||
Reference in New Issue
Block a user