LangGraph从新手到老师傅 - 8 - Checkpoint 基本会话持久化

在钱塘江
461 阅读7分钟

前言

在构建复杂的AI应用时，会话持久化是一个关键需求。想象一下，当用户与你的聊天机器人进行多轮对话时，或者当你的应用需要在多次调用之间保持状态时，如何确保这些状态不会丢失？LangGraph提供了一个强大的解决方案：Checkpoint功能。

什么是 Checkpoint？

Checkpoint 是 LangGraph 中的一个核心功能，用于保存和恢复图执行状态。它就像是程序执行过程中的一个"快照"，可以让你：

  • 在多次调用之间保存状态
  • 在不同会话之间隔离状态
  • 在程序重启后恢复之前的执行状态
  • 实现复杂的会话管理逻辑

在这篇文章中，我们将通过一个简单但实用的示例，来学习如何使用 LangGraph 的 Checkpoint 功能实现基本的会话持久化。

示例功能概述

我们将创建一个简单的计数器应用，展示 Checkpoint 的基本用法：

  • 每次调用图时，计数器的值会增加1
  • 使用 Checkpoint 保存计数器的状态
  • 通过不同的配置区分不同的会话
  • 展示如何查看和管理保存的会话状态

完整代码实现

让我们先来看完整的代码实现：

#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
LangGraph Checkpoint 示例1: 基本会话持久化

