WEBDIENSTE/langgraph
1
0
Fork 0
mirror of https://github.com/langchain-ai/langgraph.git synced 2026-08-30 19:59:40 +02:00
Code Issues Packages Projects Releases Wiki Activity

checkpoint: Add support for serializing more python types

Browse Source 
- all ip interfaces
- deque,
- ZoneInfo
- date
- time
This commit is contained in:
Nuno Campos
2024-08-11 19:27:29 -07:00
parent 4e643baaa6
commit 685db62029
2 changed files with 46 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files
libs/checkpoint/langgraph/checkpoint/serde/jsonplus.py
+29 -2
View File
@@ -1,13 +1,23 @@
import dataclasses
import importlib
import json
from datetime import datetime, timedelta, timezone
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
@@ -42,18 +52,35 @@ class JsonPlusSerializer(SerializerProtocol):
return self._encode_constructor_args(obj.__class__, kwargs=obj.dict())
elif isinstance(obj, UUID):
return self._encode_constructor_args(UUID, args=[obj.hex])
elif isinstance(obj, (set, frozenset)):
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__,
libs/checkpoint/tests/test_jsonplus.py
+17 -4
View File
@@ -1,13 +1,16 @@
import dataclasses
import sys
import uuid
from datetime import datetime, timezone
from collections import deque
from datetime import date, datetime, time, timezone
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 +63,21 @@ 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,
"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 +116,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"""{"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) == {