serde: Handle unicode issues when encoding

- any invalid utf-8 chars now removed on dumps
- fix serialization of Send
- remove serialization of NamedTuple, which doesn't work
- add test for custom serde passed to memory saver
- add test using Send and JsonPlus serde

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
+        )

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)

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)

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"]],
+    }

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"}}