LangGraph核心模块：使用 MemorySaver 进行持久化管理

使用 MemorySaver 进行持久化管理

在构建复杂的 AI 应用时，持久化管理是一个非常重要的功能，它能让我们的应用在会话过程中保持状态，甚至在应用重启后恢复之前的对话记忆。本文将详细讲解如何使用 LangGraph 框架中的 MemorySaver 来实现对话状态的持久化管理。

1. 基本概念

MemorySaver 是 LangGraph 框架提供的一种内存式状态持久化工具，它允许我们在同一应用程序执行期间保存和恢复会话状态，非常适合实现多轮对话中的状态保持。

2. 项目结构概述

我们将构建一个具有以下功能的对话式应用：

• 使用 OpenAI 的语言模型
• 定义一个天气查询工具
• 使用 LangGraph 构建对话流程图
• 通过 MemorySaver 实现会话持久化

3. 环境配置与依赖导入

首先，我们需要导入必要的模块和类型注解：

import os
from typing import Annotated, Literal
from dotenv import load_dotenv
from typing_extensions import TypedDict
from langgraph.graph.message import add_messages

4. 定义状态结构

我们使用 TypedDict 定义一个状态类，包含一个带有特殊注解的消息列表：

class State(TypedDict):
    messages: Annotated[list, add_messages]

这里的 add_messages 注解告诉 LangGraph 如何处理这个列表中的新消息。

5. 创建工具函数

接下来，我们定义一个简单的天气查询工具：

from langchain_core.tools import tool

@tool
def get_weather(region: str) -> str:
    """
    获取指定地区的天气
    :param region: 地区
    :return: 天气信息
    """
    return f"{region}晴天，25度"

# 将工具函数存入列表
tools = [get_weather]

6. 初始化语言模型

我们使用 ChatOpenAI 模型，并将工具绑定到模型上：

from langchain_openai import ChatOpenAI

load_dotenv()
# 创建一个 ChatOpenAI 模型实例
llm = ChatOpenAI(base_url="https://ark.cn-beijing.volces.com/api/v3", 
                api_key=os.getenv("OB_OPENAI_KEY"),
                model="doubao-1.5-lite-32k-250115")

# 将工具绑定到模型上
bound_model = llm.bind_tools(tools)

7. 创建节点逻辑

我们需要定义图中的节点和它们的行为：

def call_model(state: State):
    response = bound_model.invoke(state["messages"])
    # 返回一个列表，因为这将被添加到现有列表中
    return {"messages": response}

def should_continue(state: State) -> Literal["tools", "__end__"]:
    messages = state["messages"]
    last_message = messages[-1]
    # 如果最后一条消息包含工具调用，则返回 "tools"
    if last_message.tool_calls:
        return "tools"
    # 否则返回 "__end__" 以结束流程
    return "__end__"

from langgraph.prebuilt import ToolNode

# 创建一个 ToolNode 实例
tool_node = ToolNode(tools)

这里定义了两个关键函数：

• call_model: 调用语言模型处理当前状态
• should_continue: 决定下一步的流向

8. 构建状态图

使用 LangGraph 的 StateGraph 构建对话流程：

from langgraph.graph import StateGraph, MessagesState, START

# 创建一个状态图，定义状态和转换
graph = StateGraph(State)

# 添加节点
graph.add_node("tools", tool_node)
graph.add_node("agent", call_model)

# 编写链接顺序
graph.add_edge(START, "agent")
graph.add_conditional_edges("agent", should_continue)
graph.add_edge("tools", "agent")

9. 使用 MemorySaver 进行持久化管理

这是本文的核心部分，我们引入 MemorySaver 并将其应用到图的编译过程中：

from langgraph.checkpoint.memory import MemorySaver

# 创建一个 MemorySaver 实例
memory = MemorySaver()

# 编译工作流，生成一个 LangChain Runnable，并传入 MemorySaver 实例
app = graph.compile(checkpointer=memory)

# 将生成的图保存到文件（可选）
graph_png = app.get_graph().draw_mermaid_png()
with open("persistence_case.png", "wb") as f:
    f.write(graph_png)

10. 运行测试并验证持久化效果

最后，我们运行应用程序并测试持久化效果：

from langchain_core.messages import HumanMessage

# 设置配置参数，thread_id作为线程ID，做会话隔离
config = {"configurable": {"thread_id": "2"}}

# 第一轮对话：询问天气
input_message = HumanMessage(content="今天上海的天气如何？")

for event in app.stream({"messages": [input_message]}, config, stream_mode="values"):
    event["messages"][-1].pretty_print()

# 第二轮对话：验证记忆功能
input_message = HumanMessage(content="我刚问你了什么问题？")

for event in app.stream({"messages": [input_message]}, config, stream_mode="values"):
    event["messages"][-1].pretty_print()

11. 输出结果分析

执行上面的代码，我们可以看到以下输出：

================================ Human Message =================================
今天上海的天气如何？
================================== Ai Message ==================================
用户想了解上海的天气，调用get_weather函数获取相关信息。
Tool Calls:
  get_weather (call_tdbiokyoeqxwb5rurplu6z0l)
 Call ID: call_tdbiokyoeqxwb5rurplu6z0l
  Args:
    region: 上海
================================= Tool Message =================================
Name: get_weather
上海晴天，25度
================================== Ai Message ==================================
上海今天的天气是晴天，气温为25度。
================================ Human Message =================================
我刚问你了什么问题？
================================== Ai Message ==================================
你问的问题是：今天上海的天气如何？

从输出结果可以看出，AI 成功地记住了之前的交互内容。当我们问"我刚问你了什么问题？"时，它能够回忆起我们之前询问了上海的天气。这证明了 MemorySaver 成功地保存了会话状态。

12. MemorySaver 的工作原理

MemorySaver 的工作原理如下：

1. 状态存储: 它将每个会话的状态存储在内存中
2. 线程隔离: 通过 thread_id 区分不同的会话
3. 上下文维护: 在整个对话过程中保持上下文信息
4. 会话恢复: 允许后续请求基于之前的状态继续对话

13. MemorySaver 与其他持久化方案的比较

持久化方案	优点	缺点
MemorySaver	实现简单，速度快	应用重启后数据丢失
FileSaver	可跨应用重启保存数据	读写文件带来额外开销
数据库存储	可靠性高，适合生产环境	实现复杂，需要额外依赖

14. 总结与应用场景

通过使用 MemorySaver，我们成功实现了对话状态的持久化管理，使得 AI 能够在多轮对话中保持上下文记忆。这种方案特别适合：

1. 开发阶段的原型测试
2. 单会话临时应用
3. 不需要长期存储的简单应用场景

对于需要更可靠持久化的生产环境，可以考虑使用 FileSaver 或数据库存储方案。

15. 进阶应用

如果希望实现更复杂的持久化功能，可以考虑：

1. 实现自定义 Checkpointer
2. 将状态存储在数据库中
3. 添加过期策略以管理内存使用

通过 LangGraph 框架和 MemorySaver，我们可以轻松构建具有持久化对话功能的智能应用，为用户提供连贯一致的交互体验。