diff --git a/langgraph/constants.py b/langgraph/constants.py index 35e4bbd1a..89612ae9e 100644 --- a/langgraph/constants.py +++ b/langgraph/constants.py @@ -1,4 +1,4 @@ -from typing import Any, NamedTuple +from typing import Any CONFIG_KEY_SEND = "__pregel_send" CONFIG_KEY_READ = "__pregel_read" @@ -11,6 +11,23 @@ RESERVED = {INTERRUPT, TASKS, CONFIG_KEY_SEND, CONFIG_KEY_READ} TAG_HIDDEN = "langsmith:hidden" -class Send(NamedTuple): +class Send: node: str arg: Any + + def __init__(self, /, node: str, arg: Any) -> None: + self.node = node + self.arg = arg + + def __hash__(self) -> int: + return hash((self.node, self.arg)) + + def __repr__(self) -> str: + return f"Send(node={self.node!r}, arg={self.arg!r})" + + def __eq__(self, value: object) -> bool: + return ( + isinstance(value, Send) + and self.node == value.node + and self.arg == value.arg + ) diff --git a/langgraph/serde/jsonplus.py b/langgraph/serde/jsonplus.py index 6d539950d..d28e6b640 100644 --- a/langgraph/serde/jsonplus.py +++ b/langgraph/serde/jsonplus.py @@ -3,12 +3,13 @@ import importlib import json from datetime import datetime, timedelta, timezone from enum import Enum -from typing import Any, NamedTuple, Optional +from typing import Any, Optional from uuid import UUID from langchain_core.load.load import Reviver from langchain_core.load.serializable import Serializable +from langgraph.constants import Send from langgraph.serde.base import SerializerProtocol LC_REVIVER = Reviver() @@ -63,8 +64,10 @@ class JsonPlusSerializer(SerializerProtocol): ) elif isinstance(obj, Enum): return self._encode_constructor_args(obj.__class__, args=[obj.value]) - elif isinstance(obj, NamedTuple): - return self._encode_constructor_args(obj.__class__, args=[*obj]) + elif isinstance(obj, Send): + return self._encode_constructor_args( + obj.__class__, kwargs={"node": obj.node, "arg": obj.arg} + ) else: raise TypeError( f"Object of type {obj.__class__.__name__} is not JSON serializable" @@ -92,7 +95,9 @@ class JsonPlusSerializer(SerializerProtocol): return LC_REVIVER(value) def dumps(self, obj: Any) -> bytes: - return json.dumps(obj, default=self._default, sort_keys=True).encode() + return json.dumps(obj, default=self._default, ensure_ascii=False).encode( + "utf-8", "ignore" + ) def loads(self, data: bytes) -> Any: return json.loads(data, object_hook=self._reviver) diff --git a/tests/memory_assert.py b/tests/memory_assert.py index 576fde89a..db1212ae9 100644 --- a/tests/memory_assert.py +++ b/tests/memory_assert.py @@ -43,9 +43,12 @@ class MemorySaverAssertImmutable(MemorySaver): # assert checkpoint hasn't been modified since last written thread_id = config["configurable"]["thread_id"] if saved := super().get(config): - assert self.storage_for_copies[thread_id][saved["id"]] == saved - self.storage_for_copies[thread_id][checkpoint["id"]] = copy_checkpoint( - checkpoint + assert ( + self.serde.loads(self.storage_for_copies[thread_id][saved["id"]]) + == saved + ) + self.storage_for_copies[thread_id][checkpoint["id"]] = self.serde.dumps( + copy_checkpoint(checkpoint) ) # call super to write checkpoint return super().put(config, checkpoint, metadata) diff --git a/tests/test_jsonplus.py b/tests/test_jsonplus.py index c33cfdf92..d2df4c747 100644 --- a/tests/test_jsonplus.py +++ b/tests/test_jsonplus.py @@ -74,6 +74,24 @@ def test_serde_jsonplus() -> None: "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", + "text": [ + "Hello\ud83d\ude00", + "Python\ud83d\udc0d", + "Surrogate\ud834\udd1e", + "Example\ud83c\udf89", + "String\ud83c\udfa7", + "With\ud83c\udf08", + "Surrogates\ud83d\ude0e", + "Embedded\ud83d\udcbb", + "In\ud83c\udf0e", + "The\ud83d\udcd6", + "Text\ud83d\udcac", + "收花🙄·到", + ], "an_int": 1, "a_float": 1.1, "runnable_map": RunnableMap({}), @@ -85,7 +103,10 @@ def test_serde_jsonplus() -> None: assert ( dumped - == b"""{"a_bool": true, "a_float": 1.1, "a_none": null, "a_str": "foo", "an_int": 1, "my_dataclass": {"args": [], "id": ["tests", "test_jsonplus", "MyDataclass"], "kwargs": {"bar": 1, "foo": "foo"}, "lc": 2, "method": null, "type": "constructor"}, "my_enum": {"args": ["foo"], "id": ["tests", "test_jsonplus", "MyEnum"], "kwargs": {}, "lc": 2, "method": null, "type": "constructor"}, "my_funny_pydantic": {"args": [], "id": ["tests", "test_jsonplus", "MyFunnyPydantic"], "kwargs": {"bar": 1, "foo": "foo"}, "lc": 2, "method": null, "type": "constructor"}, "my_pydantic": {"args": [], "id": ["tests", "test_jsonplus", "MyPydantic"], "kwargs": {"bar": 1, "foo": "foo"}, "lc": 2, "method": null, "type": "constructor"}, "my_slotted_class": {"args": [], "id": ["tests", "test_jsonplus", "MyDataclassWSlots"], "kwargs": {"bar": 2, "foo": "bar"}, "lc": 2, "method": null, "type": "constructor"}, "person": {"args": [], "id": ["tests", "test_jsonplus", "Person"], "kwargs": {"name": "foo"}, "lc": 2, "method": null, "type": "constructor"}, "runnable_map": {"graph": {"edges": [], "nodes": [{"data": "Parallel<>Input", "id": 0, "type": "schema"}, {"data": "Parallel<>Output", "id": 1, "type": "schema"}]}, "id": ["langchain", "schema", "runnable", "RunnableParallel"], "kwargs": {"steps__": {}}, "lc": 1, "name": "RunnableParallel<>", "type": "constructor"}, "time": {"args": ["2024-04-19T23:04:57.051022+23:59"], "id": ["datetime", "datetime"], "kwargs": {}, "lc": 2, "method": "fromisoformat", "type": "constructor"}, "uid": {"args": ["00000000000000000000000000000001"], "id": ["uuid", "UUID"], "kwargs": {}, "lc": 2, "method": null, "type": "constructor"}}""" + == 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": []}}}""" ) - assert serde.loads(dumped) == to_serialize + assert serde.loads(dumped) == { + **to_serialize, + "text": [v.encode("utf-8", "ignore").decode() for v in to_serialize["text"]], + } diff --git a/tests/test_pregel.py b/tests/test_pregel.py index 5b7cbcff0..52f81dcd8 100644 --- a/tests/test_pregel.py +++ b/tests/test_pregel.py @@ -48,10 +48,13 @@ from langgraph.prebuilt.chat_agent_executor import ( from langgraph.prebuilt.tool_node import ToolNode from langgraph.pregel import Channel, GraphRecursionError, Pregel, StateSnapshot from langgraph.pregel.retry import RetryPolicy +from langgraph.serde.base import SerializerProtocol +from langgraph.serde.jsonplus import JsonPlusSerializer from tests.any_str import AnyStr from tests.memory_assert import ( MemorySaverAssertCheckpointMetadata, MemorySaverAssertImmutable, + NoopSerializer, ) @@ -3433,7 +3436,8 @@ def test_prebuilt_chat(snapshot: SnapshotAssertion) -> None: ] -def test_state_graph_packets() -> None: +@pytest.mark.parametrize("serde", [NoopSerializer(), JsonPlusSerializer()]) +def test_state_graph_packets(serde: SerializerProtocol) -> None: from langchain_core.language_models.fake_chat_models import ( FakeMessagesListChatModel, ) @@ -3666,7 +3670,7 @@ def test_state_graph_packets() -> None: ] app_w_interrupt = workflow.compile( - checkpointer=MemorySaverAssertImmutable(), + checkpointer=MemorySaverAssertImmutable(serde=serde), interrupt_after=["agent"], ) config = {"configurable": {"thread_id": "1"}}