mirror of
https://github.com/langchain-ai/langgraph.git
synced 2026-09-12 04:37:51 +02:00
Compare commits
8
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
60c9602bcc | ||
|
|
40e471cdd4 | ||
|
|
6fe30ca652 | ||
|
|
222f23a810 | ||
|
|
a568099bf2 | ||
|
|
cea0f1de87 | ||
|
|
6c16d7ab4b | ||
|
|
9a7d174380 |
@@ -56,3 +56,46 @@ jobs:
|
|||||||
# grep will exit non-zero if the target message isn't found,
|
# grep will exit non-zero if the target message isn't found,
|
||||||
# and `set -e` above will cause the step to fail.
|
# and `set -e` above will cause the step to fail.
|
||||||
echo "$STATUS" | grep 'nothing to commit, working tree clean'
|
echo "$STATUS" | grep 'nothing to commit, working tree clean'
|
||||||
|
|
||||||
|
build-strict-msgpack:
|
||||||
|
runs-on: ubuntu-latest
|
||||||
|
defaults:
|
||||||
|
run:
|
||||||
|
working-directory: libs/langgraph
|
||||||
|
name: "test strict msgpack #3.12"
|
||||||
|
steps:
|
||||||
|
- uses: actions/checkout@v6
|
||||||
|
- name: Set up Python 3.12
|
||||||
|
uses: astral-sh/setup-uv@v7
|
||||||
|
with:
|
||||||
|
python-version: "3.12"
|
||||||
|
enable-cache: true
|
||||||
|
cache-suffix: "test-langgraph-strict-msgpack"
|
||||||
|
- name: Login to Docker Hub
|
||||||
|
uses: docker/login-action@v3
|
||||||
|
if: ${{ !github.event.pull_request.head.repo.fork }}
|
||||||
|
with:
|
||||||
|
username: ${{ secrets.DOCKERHUB_USERNAME }}
|
||||||
|
password: ${{ secrets.DOCKERHUB_RO_TOKEN }}
|
||||||
|
|
||||||
|
- name: Install dependencies
|
||||||
|
shell: bash
|
||||||
|
run: uv sync --frozen --group test --no-dev
|
||||||
|
|
||||||
|
- name: Run tests (strict msgpack)
|
||||||
|
shell: bash
|
||||||
|
env:
|
||||||
|
LANGGRAPH_STRICT_MSGPACK: "1"
|
||||||
|
run: make test_parallel
|
||||||
|
|
||||||
|
- name: Ensure the tests did not create any additional files
|
||||||
|
shell: bash
|
||||||
|
run: |
|
||||||
|
set -eu
|
||||||
|
|
||||||
|
STATUS="$(git status)"
|
||||||
|
echo "$STATUS"
|
||||||
|
|
||||||
|
# grep will exit non-zero if the target message isn't found,
|
||||||
|
# and `set -e` above will cause the step to fail.
|
||||||
|
echo "$STATUS" | grep 'nothing to commit, working tree clean'
|
||||||
|
|||||||
@@ -37,6 +37,48 @@ LC_REVIVER = Reviver()
|
|||||||
EMPTY_BYTES = b""
|
EMPTY_BYTES = b""
|
||||||
logger = logging.getLogger(__name__)
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
|
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"),
|
||||||
|
# regex
|
||||||
|
("re", "compile"),
|
||||||
|
# langgraph
|
||||||
|
("langgraph.types", "Send"),
|
||||||
|
("langgraph.types", "Interrupt"),
|
||||||
|
("langgraph.types", "Command"),
|
||||||
|
("langgraph.types", "StateSnapshot"),
|
||||||
|
("langgraph.types", "PregelTask"),
|
||||||
|
}
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
class JsonPlusSerializer(SerializerProtocol):
|
class JsonPlusSerializer(SerializerProtocol):
|
||||||
"""Serializer that uses ormsgpack, with optional fallbacks.
|
"""Serializer that uses ormsgpack, with optional fallbacks.
|
||||||
@@ -54,18 +96,29 @@ class JsonPlusSerializer(SerializerProtocol):
|
|||||||
*,
|
*,
|
||||||
pickle_fallback: bool = False,
|
pickle_fallback: bool = False,
|
||||||
allowed_json_modules: Sequence[tuple[str, ...]] | Literal[True] | None = None,
|
allowed_json_modules: Sequence[tuple[str, ...]] | Literal[True] | None = None,
|
||||||
|
# TODO: change default to None once users have had time to configure allowlists
|
||||||
|
allowed_msgpack_modules: Sequence[tuple[str, ...]]
|
||||||
|
| Literal[True]
|
||||||
|
| None = True,
|
||||||
__unpack_ext_hook__: Callable[[int, bytes], Any] | None = None,
|
__unpack_ext_hook__: Callable[[int, bytes], Any] | None = None,
|
||||||
) -> None:
|
) -> None:
|
||||||
self.pickle_fallback = pickle_fallback
|
self.pickle_fallback = pickle_fallback
|
||||||
self._allowed_modules = (
|
# JSON allowlist
|
||||||
|
self._allowed_json_modules: set[tuple[str, ...]] | Literal[True] | None = (
|
||||||
{mod_and_name for mod_and_name in allowed_json_modules}
|
{mod_and_name for mod_and_name in allowed_json_modules}
|
||||||
if allowed_json_modules and allowed_json_modules is not True
|
if allowed_json_modules and allowed_json_modules is not True
|
||||||
else (allowed_json_modules if allowed_json_modules is True else None)
|
else (allowed_json_modules if allowed_json_modules is True else None)
|
||||||
)
|
)
|
||||||
|
# Msgpack allowlist
|
||||||
|
self._allowed_msgpack_modules: set[tuple[str, ...]] | Literal[True] | None = (
|
||||||
|
{mod_and_name for mod_and_name in allowed_msgpack_modules}
|
||||||
|
if allowed_msgpack_modules and allowed_msgpack_modules is not True
|
||||||
|
else (allowed_msgpack_modules if allowed_msgpack_modules is True else None)
|
||||||
|
)
|
||||||
self._unpack_ext_hook = (
|
self._unpack_ext_hook = (
|
||||||
__unpack_ext_hook__
|
__unpack_ext_hook__
|
||||||
if __unpack_ext_hook__ is not None
|
if __unpack_ext_hook__ is not None
|
||||||
else _msgpack_ext_hook
|
else _create_msgpack_ext_hook(self._allowed_msgpack_modules)
|
||||||
)
|
)
|
||||||
|
|
||||||
def _encode_constructor_args(
|
def _encode_constructor_args(
|
||||||
@@ -90,7 +143,7 @@ class JsonPlusSerializer(SerializerProtocol):
|
|||||||
return out
|
return out
|
||||||
|
|
||||||
def _reviver(self, value: dict[str, Any]) -> Any:
|
def _reviver(self, value: dict[str, Any]) -> Any:
|
||||||
if self._allowed_modules and (
|
if self._allowed_json_modules and (
|
||||||
value.get("lc", None) == 2
|
value.get("lc", None) == 2
|
||||||
and value.get("type", None) == "constructor"
|
and value.get("type", None) == "constructor"
|
||||||
and value.get("id", None) is not None
|
and value.get("id", None) is not None
|
||||||
@@ -107,7 +160,7 @@ class JsonPlusSerializer(SerializerProtocol):
|
|||||||
return LC_REVIVER(value)
|
return LC_REVIVER(value)
|
||||||
|
|
||||||
def _revive_lc2(self, value: dict[str, Any]) -> Any:
|
def _revive_lc2(self, value: dict[str, Any]) -> Any:
|
||||||
self._check_allowed_modules(value)
|
self._check_allowed_json_modules(value)
|
||||||
|
|
||||||
[*module, name] = value["id"]
|
[*module, name] = value["id"]
|
||||||
try:
|
try:
|
||||||
@@ -139,7 +192,7 @@ class JsonPlusSerializer(SerializerProtocol):
|
|||||||
except Exception:
|
except Exception:
|
||||||
return None
|
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"])
|
needed = tuple(value["id"])
|
||||||
method = value.get("method")
|
method = value.get("method")
|
||||||
if isinstance(method, list):
|
if isinstance(method, list):
|
||||||
@@ -150,7 +203,7 @@ class JsonPlusSerializer(SerializerProtocol):
|
|||||||
method_display = "<init>"
|
method_display = "<init>"
|
||||||
|
|
||||||
dotted = ".".join(needed)
|
dotted = ".".join(needed)
|
||||||
if not self._allowed_modules:
|
if not self._allowed_json_modules:
|
||||||
raise InvalidModuleError(
|
raise InvalidModuleError(
|
||||||
f"Refused to deserialize JSON constructor: {dotted} (method: {method_display}). "
|
f"Refused to deserialize JSON constructor: {dotted} (method: {method_display}). "
|
||||||
"No allowed_json_modules configured.\n\n"
|
"No allowed_json_modules configured.\n\n"
|
||||||
@@ -161,9 +214,9 @@ class JsonPlusSerializer(SerializerProtocol):
|
|||||||
"or plain-JSON representations revived without import-time side effects."
|
"or plain-JSON representations revived without import-time side effects."
|
||||||
)
|
)
|
||||||
|
|
||||||
if self._allowed_modules is True:
|
if self._allowed_json_modules is True:
|
||||||
return
|
return
|
||||||
if needed in self._allowed_modules:
|
if needed in self._allowed_json_modules:
|
||||||
return
|
return
|
||||||
|
|
||||||
raise InvalidModuleError(
|
raise InvalidModuleError(
|
||||||
@@ -448,92 +501,158 @@ def _msgpack_default(obj: Any) -> str | ormsgpack.Ext:
|
|||||||
raise TypeError(f"Object of type {obj.__class__.__name__} is not serializable")
|
raise TypeError(f"Object of type {obj.__class__.__name__} is not serializable")
|
||||||
|
|
||||||
|
|
||||||
def _msgpack_ext_hook(code: int, data: bytes) -> Any:
|
def _create_msgpack_ext_hook(
|
||||||
if code == EXT_CONSTRUCTOR_SINGLE_ARG:
|
allowed_modules: set[tuple[str, ...]] | Literal[True] | None,
|
||||||
try:
|
) -> Callable[[int, bytes], Any]:
|
||||||
tup = ormsgpack.unpackb(
|
"""Create msgpack ext hook with allowlist.
|
||||||
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
|
|
||||||
|
|
||||||
dtype_str, shape, order, buf = ormsgpack.unpackb(
|
Args:
|
||||||
data, ext_hook=_msgpack_ext_hook, option=ormsgpack.OPT_NON_STR_KEYS
|
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 SAFE_MSGPACK_TYPES:
|
||||||
|
return True
|
||||||
|
|
||||||
|
if allowed_modules is not None and allowed_modules is not True:
|
||||||
|
if key in allowed_modules:
|
||||||
|
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 True
|
||||||
return arr.reshape(shape, order=order)
|
else:
|
||||||
except Exception:
|
# strict mode blocks unregistered types
|
||||||
return
|
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 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]):
|
||||||
|
return None
|
||||||
|
# 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 None
|
||||||
|
# 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 None
|
||||||
|
# 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(tup[0], tup[1]):
|
||||||
|
return None
|
||||||
|
# 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 None
|
||||||
|
# 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 None
|
||||||
|
# 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
|
||||||
|
|
||||||
|
|
||||||
|
_msgpack_ext_hook = _create_msgpack_ext_hook(allowed_modules=None)
|
||||||
|
|
||||||
|
|
||||||
def _msgpack_ext_hook_to_json(code: int, data: bytes) -> Any:
|
def _msgpack_ext_hook_to_json(code: int, data: bytes) -> Any:
|
||||||
|
|||||||
@@ -37,6 +37,10 @@ class MyPydantic(BaseModel):
|
|||||||
inner: InnerPydantic
|
inner: InnerPydantic
|
||||||
|
|
||||||
|
|
||||||
|
class AnotherPydantic(BaseModel):
|
||||||
|
foo: str
|
||||||
|
|
||||||
|
|
||||||
class InnerPydanticV1(BaseModelV1):
|
class InnerPydanticV1(BaseModelV1):
|
||||||
hello: str
|
hello: str
|
||||||
|
|
||||||
@@ -512,5 +516,240 @@ def test_serde_jsonplus_pandas_series(series: pd.Series) -> None:
|
|||||||
|
|
||||||
assert dumped[0] == "pickle"
|
assert dumped[0] == "pickle"
|
||||||
result = serde.loads_typed(dumped)
|
result = serde.loads_typed(dumped)
|
||||||
|
|
||||||
assert result.equals(series)
|
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.
|
||||||
|
|
||||||
|
TODO: We'll want to change this to block unregistered types in the future."""
|
||||||
|
|
||||||
|
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
|
||||||
|
|
||||||
|
|
||||||
|
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.
|
||||||
|
|
||||||
|
TODO: This will be the default behavior in the future."""
|
||||||
|
|
||||||
|
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()
|
||||||
|
assert result is None
|
||||||
|
|
||||||
|
|
||||||
|
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()
|
||||||
|
assert result is None
|
||||||
|
|
||||||
|
|
||||||
|
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
|
||||||
|
|
||||||
|
|
||||||
|
@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
|
from typing import Any
|
||||||
|
|
||||||
import pytest
|
import pytest
|
||||||
from langchain_core.runnables import RunnableConfig
|
from langchain_core.runnables import RunnableConfig
|
||||||
|
from pydantic import BaseModel
|
||||||
|
|
||||||
from langgraph.checkpoint.base import (
|
from langgraph.checkpoint.base import (
|
||||||
Checkpoint,
|
Checkpoint,
|
||||||
@@ -10,6 +12,11 @@ from langgraph.checkpoint.base import (
|
|||||||
empty_checkpoint,
|
empty_checkpoint,
|
||||||
)
|
)
|
||||||
from langgraph.checkpoint.memory import InMemorySaver
|
from langgraph.checkpoint.memory import InMemorySaver
|
||||||
|
from langgraph.checkpoint.serde.jsonplus import JsonPlusSerializer
|
||||||
|
|
||||||
|
|
||||||
|
class MemoryPydantic(BaseModel):
|
||||||
|
foo: str
|
||||||
|
|
||||||
|
|
||||||
class TestMemorySaver:
|
class TestMemorySaver:
|
||||||
@@ -199,3 +206,77 @@ async def test_memory_saver() -> None:
|
|||||||
|
|
||||||
with memory_saver as sync_memory_saver:
|
with memory_saver as sync_memory_saver:
|
||||||
assert sync_memory_saver is 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()
|
||||||
|
assert result.checkpoint["channel_values"]["foo"] is None
|
||||||
|
|||||||
@@ -126,7 +126,7 @@ class SerdeConfig(TypedDict, total=False):
|
|||||||
If omitted, no serde is set up (the object store will still be present, however)."""
|
If omitted, no serde is set up (the object store will still be present, however)."""
|
||||||
|
|
||||||
allowed_json_modules: list[list[str]] | bool | None
|
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 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
|
If omitted, no modules are allowed, and the object returned will simply be a json object OR
|
||||||
@@ -146,7 +146,34 @@ class SerdeConfig(TypedDict, total=False):
|
|||||||
Example:
|
Example:
|
||||||
{...
|
{...
|
||||||
"serde": {
|
"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
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -306,8 +333,7 @@ class EncryptionConfig(TypedDict, total=False):
|
|||||||
"""Configuration for custom at-rest encryption logic.
|
"""Configuration for custom at-rest encryption logic.
|
||||||
|
|
||||||
Allows you to implement custom encryption for sensitive data stored in the database,
|
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
|
path: str
|
||||||
"""Required. Path to an instance of the Encryption() class that implements custom encryption handlers.
|
"""Required. Path to an instance of the Encryption() class that implements custom encryption handlers.
|
||||||
|
|||||||
@@ -591,7 +591,27 @@
|
|||||||
"type": "null"
|
"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": {
|
"pickle_fallback": {
|
||||||
"type": "boolean",
|
"type": "boolean",
|
||||||
|
|||||||
@@ -591,7 +591,27 @@
|
|||||||
"type": "null"
|
"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": {
|
"pickle_fallback": {
|
||||||
"type": "boolean",
|
"type": "boolean",
|
||||||
|
|||||||
@@ -1,3 +1,4 @@
|
|||||||
|
import os
|
||||||
from contextlib import asynccontextmanager, contextmanager
|
from contextlib import asynccontextmanager, contextmanager
|
||||||
from uuid import uuid4
|
from uuid import uuid4
|
||||||
|
|
||||||
@@ -5,6 +6,7 @@ import pytest
|
|||||||
from langgraph.checkpoint.postgres import PostgresSaver
|
from langgraph.checkpoint.postgres import PostgresSaver
|
||||||
from langgraph.checkpoint.postgres.aio import AsyncPostgresSaver
|
from langgraph.checkpoint.postgres.aio import AsyncPostgresSaver
|
||||||
from langgraph.checkpoint.serde.encrypted import EncryptedSerializer
|
from langgraph.checkpoint.serde.encrypted import EncryptedSerializer
|
||||||
|
from langgraph.checkpoint.serde.jsonplus import JsonPlusSerializer
|
||||||
from langgraph.checkpoint.sqlite import SqliteSaver
|
from langgraph.checkpoint.sqlite import SqliteSaver
|
||||||
from langgraph.checkpoint.sqlite.aio import AsyncSqliteSaver
|
from langgraph.checkpoint.sqlite.aio import AsyncSqliteSaver
|
||||||
from psycopg import AsyncConnection, Connection
|
from psycopg import AsyncConnection, Connection
|
||||||
@@ -18,30 +20,60 @@ from tests.memory_assert import ( # noqa: E402
|
|||||||
)
|
)
|
||||||
|
|
||||||
DEFAULT_POSTGRES_URI = "postgres://postgres:postgres@localhost:5442/"
|
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
|
@contextmanager
|
||||||
def _checkpointer_memory():
|
def _checkpointer_memory():
|
||||||
yield MemorySaverAssertImmutable()
|
if STRICT_MSGPACK:
|
||||||
|
yield MemorySaverAssertImmutable(serde=_strict_msgpack_serde())
|
||||||
|
else:
|
||||||
|
yield MemorySaverAssertImmutable()
|
||||||
|
|
||||||
|
|
||||||
@contextmanager
|
@contextmanager
|
||||||
def _checkpointer_memory_migrate_sends():
|
def _checkpointer_memory_migrate_sends():
|
||||||
yield MemorySaverNeedsPendingSendsMigration()
|
checkpointer = MemorySaverNeedsPendingSendsMigration()
|
||||||
|
_apply_strict_msgpack(checkpointer)
|
||||||
|
yield checkpointer
|
||||||
|
|
||||||
|
|
||||||
@contextmanager
|
@contextmanager
|
||||||
def _checkpointer_sqlite():
|
def _checkpointer_sqlite():
|
||||||
with SqliteSaver.from_conn_string(":memory:") as checkpointer:
|
with SqliteSaver.from_conn_string(":memory:") as checkpointer:
|
||||||
|
_apply_strict_msgpack(checkpointer)
|
||||||
yield checkpointer
|
yield checkpointer
|
||||||
|
|
||||||
|
|
||||||
@contextmanager
|
@contextmanager
|
||||||
def _checkpointer_sqlite_aes():
|
def _checkpointer_sqlite_aes():
|
||||||
with SqliteSaver.from_conn_string(":memory:") as checkpointer:
|
with SqliteSaver.from_conn_string(":memory:") as checkpointer:
|
||||||
checkpointer.serde = EncryptedSerializer.from_pycryptodome_aes(
|
if STRICT_MSGPACK:
|
||||||
key=b"1234567890123456"
|
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
|
yield checkpointer
|
||||||
|
|
||||||
|
|
||||||
@@ -57,6 +89,7 @@ def _checkpointer_postgres():
|
|||||||
DEFAULT_POSTGRES_URI + database
|
DEFAULT_POSTGRES_URI + database
|
||||||
) as checkpointer:
|
) as checkpointer:
|
||||||
checkpointer.setup()
|
checkpointer.setup()
|
||||||
|
_apply_strict_msgpack(checkpointer)
|
||||||
yield checkpointer
|
yield checkpointer
|
||||||
finally:
|
finally:
|
||||||
# drop unique db
|
# drop unique db
|
||||||
@@ -79,6 +112,7 @@ def _checkpointer_postgres_pipe():
|
|||||||
# setup can't run inside pipeline because of implicit transaction
|
# setup can't run inside pipeline because of implicit transaction
|
||||||
with checkpointer.conn.pipeline() as pipe:
|
with checkpointer.conn.pipeline() as pipe:
|
||||||
checkpointer.pipe = pipe
|
checkpointer.pipe = pipe
|
||||||
|
_apply_strict_msgpack(checkpointer)
|
||||||
yield checkpointer
|
yield checkpointer
|
||||||
finally:
|
finally:
|
||||||
# drop unique db
|
# drop unique db
|
||||||
@@ -99,6 +133,7 @@ def _checkpointer_postgres_pool():
|
|||||||
) as pool:
|
) as pool:
|
||||||
checkpointer = PostgresSaver(pool)
|
checkpointer = PostgresSaver(pool)
|
||||||
checkpointer.setup()
|
checkpointer.setup()
|
||||||
|
_apply_strict_msgpack(checkpointer)
|
||||||
yield checkpointer
|
yield checkpointer
|
||||||
finally:
|
finally:
|
||||||
# drop unique db
|
# drop unique db
|
||||||
@@ -109,6 +144,7 @@ def _checkpointer_postgres_pool():
|
|||||||
@asynccontextmanager
|
@asynccontextmanager
|
||||||
async def _checkpointer_sqlite_aio():
|
async def _checkpointer_sqlite_aio():
|
||||||
async with AsyncSqliteSaver.from_conn_string(":memory:") as checkpointer:
|
async with AsyncSqliteSaver.from_conn_string(":memory:") as checkpointer:
|
||||||
|
_apply_strict_msgpack(checkpointer)
|
||||||
yield checkpointer
|
yield checkpointer
|
||||||
|
|
||||||
|
|
||||||
@@ -126,6 +162,7 @@ async def _checkpointer_postgres_aio():
|
|||||||
DEFAULT_POSTGRES_URI + database
|
DEFAULT_POSTGRES_URI + database
|
||||||
) as checkpointer:
|
) as checkpointer:
|
||||||
await checkpointer.setup()
|
await checkpointer.setup()
|
||||||
|
_apply_strict_msgpack(checkpointer)
|
||||||
yield checkpointer
|
yield checkpointer
|
||||||
finally:
|
finally:
|
||||||
# drop unique db
|
# drop unique db
|
||||||
@@ -152,6 +189,7 @@ async def _checkpointer_postgres_aio_pipe():
|
|||||||
# setup can't run inside pipeline because of implicit transaction
|
# setup can't run inside pipeline because of implicit transaction
|
||||||
async with checkpointer.conn.pipeline() as pipe:
|
async with checkpointer.conn.pipeline() as pipe:
|
||||||
checkpointer.pipe = pipe
|
checkpointer.pipe = pipe
|
||||||
|
_apply_strict_msgpack(checkpointer)
|
||||||
yield checkpointer
|
yield checkpointer
|
||||||
finally:
|
finally:
|
||||||
# drop unique db
|
# drop unique db
|
||||||
@@ -176,6 +214,7 @@ async def _checkpointer_postgres_aio_pool():
|
|||||||
) as pool:
|
) as pool:
|
||||||
checkpointer = AsyncPostgresSaver(pool)
|
checkpointer = AsyncPostgresSaver(pool)
|
||||||
await checkpointer.setup()
|
await checkpointer.setup()
|
||||||
|
_apply_strict_msgpack(checkpointer)
|
||||||
yield checkpointer
|
yield checkpointer
|
||||||
finally:
|
finally:
|
||||||
# drop unique db
|
# drop unique db
|
||||||
|
|||||||
Reference in New Issue
Block a user