mirror of
https://github.com/langchain-ai/langgraph.git
synced 2026-09-03 02:28:44 +02:00
Merge pull request #1311 from langchain-ai/nc/11aug/serde-more
checkpoint: Add support for serializing more python types
This commit is contained in:
Nuno Campos
@@ -1,13 +1,26 @@
|
||||
import dataclasses
|
||||
import decimal
|
||||
import importlib
|
||||
import json
|
||||
from datetime import datetime, timedelta, timezone
|
||||
import pathlib
|
||||
import re
|
||||
from collections import deque
|
||||
from datetime import date, datetime, time, timedelta, timezone
|
||||
from enum import Enum
|
||||
from ipaddress import (
|
||||
IPv4Address,
|
||||
IPv4Interface,
|
||||
IPv4Network,
|
||||
IPv6Address,
|
||||
IPv6Interface,
|
||||
IPv6Network,
|
||||
)
|
||||
from typing import Any, Optional
|
||||
from uuid import UUID
|
||||
|
||||
from langchain_core.load.load import Reviver
|
||||
from langchain_core.load.serializable import Serializable
|
||||
from zoneinfo import ZoneInfo
|
||||
|
||||
from langgraph.checkpoint.serde.base import SerializerProtocol
|
||||
from langgraph.checkpoint.serde.types import SendProtocol
|
||||
@@ -40,20 +53,45 @@ class JsonPlusSerializer(SerializerProtocol):
|
||||
return self._encode_constructor_args(obj.__class__, kwargs=obj.model_dump())
|
||||
elif hasattr(obj, "dict") and callable(obj.dict):
|
||||
return self._encode_constructor_args(obj.__class__, kwargs=obj.dict())
|
||||
elif isinstance(obj, pathlib.Path):
|
||||
return self._encode_constructor_args(pathlib.Path, args=obj.parts)
|
||||
elif isinstance(obj, re.Pattern):
|
||||
return self._encode_constructor_args(
|
||||
re.compile, args=[obj.pattern, obj.flags]
|
||||
)
|
||||
elif isinstance(obj, UUID):
|
||||
return self._encode_constructor_args(UUID, args=[obj.hex])
|
||||
elif isinstance(obj, (set, frozenset)):
|
||||
elif isinstance(obj, decimal.Decimal):
|
||||
return self._encode_constructor_args(decimal.Decimal, args=[str(obj)])
|
||||
elif isinstance(obj, (set, frozenset, deque)):
|
||||
return self._encode_constructor_args(type(obj), args=[list(obj)])
|
||||
elif isinstance(obj, (IPv4Address, IPv4Interface, IPv4Network)):
|
||||
return self._encode_constructor_args(obj.__class__, args=[str(obj)])
|
||||
elif isinstance(obj, (IPv6Address, IPv6Interface, IPv6Network)):
|
||||
return self._encode_constructor_args(obj.__class__, args=[str(obj)])
|
||||
|
||||
elif isinstance(obj, datetime):
|
||||
return self._encode_constructor_args(
|
||||
datetime, method="fromisoformat", args=[obj.isoformat()]
|
||||
)
|
||||
elif isinstance(obj, timezone):
|
||||
return self._encode_constructor_args(timezone, args=obj.__getinitargs__())
|
||||
elif isinstance(obj, ZoneInfo):
|
||||
return self._encode_constructor_args(ZoneInfo, args=[obj.key])
|
||||
elif isinstance(obj, timedelta):
|
||||
return self._encode_constructor_args(
|
||||
timedelta, args=[obj.days, obj.seconds, obj.microseconds]
|
||||
)
|
||||
elif isinstance(obj, date):
|
||||
return self._encode_constructor_args(
|
||||
date, args=[obj.year, obj.month, obj.day]
|
||||
)
|
||||
elif isinstance(obj, time):
|
||||
return self._encode_constructor_args(
|
||||
time,
|
||||
args=[obj.hour, obj.minute, obj.second, obj.microsecond, obj.tzinfo],
|
||||
kwargs={"fold": obj.fold},
|
||||
)
|
||||
elif dataclasses.is_dataclass(obj):
|
||||
return self._encode_constructor_args(
|
||||
obj.__class__,
|
||||
|
||||
@@ -1,13 +1,19 @@
|
||||
import dataclasses
|
||||
import pathlib
|
||||
import re
|
||||
import sys
|
||||
import uuid
|
||||
from datetime import datetime, timezone
|
||||
from collections import deque
|
||||
from datetime import date, datetime, time, timezone
|
||||
from decimal import Decimal
|
||||
from enum import Enum
|
||||
from ipaddress import IPv4Address
|
||||
|
||||
import dataclasses_json
|
||||
from langchain_core.pydantic_v1 import BaseModel as LcBaseModel
|
||||
from langchain_core.runnables import RunnableMap
|
||||
from pydantic import BaseModel
|
||||
from zoneinfo import ZoneInfo
|
||||
|
||||
from langgraph.checkpoint.serde.jsonplus import JsonPlusSerializer
|
||||
|
||||
@@ -60,11 +66,24 @@ class Person:
|
||||
|
||||
def test_serde_jsonplus() -> None:
|
||||
uid = uuid.UUID(int=1)
|
||||
current_time = datetime(2024, 4, 19, 23, 4, 57, 51022, timezone.max)
|
||||
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 = {
|
||||
"uid": uid,
|
||||
"path": pathlib.Path("foo", "bar"),
|
||||
"re": re.compile(r"foo", re.DOTALL),
|
||||
"decimal": Decimal("1.10101"),
|
||||
"ip4": ip4,
|
||||
"deque": deque_instance,
|
||||
"tzn": tzn,
|
||||
"date": current_date,
|
||||
"time": current_time,
|
||||
"uid": uid,
|
||||
"timestamp": current_timestamp,
|
||||
"my_slotted_class": MyDataclassWSlots("bar", 2),
|
||||
"my_dataclass": MyDataclass("foo", 1),
|
||||
"my_enum": MyEnum.FOO,
|
||||
@@ -103,7 +122,7 @@ def test_serde_jsonplus() -> None:
|
||||
|
||||
assert dumped == (
|
||||
"json",
|
||||
b"""{"uid": {"lc": 2, "type": "constructor", "id": ["uuid", "UUID"], "method": null, "args": ["00000000000000000000000000000001"], "kwargs": {}}, "time": {"lc": 2, "type": "constructor", "id": ["datetime", "datetime"], "method": "fromisoformat", "args": ["2024-04-19T23:04:57.051022+23:59"], "kwargs": {}}, "my_slotted_class": {"lc": 2, "type": "constructor", "id": ["tests", "test_jsonplus", "MyDataclassWSlots"], "method": null, "args": [], "kwargs": {"foo": "bar", "bar": 2}}, "my_dataclass": {"lc": 2, "type": "constructor", "id": ["tests", "test_jsonplus", "MyDataclass"], "method": null, "args": [], "kwargs": {"foo": "foo", "bar": 1}}, "my_enum": {"lc": 2, "type": "constructor", "id": ["tests", "test_jsonplus", "MyEnum"], "method": null, "args": ["foo"], "kwargs": {}}, "my_pydantic": {"lc": 2, "type": "constructor", "id": ["tests", "test_jsonplus", "MyPydantic"], "method": null, "args": [], "kwargs": {"foo": "foo", "bar": 1}}, "my_funny_pydantic": {"lc": 2, "type": "constructor", "id": ["tests", "test_jsonplus", "MyFunnyPydantic"], "method": null, "args": [], "kwargs": {"foo": "foo", "bar": 1}}, "person": {"lc": 2, "type": "constructor", "id": ["tests", "test_jsonplus", "Person"], "method": null, "args": [], "kwargs": {"name": "foo"}}, "a_bool": true, "a_none": null, "a_str": "foo", "a_str_nuc": "foo\\u0000", "a_str_uc": "foo \xe2\x9b\xb0\xef\xb8\x8f", "a_str_ucuc": "foo \xe2\x9b\xb0\xef\xb8\x8f\\u0000", "a_str_ucucuc": "foo \\\\u26f0\\\\ufe0f", "text": ["Hello", "Python", "Surrogate", "Example", "String", "With", "Surrogates", "Embedded", "In", "The", "Text", "\xe6\x94\xb6\xe8\x8a\xb1\xf0\x9f\x99\x84\xc2\xb7\xe5\x88\xb0"], "an_int": 1, "a_float": 1.1, "runnable_map": {"lc": 1, "type": "constructor", "id": ["langchain", "schema", "runnable", "RunnableParallel"], "kwargs": {"steps__": {}}, "name": "RunnableParallel<>", "graph": {"nodes": [{"id": 0, "type": "schema", "data": "Parallel<>Input"}, {"id": 1, "type": "schema", "data": "Parallel<>Output"}], "edges": []}}}""",
|
||||
b"""{"path": {"lc": 2, "type": "constructor", "id": ["pathlib", "Path"], "method": null, "args": ["foo", "bar"], "kwargs": {}}, "re": {"lc": 2, "type": "constructor", "id": ["re", "compile"], "method": null, "args": ["foo", 48], "kwargs": {}}, "decimal": {"lc": 2, "type": "constructor", "id": ["decimal", "Decimal"], "method": null, "args": ["1.10101"], "kwargs": {}}, "ip4": {"lc": 2, "type": "constructor", "id": ["ipaddress", "IPv4Address"], "method": null, "args": ["192.168.0.1"], "kwargs": {}}, "deque": {"lc": 2, "type": "constructor", "id": ["collections", "deque"], "method": null, "args": [[1, 2, 3]], "kwargs": {}}, "tzn": {"lc": 2, "type": "constructor", "id": ["zoneinfo", "ZoneInfo"], "method": null, "args": ["America/New_York"], "kwargs": {}}, "date": {"lc": 2, "type": "constructor", "id": ["datetime", "date"], "method": null, "args": [2024, 4, 19], "kwargs": {}}, "time": {"lc": 2, "type": "constructor", "id": ["datetime", "time"], "method": null, "args": [23, 4, 57, 51022, {"lc": 2, "type": "constructor", "id": ["datetime", "timezone"], "method": null, "args": [{"lc": 2, "type": "constructor", "id": ["datetime", "timedelta"], "method": null, "args": [0, 86340, 0], "kwargs": {}}], "kwargs": {}}], "kwargs": {"fold": 0}}, "uid": {"lc": 2, "type": "constructor", "id": ["uuid", "UUID"], "method": null, "args": ["00000000000000000000000000000001"], "kwargs": {}}, "timestamp": {"lc": 2, "type": "constructor", "id": ["datetime", "datetime"], "method": "fromisoformat", "args": ["2024-04-19T23:04:57.051022+23:59"], "kwargs": {}}, "my_slotted_class": {"lc": 2, "type": "constructor", "id": ["tests", "test_jsonplus", "MyDataclassWSlots"], "method": null, "args": [], "kwargs": {"foo": "bar", "bar": 2}}, "my_dataclass": {"lc": 2, "type": "constructor", "id": ["tests", "test_jsonplus", "MyDataclass"], "method": null, "args": [], "kwargs": {"foo": "foo", "bar": 1}}, "my_enum": {"lc": 2, "type": "constructor", "id": ["tests", "test_jsonplus", "MyEnum"], "method": null, "args": ["foo"], "kwargs": {}}, "my_pydantic": {"lc": 2, "type": "constructor", "id": ["tests", "test_jsonplus", "MyPydantic"], "method": null, "args": [], "kwargs": {"foo": "foo", "bar": 1}}, "my_funny_pydantic": {"lc": 2, "type": "constructor", "id": ["tests", "test_jsonplus", "MyFunnyPydantic"], "method": null, "args": [], "kwargs": {"foo": "foo", "bar": 1}}, "person": {"lc": 2, "type": "constructor", "id": ["tests", "test_jsonplus", "Person"], "method": null, "args": [], "kwargs": {"name": "foo"}}, "a_bool": true, "a_none": null, "a_str": "foo", "a_str_nuc": "foo\\u0000", "a_str_uc": "foo \xe2\x9b\xb0\xef\xb8\x8f", "a_str_ucuc": "foo \xe2\x9b\xb0\xef\xb8\x8f\\u0000", "a_str_ucucuc": "foo \\\\u26f0\\\\ufe0f", "text": ["Hello", "Python", "Surrogate", "Example", "String", "With", "Surrogates", "Embedded", "In", "The", "Text", "\xe6\x94\xb6\xe8\x8a\xb1\xf0\x9f\x99\x84\xc2\xb7\xe5\x88\xb0"], "an_int": 1, "a_float": 1.1, "runnable_map": {"lc": 1, "type": "constructor", "id": ["langchain", "schema", "runnable", "RunnableParallel"], "kwargs": {"steps__": {}}, "name": "RunnableParallel<>", "graph": {"nodes": [{"id": 0, "type": "schema", "data": "Parallel<>Input"}, {"id": 1, "type": "schema", "data": "Parallel<>Output"}], "edges": []}}}""",
|
||||
)
|
||||
|
||||
assert serde.loads_typed(dumped) == {
|
||||
|
||||
Reference in New Issue
Block a user