Custom serde not used for agent state (metadata column)

[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

[vermapratyush]

cc @nfcampos @vbarda

[dosubot[bot]]

Hey @vermapratyush! 👋 I'm here to help you out with any bugs or questions you have while waiting for a human maintainer. Let's squash those bugs together!

To ensure that a custom serde is used for the metadata column in LangGraph when the agent state contains proto objects, and to avoid the TypeError related to non-JSON serializable objects, you can define a custom serialization and deserialization method for the AgentState. Here is an example of how you can achieve this:

1. Define Custom Serialization and Deserialization Methods:

import json
from langchain_core.messages import BaseMessage

class CustomJSONEncoder(json.JSONEncoder):
    def default(self, obj):
        if isinstance(obj, BaseMessage):
            return obj.to_dict()  # Assuming BaseMessage has a to_dict method
        # Add other custom serialization logic here
        return super().default(obj)

def custom_json_serializer(obj):
    return json.dumps(obj, cls=CustomJSONEncoder)

def custom_json_deserializer(json_str):
    return json.loads(json_str)  # Add custom deserialization logic if needed

2. Use Custom Serde in LangGraph:

When defining the AgentState, you can specify the custom serializer and deserializer:

class AgentState(TypedDict):
    messages: Annotated[Sequence[BaseMessage], operator.add]

# Example of using custom serde in LangGraph
workflow = StateGraph(AgentState, serializer=custom_json_serializer, deserializer=custom_json_deserializer)

3. Ensure Custom Serde is Applied to Metadata Column:

When you are working with the metadata column, ensure that the custom serde methods are used to serialize and deserialize the data:

# Example of serializing metadata
metadata = {"some_key": "some_value", "messages": state["messages"]}
serialized_metadata = custom_json_serializer(metadata)

# Example of deserializing metadata
deserialized_metadata = custom_json_deserializer(serialized_metadata)

By implementing these custom serialization and deserialization methods, you can handle non-JSON serializable objects like proto objects and avoid TypeError ^[1].

It looks like this issue needs some specialized attention, @eyurtsev could you please help out?

To continue talking to Dosu, mention @dosu.

[vermapratyush]

Created a new discussion thread in langgraph repo: langchain-ai/langgraph#1755