topic Re: LangGraph MemorySaver checkpointer usage with MLflow in Generative AI

LangGraph MemorySaver checkpointer usage with MLflow

sebascardonal — Tue, 11 Feb 2025 21:38:54 GMT

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!

Re: LangGraph MemorySaver checkpointer usage with MLflow

Pawan1979 — Wed, 12 Feb 2025 10:59:10 GMT

To register a LangGraph graph in MLflow using MemorySaver for short-term memory, you can follow these steps:

1. **Set up MemorySaver:** Create a `MemorySaver` checkpointer to enable persistent checkpointing, which saves the state of the graph after each step. This allows the chatbot to remember the context of previous interactions[3][5].

```python
from langgraph.checkpoint.memory import MemorySaver
memory = MemorySaver()
```
2. **Compile the graph with the checkpointer:** When compiling the graph, provide the `MemorySaver` checkpointer. Also, specify a `thread_id` when calling your graph, which allows LangGraph to load the saved state when the graph is invoked again[3].

```python
from langgraph.checkpoint.memory import MemorySaver

# We need this because we want to enable threads (conversations)
checkpointer = MemorySaver()

# ... Define the graph ...

# Compile the graph with the checkpointer and store
graph = graph.compile(checkpointer=checkpointer, store=in_memory_store)
```

When you invoke the graph, use a `thread_id` and a `user_id` to namespace memories to a particular user[2].

```python
# Invoke the graph
user_id = "1"
config = {"configurable": {"thread_id": "1", "user_id": user_id}}

# First let's just say hi to the AI
for update in graph.stream(
{"messages": [{"role": "user", "content": "hi"}]}, config, stream_mode="updates"
😞
print(update)
```
3. **Define the LangGraph agent in a script:** Create a Python script (e.g., `langgraph.py`) that defines your LangGraph agent[1].
4. **Log the LangGraph model to MLflow:** Use `mlflow.langchain.log_model` to log the LangGraph agent to MLflow. Specify the path to the script defining the agent, the artifact path, and an input example[1].

```python
import mlflow

input_example = {
"messages": [{"role": "user", "content": "what is the weather in seattle today?"}]
}

with mlflow.start_run():
model_info = mlflow.langchain.log_model(
lc_model="./langgraph.py", # specify the path to the LangGraph agent script definition
artifact_path="langgraph",
input_example=input_example,
)
```
5. **Load the model from MLflow and use it for inference:** Load the LangGraph agent from MLflow using `mlflow.langchain.load_model` and then invoke it with a query[1].

```python
agent = mlflow.langchain.load_model(model_info.model_uri)
query = {
"messages": [
{
"role": "user",
"content": "Should I bring an umbrella today when I go to work in San Francisco?",
}
]
}
agent.invoke(query)
```

By using MemorySaver, LangGraph can store and recall information between conversations, allowing the agent to learn from feedback and adapt to user preferences[8].

Citations:
[1] https://mlflow.org/docs/latest/llms/langchain/index.html
[2] https://langchain-ai.github.io/langgraph/concepts/persistence/
[3] https://langchain-ai.github.io/langgraph/tutorials/introduction/
[4] https://mlflow.org/docs/latest/model-registry.html
[5] https://www.youtube.com/watch?v=GMaGG8UBek8
[6] https://langchain-ai.github.io/langgraph/how-tos/memory/manage-conversation-history/
[7] https://mlflow.org/blog/langgraph-model-from-code
[8] https://blog.langchain.dev/launching-long-term-memory-support-in-langgraph/
[9] https://www.getzep.com/ai-agents/langgraph-tutorial
[10] https://mlflow.org/docs/latest/python_api/mlflow.langchain.html
[11] https://github.com/langchain-ai/langgraph/discussions/352

Re: LangGraph MemorySaver checkpointer usage with MLflow

sebascardonal — Wed, 12 Feb 2025 14:12:43 GMT

Hi Pawan.

Thanks for your answer. However, the issue is more related to MLflow than LangGraph. My graph is working perfectly in the notebook, but I need to put it into production and create the API, and to do that, I need to register the graph and be able to use it.

The question is, how MLflow or Databricks would know that a threat_id is needed for every chat session? How should I register the graph to be able to use the short-term memory in the API?

Thanks!

Re: LangGraph MemorySaver checkpointer usage with MLflow

moemedina — Wed, 19 Mar 2025 14:53:21 GMT

Hi @sebascardonal, I am having the same issue. Did you find a way to proceed?

Re: LangGraph MemorySaver checkpointer usage with MLflow

sebascardonal — Wed, 19 Mar 2025 18:24:41 GMT

Hi @moemedina. No, I didn't.

I'm considering using ChatModel/ChatAgent class to wrap the graph and be able to move on. However, the MLflow documentation is still referring to ChatModel where Chat Agent is the latest recommendation:

MLflow ChatModel Doc: https://mlflow.org/docs/latest/llms/chat-model-intro/
Databricks Doc: https://docs.databricks.com/aws/en/generative-ai/agent-framework/author-agent
ChatAgent API Doc: https://mlflow.org/docs/latest/api_reference/python_api/mlflow.pyfunc.html#mlflow.pyfunc.ChatAgent