From 723d6f8a8421ae6248af5b3580047fc49f058eaf Mon Sep 17 00:00:00 2001 From: William Fu-Hinthorn <13333726+hinthornw@users.noreply.github.com> Date: Mon, 24 Mar 2025 16:15:48 -0700 Subject: [PATCH] Add json mode msgpack unpacker --- .../langgraph/checkpoint/serde/jsonplus.py | 73 +++++++++++- libs/checkpoint/tests/test_jsonplus.py | 112 +++++++++++++++++- 2 files changed, 183 insertions(+), 2 deletions(-) diff --git a/libs/checkpoint/langgraph/checkpoint/serde/jsonplus.py b/libs/checkpoint/langgraph/checkpoint/serde/jsonplus.py index 619fa2d63..fb2d40d80 100644 --- a/libs/checkpoint/langgraph/checkpoint/serde/jsonplus.py +++ b/libs/checkpoint/langgraph/checkpoint/serde/jsonplus.py @@ -33,6 +33,13 @@ LC_REVIVER = Reviver() class JsonPlusSerializer(SerializerProtocol): + def __init__(self, *, __unpack_ext_hook__: Optional[Callable] = None) -> None: + self._unpack_ext_hook = ( + __unpack_ext_hook__ + if __unpack_ext_hook__ is not None + else _msgpack_ext_hook + ) + def _encode_constructor_args( self, constructor: Union[Callable, type[Any]], @@ -208,7 +215,7 @@ class JsonPlusSerializer(SerializerProtocol): return self.loads(data_) elif type_ == "msgpack": return msgpack.unpackb( - data_, ext_hook=_msgpack_ext_hook, strict_map_key=False + data_, ext_hook=self._unpack_ext_hook, strict_map_key=False ) else: raise NotImplementedError(f"Unknown serialization type: {type_}") @@ -513,5 +520,69 @@ def _msgpack_ext_hook(code: int, data: bytes) -> Any: return +def _msgpack_ext_hook_to_json(code: int, data: bytes) -> Any: + if code == EXT_CONSTRUCTOR_SINGLE_ARG: + try: + tup = msgpack.unpackb( + data, ext_hook=_msgpack_ext_hook_to_json, strict_map_key=False + ) + if tup[0] == "uuid" and tup[1] == "UUID": + hex_ = tup[2] + return ( + f"{hex_[:8]}-{hex_[8:12]}-{hex_[12:16]}-{hex_[16:20]}-{hex_[20:]}" + ) + # module, name, arg + return tup[2] + except Exception: + return + elif code == EXT_CONSTRUCTOR_POS_ARGS: + try: + tup = msgpack.unpackb( + data, ext_hook=_msgpack_ext_hook_to_json, strict_map_key=False + ) + # module, name, args + return tup[2] + except Exception: + return + elif code == EXT_CONSTRUCTOR_KW_ARGS: + try: + tup = msgpack.unpackb( + data, ext_hook=_msgpack_ext_hook_to_json, strict_map_key=False + ) + # module, name, args + return tup[2] + except Exception: + return + elif code == EXT_METHOD_SINGLE_ARG: + try: + tup = msgpack.unpackb( + data, ext_hook=_msgpack_ext_hook_to_json, strict_map_key=False + ) + # module, name, arg, method + return tup[2] + except Exception: + return + elif code == EXT_PYDANTIC_V1: + try: + tup = msgpack.unpackb( + data, ext_hook=_msgpack_ext_hook_to_json, strict_map_key=False + ) + # module, name, kwargs + return tup[2] + except Exception: + # for pydantic objects we can't find/reconstruct + # let's return the kwargs dict instead + return + elif code == EXT_PYDANTIC_V2: + try: + tup = msgpack.unpackb( + data, ext_hook=_msgpack_ext_hook_to_json, strict_map_key=False + ) + # module, name, kwargs, method + return tup[2] + except Exception: + return + + def _msgpack_enc(data: Any) -> bytes: return msgpack.packb(data, default=_msgpack_default) diff --git a/libs/checkpoint/tests/test_jsonplus.py b/libs/checkpoint/tests/test_jsonplus.py index b75962308..34f7b2ec5 100644 --- a/libs/checkpoint/tests/test_jsonplus.py +++ b/libs/checkpoint/tests/test_jsonplus.py @@ -15,7 +15,10 @@ from pydantic.v1 import BaseModel as BaseModelV1 from pydantic.v1 import SecretStr as SecretStrV1 from zoneinfo import ZoneInfo -from langgraph.checkpoint.serde.jsonplus import JsonPlusSerializer +from langgraph.checkpoint.serde.jsonplus import ( + JsonPlusSerializer, + _msgpack_ext_hook_to_json, +) from langgraph.store.base import Item @@ -164,6 +167,113 @@ def test_serde_jsonplus() -> None: ] +def test_serde_jsonplus_json_mode() -> None: + uid = uuid.UUID(int=1) + deque_instance = deque([1, 2, 3]) + tzn = ZoneInfo("America/New_York") + ip4 = IPv4Address("192.168.0.1") + current_date = date(2024, 4, 19) + current_time = time(23, 4, 57, 51022, timezone.max) + current_timestamp = datetime(2024, 4, 19, 23, 4, 57, 51022, timezone.max) + + to_serialize = { + "path": pathlib.Path("foo", "bar"), + "re": re.compile(r"foo", re.DOTALL), + "decimal": Decimal("1.10101"), + "set": {1, 2, frozenset({1, 2})}, + "frozen_set": frozenset({1, 2, 3}), + "ip4": ip4, + "deque": deque_instance, + "tzn": tzn, + "date": current_date, + "time": current_time, + "uid": uid, + "timestamp": current_timestamp, + "my_slotted_class": MyDataclassWSlots("bar", 2, InnerDataclass("hello")), + "my_dataclass": MyDataclass("foo", 1, InnerDataclass("hello")), + "my_enum": MyEnum.FOO, + "my_pydantic": MyPydantic(foo="foo", bar=1, inner=InnerPydantic(hello="hello")), + "my_pydantic_v1": MyPydanticV1( + foo="foo", bar=1, inner=InnerPydanticV1(hello="hello") + ), + "my_secret_str": SecretStr("meow"), + "my_secret_str_v1": SecretStrV1("meow"), + "person": Person(name="foo"), + "a_bool": True, + "a_none": None, + "a_str": "foo", + "a_str_nuc": "foo\u0000", + "a_str_uc": "foo ⛰️", + "a_str_ucuc": "foo \u26f0\ufe0f\u0000", + "a_str_ucucuc": "foo \\u26f0\\ufe0f", + "an_int": 1, + "a_float": 1.1, + "a_bytes": b"my bytes", + "a_bytearray": bytearray([42]), + "my_item": Item( + value={}, + key="my-key", + namespace=("a", "name", " "), + created_at=datetime(2024, 9, 24, 17, 29, 10, 128397), + updated_at=datetime(2024, 9, 24, 17, 29, 11, 128397), + ), + } + + serde = JsonPlusSerializer(__unpack_ext_hook__=_msgpack_ext_hook_to_json) + + dumped = serde.dumps_typed(to_serialize) + + assert dumped[0] == "msgpack" + result = serde.loads_typed(dumped) + assert result == { + "path": ["foo", "bar"], + "re": ["foo", 48], + "decimal": "1.10101", + "set": [1, 2, [1, 2]], + "frozen_set": [1, 2, 3], + "ip4": "192.168.0.1", + "deque": [1, 2, 3], + "tzn": "America/New_York", + "date": [2024, 4, 19], + "time": { + "hour": 23, + "minute": 4, + "second": 57, + "microsecond": 51022, + "tzinfo": [[0, 86340, 0]], + "fold": 0, + }, + "uid": "00000000-0000-0000-0000-000000000001", + "timestamp": "2024-04-19T23:04:57.051022+23:59", + "my_slotted_class": {"foo": "bar", "bar": 2, "inner": {"hello": "hello"}}, + "my_dataclass": {"foo": "foo", "bar": 1, "inner": {"hello": "hello"}}, + "my_enum": "foo", + "my_pydantic": {"foo": "foo", "bar": 1, "inner": {"hello": "hello"}}, + "my_pydantic_v1": {"foo": "foo", "bar": 1, "inner": {"hello": "hello"}}, + "my_secret_str": "meow", + "my_secret_str_v1": "meow", + "person": {"name": "foo"}, + "a_bool": True, + "a_none": None, + "a_str": "foo", + "a_str_nuc": "foo\x00", + "a_str_uc": "foo ⛰️", + "a_str_ucuc": "foo ⛰️\x00", + "a_str_ucucuc": "foo \\u26f0\\ufe0f", + "an_int": 1, + "a_float": 1.1, + "a_bytes": b"my bytes", + "a_bytearray": b"*", + "my_item": { + "namespace": ["a", "name", " "], + "key": "my-key", + "value": {}, + "created_at": "2024-09-24T17:29:10.128397", + "updated_at": "2024-09-24T17:29:11.128397", + }, + } + + def test_serde_jsonplus_bytes() -> None: serde = JsonPlusSerializer()