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 short-term 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 ...