这个示例展示了如何使用InMemorySaver来保存和恢复图执行状态，实现会话持久化。
应用功能：创建一个简单的计数器，每次调用增加计数，并通过checkpoint保存状态。
"""

from typing_extensions import TypedDict
from langgraph.graph import StateGraph, START, END
from langgraph.checkpoint.memory import InMemorySaver

# 定义状态类型
class CounterState(TypedDict):
    count: int

# 定义节点函数
def increment_counter(state: CounterState) -> CounterState:
    """增加计数器的值"""
    current_count = state.get("count", 0)
    return {"count": current_count + 1}

# 创建StateGraph
counter_graph = StateGraph(CounterState)

# 添加节点
counter_graph.add_node("increment", increment_counter)

# 设置入口点和出口点
counter_graph.add_edge(START, "increment")
counter_graph.add_edge("increment", END)

# 创建checkpoint saver
checkpointer = InMemorySaver()

# 编译图并添加checkpointer
compiled_graph = counter_graph.compile(checkpointer=checkpointer)

# 创建两个不同的会话配置
session1_config = {"configurable": {"thread_id": "session-1"}}
session2_config = {"configurable": {"thread_id": "session-2"}}

# 执行多次调用，观察状态如何在会话间保持
print("--- 执行会话1的调用 ---")

# 第一次调用传入初始值
result = compiled_graph.invoke({"count": 0}, config=session1_config)
print(f"会话1 - 第1次调用: count = {result['count']}")

# 后续调用不传入初始值，使用保存的状态
for i in range(2):
    result = compiled_graph.invoke({}, config=session1_config)
    print(f"会话1 - 第{i+2}次调用: count = {result['count']}")

print("--- 执行会话2的调用 ---")

# 第一次调用传入初始值
result = compiled_graph.invoke({"count": 0}, config=session2_config)
print(f"会话2 - 第1次调用: count = {result['count']}")

# 后续调用不传入初始值，使用保存的状态
result = compiled_graph.invoke({}, config=session2_config)
print(f"会话2 - 第2次调用: count = {result['count']}")

print("--- 再次执行会话1的调用 ---")
result = compiled_graph.invoke({}, config=session1_config)
print(f"会话1 - 第4次调用: count = {result['count']}")

执行结果

--- 执行会话1的调用 ---
会话1 - 第1次调用: count = 1
会话1 - 第2次调用: count = 2
会话1 - 第3次调用: count = 3
--- 执行会话2的调用 ---
会话2 - 第1次调用: count = 1
会话2 - 第2次调用: count = 2
--- 再次执行会话1的调用 ---
会话1 - 第4次调用: count = 4

代码解析

1. 状态类型定义

class CounterState(TypedDict):
    count: int

我们定义了一个简单的状态类型CounterState，只包含一个整数字段count，用于存储计数器的值。

2. 节点函数定义

def increment_counter(state: CounterState) -> CounterState:
    """增加计数器的值"""
    current_count = state.get("count", 0)
    return {"count": current_count + 1}

节点函数increment_counter接收当前状态，获取当前的计数值（如果不存在则默认为0），然后返回一个新的状态，其中计数器的值增加1。

3. 创建图结构

# 创建StateGraph
counter_graph = StateGraph(CounterState)

# 添加节点
counter_graph.add_node("increment", increment_counter)

# 设置入口点和出口点
counter_graph.add_edge(START, "increment")
counter_graph.add_edge("increment", END)

这部分代码创建了一个简单的图结构，包含一个increment节点，连接了START和END。

4. Checkpoint 配置

# 创建checkpoint saver
checkpointer = InMemorySaver()

# 编译图并添加checkpointer
compiled_graph = counter_graph.compile(checkpointer=checkpointer)

这是实现会话持久化的关键部分：

  • 我们创建了一个InMemorySaver实例，它会将状态保存在内存中
  • 在编译图时，我们通过checkpointer参数将这个saver传递给图

5. 会话配置

# 创建两个不同的会话配置
session1_config = {"configurable": {"thread_id": "session-1"}}
session2_config = {"configurable": {"thread_id": "session-2"}}

我们创建了两个不同的配置对象，用于标识不同的会话。每个配置对象都包含一个configurable字段，其中包含一个唯一的thread_id。

6. 执行图并观察结果

# 执行多次调用，观察状态如何在会话间保持
print("--- 执行会话1的调用 ---")
# 第一次调用传入初始值
result = compiled_graph.invoke({"count": 0}, config=session1_config)
print(f"会话1 - 第1次调用: count = {result['count']}")
# 后续调用不传入初始值，使用保存的状态
for i in range(2):
    result = compiled_graph.invoke({}, config=session1_config)
    print(f"会话1 - 第{i+2}次调用: count = {result['count']}")

我们通过invoke方法执行图，并通过config参数指定使用哪个会话配置：

  • 第一次调用时传入初始值{"count": 0}
  • 后续调用不传入值{}，这样LangGraph会使用之前保存的状态
  • 由于Checkpoint的存在，计数器的值会在多次调用之间保持并递增

执行结果分析

执行上述代码，你会看到类似以下的输出：

--- 执行会话1的调用 ---
会话1 - 第1次调用: count = 1
会话1 - 第2次调用: count = 2
会话1 - 第3次调用: count = 3

--- 执行会话2的调用 ---
会话2 - 第1次调用: count = 1
会话2 - 第2次调用: count = 2

--- 再次执行会话1的调用 ---
会话1 - 第4次调用: count = 4

这个输出展示了Checkpoint的核心特性：

  • 会话隔离：会话1和会话2的状态是完全独立的
  • 状态持久化：在多次调用之间，每个会话的状态都被保存下来
  • 状态恢复：当再次使用同一个会话配置调用图时，会从上次保存的状态继续执行

Checkpoint 的优势

通过上面的示例，我们可以看到Checkpoint的几个主要优势：

  1. 简化状态管理：不需要手动管理和传递状态，LangGraph会自动为你处理
  2. 会话隔离：不同会话之间的状态互不影响
  3. 灵活配置：通过configurable参数可以灵活地标识和管理会话
  4. 可扩展性：支持多种存储后端，可以根据需要选择合适的存储方式

管理 Checkpoint

除了基本的保存和恢复功能，我们还可以通过一些方法来管理Checkpoint。需要注意的是，不同版本的LangGraph API可能有所不同：

# 从特定会话恢复状态
session1_state = checkpointer.get(session1_config)
print(f"会话1的当前状态: {session1_state}")

# 验证状态是否正确恢复
print("--- 验证会话状态独立性 ---")
# 即使会话1已执行多次，会话2的状态仍然独立
session2_latest = compiled_graph.invoke({}, config=session2_config)
print(f"会话2最新状态: count = {session2_latest['count']}")

# 注意：在不同版本的LangGraph中，API可能有所不同

在实际应用中，你可以通过这些方法来查看、获取和验证保存的会话状态，为会话管理提供更大的灵活性。

实际应用场景

Checkpoint功能在很多实际应用场景中都非常有用：

  1. 多轮对话应用：保存聊天历史，实现会话上下文管理
  2. 长期运行的任务：在任务执行过程中保存中间状态，支持断点续传
  3. 用户会话管理：为不同用户维护独立的应用状态
  4. 复杂工作流：在复杂工作流中保存执行进度，支持流程恢复

总结

Checkpoint是LangGraph中实现会话持久化的关键功能，它通过保存和恢复图执行状态，使得我们可以构建更加复杂和强大的AI应用。在本文中，我们通过一个简单的计数器示例，学习了如何使用InMemorySaver来实现基本的会话持久化。

在接下来的文章中，我们将深入探讨更多Checkpoint的高级用法，包括异步环境中的Checkpoint、多轮对话持久化以及自定义序列化等。

如果你觉得这篇文章对你有帮助，请点赞、收藏并关注我，以获取更多关于LangGraph的教程和技巧！

avatar
文章
阅读
粉丝