From b3ea406e813f1e17eae90b6fbc8a28fcdfa221f1 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Wed, 30 Apr 2025 18:01:31 -0700 Subject: [PATCH] Add pickle_fallback for json plus serializer --- libs/checkpoint/langgraph/cache/base/__init__.py | 2 +- .../checkpoint/langgraph/checkpoint/serde/jsonplus.py | 11 ++++++++++- 2 files changed, 11 insertions(+), 2 deletions(-) diff --git a/libs/checkpoint/langgraph/cache/base/__init__.py b/libs/checkpoint/langgraph/cache/base/__init__.py index a81b48cf6..9c8653d23 100644 --- a/libs/checkpoint/langgraph/cache/base/__init__.py +++ b/libs/checkpoint/langgraph/cache/base/__init__.py @@ -11,7 +11,7 @@ T = TypeVar("T") class BaseCache(ABC, Generic[T]): """Base class for a cache.""" - serde: SerializerProtocol = JsonPlusSerializer() + serde: SerializerProtocol = JsonPlusSerializer(pickle_fallback=True) def __init__(self, *, serde: SerializerProtocol | None = None) -> None: """Initialize the cache with a serializer.""" diff --git a/libs/checkpoint/langgraph/checkpoint/serde/jsonplus.py b/libs/checkpoint/langgraph/checkpoint/serde/jsonplus.py index 3a9ab50e2..53b96065c 100644 --- a/libs/checkpoint/langgraph/checkpoint/serde/jsonplus.py +++ b/libs/checkpoint/langgraph/checkpoint/serde/jsonplus.py @@ -3,6 +3,7 @@ import decimal import importlib import json import pathlib +import pickle import re from collections import deque from collections.abc import Sequence @@ -37,8 +38,12 @@ class JsonPlusSerializer(SerializerProtocol): """Serializer that uses ormsgpack, with a fallback to extended JSON serializer.""" def __init__( - self, *, __unpack_ext_hook__: Optional[Callable[[int, bytes], Any]] = None + self, + *, + pickle_fallback: bool = False, + __unpack_ext_hook__: Optional[Callable[[int, bytes], Any]] = None, ) -> None: + self.pickle_fallback = pickle_fallback self._unpack_ext_hook = ( __unpack_ext_hook__ if __unpack_ext_hook__ is not None @@ -209,6 +214,8 @@ class JsonPlusSerializer(SerializerProtocol): except ormsgpack.MsgpackEncodeError as exc: if "valid UTF-8" in str(exc): return "json", self.dumps(obj) + elif self.pickle_fallback: + return "pickle", pickle.dumps(obj) raise exc def loads(self, data: bytes) -> Any: @@ -228,6 +235,8 @@ class JsonPlusSerializer(SerializerProtocol): return ormsgpack.unpackb( data_, ext_hook=self._unpack_ext_hook, option=ormsgpack.OPT_NON_STR_KEYS ) + elif type_ == "pickle": + return pickle.loads(data_) else: raise NotImplementedError(f"Unknown serialization type: {type_}")