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
