LangGraph从新手到老师傅 - 9 - 异步SQLite Checkpoint持久化

在钱塘江
422 阅读7分钟

前言

在上一篇文章中，我们学习了如何使用InMemorySaver来实现基本的会话持久化。然而，在实际应用中，我们经常需要处理两个重要需求：

  1. 一是在异步环境中高效运行
  2. 二是将状态持久化到磁盘以便程序重启后能够恢复。

今天，我们将学习如何使用LangGraph的AsyncSqliteSaver来同时满足这两个需求。

异步编程与持久化存储的重要性

在现代Python应用开发中，异步编程已经成为处理高并发场景的标准方法。而对于需要长期运行或需要保存重要状态的应用来说，将状态持久化到数据库中也是必不可少的。

LangGraph提供了AsyncSqliteSaver，这是一个专为异步环境设计的Checkpoint存储实现，它将状态保存到SQLite数据库中，既支持异步操作，又提供了持久化存储能力。

示例功能概述

我们将创建一个异步的计数器应用，展示如何在异步环境中使用SQLite进行状态持久化：

  • 使用异步节点函数和异步调用
  • 通过AsyncSqliteSaver将状态保存到SQLite数据库
  • 在多次调用之间保持状态
  • 展示如何管理异步环境中的Checkpoint

完整代码实现

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

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

"""
LangGraph Checkpoint 示例2: 异步SQLite Checkpoint

这个示例展示了如何使用AsyncSqliteSaver在异步环境中保存和恢复图执行状态。
应用功能：创建一个异步的计数器，每次调用增加计数，并通过SQLite数据库保存状态。

知识点：
- 创建并使用AsyncSqliteSaver
- 在异步环境中使用checkpoint功能
- 通过SQLite数据库实现持久化存储
- 理解异步checkpoint的工作原理

注意：需要安装aiosqlite包: pip install aiosqlite
"""

import asyncio
import aiosqlite
from typing_extensions import TypedDict
from langgraph.graph import StateGraph, START, END
from langgraph.checkpoint.sqlite.aio import AsyncSqliteSaver

print("======= Checkpoint示例2: 异步SQLite Checkpoint =======")

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

# 定义异步节点函数
async def async_increment_counter(state: AsyncCounterState) -> AsyncCounterState:
    """异步增加计数器的值"""
    current_count = state.get("count", 0)
    # 模拟异步操作
    await asyncio.sleep(0.1)
    return {"count": current_count + 1}

# 创建异步执行的主函数
async def main():
    # 创建StateGraph
    counter_graph = StateGraph(AsyncCounterState)
    
    # 添加异步节点
    counter_graph.add_node("increment", async_increment_counter)
    
    # 设置入口点和出口点
    counter_graph.add_edge(START, "increment")
    counter_graph.add_edge("increment", END)
    
    # 创建SQLite数据库连接并初始化AsyncSqliteSaver
    print("初始化SQLite数据库...")
    async with AsyncSqliteSaver.from_conn_string("checkpoints.db") as checkpointer:
        # 编译图并添加checkpointer
        compiled_graph = counter_graph.compile(checkpointer=checkpointer)
        
        # 创建两个不同的会话配置
        session1_config = {"configurable": {"thread_id": "async-session-1"}}
        session2_config = {"configurable": {"thread_id": "async-session-2"}}
        
        # 执行多次调用，观察状态如何在会话间保持
        print("--- 执行会话1的调用 ---")
        for i in range(3):
            result = await compiled_graph.ainvoke({},config=session1_config)
            print(f"会话1 - 第{i+1}次调用: count = {result['count']}")
        
        print("--- 执行会话2的调用 ---")
        for i in range(2):
            result = await compiled_graph.ainvoke({}, config=session2_config)
            print(f"会话2 - 第{i+1}次调用: count = {result['count']}")
        
        print("--- 再次执行会话1的调用 ---")
        result = await compiled_graph.ainvoke({}, config=session1_config)
        print(f"会话1 - 第4次调用: count = {result['count']}")

# 运行异步主函数
if __name__ == "__main__":
    asyncio.run(main())

代码解析

1. 异步环境准备

import asyncio
import aiosqlite
from langgraph.checkpoint.sqlite.aio import AsyncSqliteSaver

我们导入了异步编程所需的库：asyncio用于异步操作，aiosqlite是SQLite的异步接口，AsyncSqliteSaver是LangGraph提供的异步SQLite Checkpoint实现。

2. 异步节点函数

