LangGraph MemorySaver checkpointer usage with MLflow

sebascardonal
Databricks Partner

Hi everyone.

I am working on a graph that utilizes the MemorySaver class to incorporate short-term memory. This will enable me to maintain a multi-turn conversation with the user by storing the chat history.

I am using the MLflow "models from code" feature but I'm getting an error because when the model is invoked it requires the config parameter with a thread_id:

ValueError("Checkpointer requires one or more of the following 'configurable' keys: ['thread_id', 'checkpoint_ns', 'checkpoint_id']")Traceback (most recent call last)
 The graph compilation is:

# Compile

memory = MemorySaver()
graph = graph_builder.compile(checkpointer=memory)

How to register a LangGraph graph in MLflow that uses the MemorySaver to store the chat history in the short-term memory?

Thanks!