Using custom serde for metadata as well

[vermapratyush]

Hi,
I am using langgraph with a custom serde for checkpointer. The custom serde knows how to serialize and deserialize a proto object.
I am invoking the graph like this:

async for event in graph.astream_events(
                agent_state,
                version="v2",
                config={
                    "recursion_limit": 10,
                    "configurable": {"thread_id": msg.chat_session_id},
                },
            ):

(see logs below)

I see that the metadata column (in checkpoints table), the first entry contains a serialised version of the agent state. Subsequent rows (for the graph invocation) only contain the delta of the agent state. Meaning, they only contain fields which which are returned by the node.
This is all good for the first message, however on subsequent messages when agent_state is not empty (and contains proto objects), the serializer for metadata column won't be any good.

I created a custom build of with the following changes and it worked file.
file: langgraph/checkpoint/postgres/base.py

218c217
<         return self.serde.loads(self.serde.dumps(metadata))
---
>         return self.jsonplus_serde.loads(self.serde.dumps(metadata))
221c220,222
<         serialized_metadata_type, serialized_metadata = self.serde.dumps_typed(metadata)
---
>         serialized_metadata_type, serialized_metadata = self.jsonplus_serde.dumps_typed(
>             metadata
>         )

But with a custom serde, the filtering function _search_where may not work as expected.

Logs:

    await self.checkpointer.aput(config, checkpoint, metadata, new_versions)
  File "/usr/local/lib/python3.12/site-packages/langgraph/checkpoint/postgres/aio.py", line 275, in aput
    self._dump_metadata(metadata),
    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/site-packages/langgraph/checkpoint/postgres/base.py", line 220, in _dump_metadata
    serialized_metadata_type, serialized_metadata = self.jsonplus_serde.dumps_typed(
                                                    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/site-packages/langgraph/checkpoint/serde/jsonplus.py", line 178, in dumps_typed
    return "json", self.dumps(obj)
                   ^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/site-packages/langgraph/checkpoint/serde/jsonplus.py", line 168, in dumps
    return json.dumps(obj, default=self._default, ensure_ascii=False).encode(
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/json/__init__.py", line 238, in dumps
    **kw).encode(obj)
          ^^^^^^^^^^^
  File "/usr/local/lib/python3.12/json/encoder.py", line 200, in encode
    chunks = self.iterencode(o, _one_shot=True)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/json/encoder.py", line 258, in iterencode
    return _iterencode(o, 0)
           ^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/site-packages/langgraph/checkpoint/serde/jsonplus.py", line 121, in _default
    raise TypeError(
TypeError: Object of type MetricsQueryMessage is not JSON serializable

[rcwylie]

Hi @vermapratyush,

I think I'm experiencing a similar or related problem. I'm using the Postgres checkpoint saver and receive serialization warnings/errors which I think is linked to your comments above.

2024-09-28 14:35:43,705 DEBUG:Orjson serialization failed: TypeError('Type is not JSON serializable: frozenset'). Falling back to json.
2024-09-28 14:35:43,705 DEBUG:Failed to serialize <class 'frozenset'> to JSON: Object of type frozenset is not JSON serializable
2024-09-28 14:35:43,706 DEBUG:Orjson serialization failed: TypeError('Type is not JSON serializable: CaseInsensitiveDict'). Falling back to json.
2024-09-28 14:35:43,706 DEBUG:Failed to serialize <class 'requests.structures.CaseInsensitiveDict'> to JSON: Object of type CaseInsensitiveDict is not JSON serializable

Do you have any advice for resolving these?

Many thanks

[vermapratyush]

I am using my own checkpointing implementation for now.