async def async_increment_counter(state: AsyncCounterState) -> AsyncCounterState:
    """异步增加计数器的值"""
    current_count = state.get("count", 0)
    # 模拟异步操作
    await asyncio.sleep(0.1)
    return {"count": current_count + 1}

与之前的同步节点函数不同，这里我们使用async def定义了一个异步节点函数，并在其中使用await来模拟异步操作。在实际应用中，这可能是数据库查询、API调用等异步操作。

3. 异步主函数

async def main():
    # 创建StateGraph和添加节点...

我们创建了一个异步主函数来包含所有的异步操作，这是Python异步编程的标准模式。

4. AsyncSqliteSaver配置

# 创建SQLite数据库连接并初始化AsyncSqliteSaver
async with AsyncSqliteSaver.from_conn_string("checkpoints.db") as checkpointer:
    # 使用checkpointer...

这是使用AsyncSqliteSaver的关键部分：

  • 我们使用AsyncSqliteSaver.from_conn_string()方法创建一个AsyncSqliteSaver实例，指定数据库文件为"checkpoints.db"
  • 使用async with语句来确保数据库连接在使用完毕后正确关闭
  • checkpoints.db文件会自动创建，如果文件已存在，则会使用现有的数据库

5. 异步调用图

# 异步调用图
result = await compiled_graph.ainvoke({}, config=session1_config)

在异步环境中，我们使用ainvoke方法而不是invoke方法来调用图，并使用await等待调用完成。

执行结果分析

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

======= Checkpoint示例2: 异步SQLite Checkpoint =======
初始化SQLite数据库...

--- 执行会话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

更重要的是，你会发现在程序目录下创建了一个名为checkpoints.db的文件，这就是保存Checkpoint数据的SQLite数据库文件。即使你关闭程序并重新运行，之前保存的状态也会被恢复。

执行第二次的话，结果会不同，状态数据被存储到数据库中了

======= Checkpoint示例2: 异步SQLite Checkpoint =======
初始化SQLite数据库...
--- 执行会话1的调用 ---
会话1 - 第1次调用: count = 5
会话1 - 第2次调用: count = 6
会话1 - 第3次调用: count = 7
--- 执行会话2的调用 ---
会话2 - 第1次调用: count = 3
会话2 - 第2次调用: count = 4
--- 再次执行会话1的调用 ---
会话1 - 第4次调用: count = 8

AsyncSqliteSaver的工作原理

AsyncSqliteSaver的工作原理可以简要概括为：

  1. 当你调用compiled_graph.ainvoke()时，LangGraph会执行图计算
  2. 计算完成后，LangGraph会调用checkpointer.put()方法，将当前状态保存到SQLite数据库
  3. 当你再次使用相同的会话配置调用图时，LangGraph会调用checkpointer.get()方法，从数据库中读取之前保存的状态
  4. 图计算会基于读取的状态继续执行

实际应用场景

异步SQLite Checkpoint在很多实际应用场景中都非常有用：

  1. 异步Web应用：在FastAPI或Starlette等异步Web框架中使用，处理并发请求
  2. 长时间运行的异步任务：在异步任务中保存进度，支持断点续传
  3. 数据处理管道：在异步数据处理管道中保存中间状态
  4. 微服务架构：在异步微服务中管理会话状态

生产环境注意事项

虽然AsyncSqliteSaver适合开发和测试环境，但在生产环境中，你可能需要考虑以下几点：

  1. 并发性能：SQLite在高并发场景下的性能可能不如其他数据库（如PostgreSQL）
  2. 数据备份：定期备份SQLite数据库文件，以防止数据丢失
  3. 数据库大小：注意监控数据库文件大小，避免无限增长
  4. 迁移策略：考虑未来可能需要将数据迁移到其他数据库的策略

对于生产环境，LangGraph还提供了其他Checkpoint实现，如PostgresSaver，它基于PostgreSQL数据库，提供了更好的并发性能和可靠性。

总结

在本文中，我们学习了如何使用LangGraph的AsyncSqliteSaver在异步环境中实现状态持久化。通过将异步编程与SQLite持久化存储相结合，我们可以构建既高效又可靠的AI应用。

AsyncSqliteSaver为我们提供了一个简单而强大的解决方案，它既支持异步操作，又提供了持久化存储能力，非常适合开发和测试环境。在实际应用中，我们需要根据具体需求选择合适的Checkpoint实现，并注意相关的优化和最佳实践。

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

avatar
文章
阅读
粉丝