cancel
Showing results forย 
Search instead forย 
Did you mean:ย 
Generative AI
Explore discussions on generative artificial intelligence techniques and applications within the Databricks Community. Share ideas, challenges, and breakthroughs in this cutting-edge field.
cancel
Showing results forย 
Search instead forย 
Did you mean:ย 

LangGraph MemorySaver checkpointer usage with MLflow

sebascardonal
New Contributor II

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!

2 REPLIES 2

Pawan1979
New Contributor II

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

In today's tutorial, we're going to add memory to the LangGraph ReAct agent using the Tavily tool to get a web connection and an OpenAI LLM that we built in the previous LangChain tutorial. We'll use LangGraph's MemorySaver class to implement checkpointers, which is a way to add in-memory storage

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!

Connect with Databricks Users in Your Area

Join a Regional User Group to connect with local Databricks users. Events will be happening in your city, and you wonโ€™t want to miss the chance to attend and share knowledge.

If there isnโ€™t a group near you, start one and help create a community that brings people together.

Request a New Group