LangGraph 构建的工作流调用数据库的时候怎么添加重试机制

在工作流许多使用场景中，我们可能希望我们的节点具有自定义的重试策略，例如在调用API、查询数据库或调用LLM等情况下。这里我给大家介绍一种方式叫 RetryPolicy。

from langgraph.pregel import RetryPolicy
RetryPolicy()

RetryPolicy(initial_interval=0.5, backoff_factor=2.0, max_interval=128.0, max_attempts=3, jitter=True, retry_on=<function default_retry_on at 0x78b964b89940>)

默认情况下，retry_on 参数使用的是 default_retry_on 函数，该函数会对所有异常进行重试，除了以下几种情况：

ValueError TypeError ArithmeticError ImportError LookupError NameError SyntaxError RuntimeError ReferenceError StopIteration StopAsyncIteration OSError

此外，对于来自流行的 HTTP 请求库（如 requests 和 httpx）的异常，它仅对 5xx 状态码 进行重试。
我们下面通过一个示例代码，展示了如何使用 LangGraph 构建一个基于状态的多节点工作流。代码中结合了数据库查询和模型调用，并通过重试策略增强了系统的健壮性。LangGraph 的灵活性和可扩展性使其非常适合构建复杂的 AI 应用和工作流。

导入依赖库

下面使我们具体的例子：

import operator
import sqlite3
from typing import Annotated, Sequence
from typing_extensions import TypedDict

from langchain_openai import ChatOpenAI
from langchain_core.messages import BaseMessage

from langgraph.graph import END, StateGraph, START
from langchain_community.utilities import SQLDatabase
from langchain_core.messages import AIMessage

初始化数据库和模型

db = SQLDatabase.from_uri("sqlite:///:memory:")
model = ChatOpenAI(model_name="gpt-4o-mini")

db: 使用 SQLDatabase.from_uri 方法初始化一个 SQLite 内存数据库。sqlite:///:memory: 表示使用内存数据库，数据不会持久化。model: 初始化 OpenAi的 gpt-4o-mini 模型，用于后续的模型调用。
由于这里用到了 sqlite，我们先向里面创建表和添加一些数据：

db = SQLDatabase.from_uri("sqlite:///:memory:")
# 获取底层引擎
engine = db._engine

# 创建艺术家表
db.run("""
CREATE TABLE artists (
    id INTEGER PRIMARY KEY,
    name TEXT NOT NULL,
    birth_year INTEGER,
    nationality TEXT,
    style TEXT
)
""")

# 创建文章表
db.run("""
CREATE TABLE articles (
    id INTEGER PRIMARY KEY,
    artist_id INTEGER,
    title TEXT NOT NULL,
    content TEXT NOT NULL,
    publication_date TEXT,
    FOREIGN KEY (artist_id) REFERENCES artists (id)
)
""")

# 插入艺术家数据
db.run("""
INSERT INTO artists (name, birth_year, nationality, style) VALUES 
('毕加索', 1881, '西班牙', '立体主义'),
('梵高', 1853, '荷兰', '后印象派'),
('莫奈', 1840, '法国', '印象派')
""")

# 插入文章数据 - 一次插入一篇
db.run("""
INSERT INTO articles (artist_id, title, content, publication_date) VALUES 
(1, '毕加索的艺术生涯', '毕加索(Pablo Picasso)是20世纪最具影响力的艺术家之一，他的作品跨越了多种风格。他与乔治·布拉克一起创立了立体主义，彻底改变了现代艺术的面貌。他的代表作包括《亚维农少女》、《格尔尼卡》等。毕加索不仅是一位画家，还是一位雕塑家、陶艺家和舞台设计师。', '2023-01-15')
""")

db.run("""
INSERT INTO articles (artist_id, title, content, publication_date) VALUES 
(1, '解析《格尔尼卡》', '《格尔尼卡》是毕加索最著名的作品之一，创作于1937年，是对西班牙内战中格尔尼卡轰炸事件的强烈控诉。这幅巨作通过扭曲的形象和强烈的黑白对比，表达了战争的恐怖和无辜者的苦难。作品中充满了象征性的图像，如哭泣的女人、死去的战士和受伤的马。', '2023-02-20')
""")

db.run("""
INSERT INTO articles (artist_id, title, content, publication_date) VALUES 
(2, '梵高的艺术与生活', '文森特·梵高(Vincent van Gogh)虽然生前默默无闻，但死后成为了艺术史上最重要的画家之一。他短暂而动荡的一生创作了约2100幅作品。梵高的画作以鲜艳的色彩和富有情感的笔触著称，《星空》、《向日葵》和《夜间咖啡馆》等都是他的代表作。', '2023-03-10')
""")

db.run("""
INSERT INTO articles (artist_id, title, content, publication_date) VALUES 
(3, '莫奈与印象派', '克劳德·莫奈(Claude Monet)是印象派运动的创始人之一，他的作品《日出·印象》给这一艺术流派命名。莫奈专注于捕捉光线和时间流逝对景物的影响，他的睡莲系列画作尤为著名。莫奈晚年在吉维尼的家园创造了著名的水上花园，成为了他创作的重要灵感来源。', '2023-04-05')
""")

# 验证数据 - 查询艺术家及其文章
result = db.run("""
SELECT a.name, a.nationality, a.style, b.title 
FROM artists a
JOIN articles b ON a.id = b.artist_id
ORDER BY a.name
""")

# 打印结果
print("艺术家及其文章列表:")
for row in result:
    print(row)

定义状态类型

class AgentState(TypedDict):
    messages: Annotated[Sequence[BaseMessage], operator.add]

gentState: 定义一个类型化的字典，表示工作流的状态。其中：messages: 是一个 BaseMessage 的序列，表示消息列表。Annotated: 使用 operator.add 注解，表示在状态更新时，messages 会通过追加的方式合并。

定义节点函数

def query_database(state):
    query_result = db.run("SELECT * FROM Artist LIMIT 10;")
    return {"messages": [AIMessage(content=query_result)]}
    
def call_model(state):
    response = model.invoke(state["messages"])
    return {"messages": [response]}

query_database: 该节点函数从数据库中查询数据，并将结果封装为 AIMessage 返回。db.run: 执行 SQL 查询，返回查询结果。AIMessage: 将查询结果包装为 AI 消息，以便后续节点处理。
call_model: 该节点函数调用 OpenAI GPT 模型，并将模型返回的结果添加到状态中。model.invoke: 调用 OpenAI 模型，传入当前状态中的消息列表。response: 模型的返回结果，会被包装为消息并返回。

构建工作流

builder = StateGraph(AgentState)
builder.add_node(
    "query_database",
    query_database,
    retry=RetryPolicy(retry_on=sqlite3.OperationalError),
)
builder.add_node("model", call_model, retry=RetryPolicy(max_attempts=5))
builder.add_edge(START, "model")
builder.add_edge("model", "query_database")
builder.add_edge("query_database", END)

graph = builder.compile()

StateGraph: 初始化一个状态图，工作流的状态类型为 AgentState。
add_node: 添加节点到工作流中：query_database 节点: 设置重试策略，当遇到 sqlite3.OperationalError 时重试。
model 节点: 设置最大重试次数为 5 次。
add_edge: 定义节点之间的连接关系：START -> model: 工作流从 model 节点开始。
model -> query_database: 模型调用完成后，进入数据库查询节点。
query_database -> END: 数据库查询完成后，工作流结束。
compile: 编译工作流，生成可执行的图对象 graph。
下面是展示出来的图：

最后我们来调用上面的graph

graph.invoke({"messages": [HumanMessage(content="有哪些艺术家")]})

得到结果

{'messages': [HumanMessage(content='有哪些艺术家', additional_kwargs={}, response_metadata={}),
  AIMessage(content='有很多著名的艺术家，涵盖各种艺术形式和历史时期。以下是一些著名艺术家的例子：\n\n### 绘画\n1. **列奥纳多·达·芬奇**（Leonardo da Vinci） - 意大利文艺复兴时期的画家，以《蒙娜丽莎》和《最后的晚餐》闻名。\n2. **米开朗基罗**（Michelangelo） - 同样是意大利文艺复兴时期的艺术家，以其雕塑和壁画著称，尤其是西斯廷小教堂的天顶画。\n3. **文森特·梵高**（Vincent van Gogh） - 荷兰后印象派画家，以其独特的色彩运用和笔触风格闻名。\n4. **巴勃罗·毕加索**（Pablo Picasso） - 西班牙画家和雕塑家，被认为是现代艺术的先锋之一，创造了立体主义。\n5. **克劳德·莫奈**（Claude Monet） - 法国印象派画家，以其光影和色彩变化的表现手法而闻名。\n\n### 雕塑\n1. **亨利·摩尔**（Henry Moore） - 英国雕塑家，以其大型抽象雕塑而闻名。\n2. **阿尔贝托·贾科梅蒂**（Alberto Giacometti） - 瑞士雕塑家和画家，以其纤细的人物雕塑著称。\n\n### 现代艺术\n1. **安迪·沃霍尔**（Andy Warhol） - 美国波普艺术家，以其对大众文化和消费主义的探索而著称。\n2. **杰夫·昆斯**（Jeff Koons） - 现代艺术家，以其大型不锈钢雕塑和玩具般的艺术作品而闻名。\n\n### 中国艺术家\n1. **齐白石** - 现代中国画家，以花鸟、虫鱼及山水画闻名。\n2. **徐悲鸿** - 以融合中西艺术的手法而受到推崇，尤其是其马的绘画作品。\n\n这只是一些著名艺术家的例子，还有很多其他艺术家在不同的领域和风格中留下了深远的影响。', additional_kwargs={'refusal': None}, response_metadata={'token_usage': {'completion_tokens': 502, 'prompt_tokens': 10, 'total_tokens': 512, 'completion_tokens_details': {'accepted_prediction_tokens': 0, 'audio_tokens': 0, 'reasoning_tokens': 0, 'rejected_prediction_tokens': 0}, 'prompt_tokens_details': {'audio_tokens': 0, 'cached_tokens': 0}}, 'model_name': 'gpt-4o-mini-2024-07-18', 'system_fingerprint': 'fp_b705f0c291', 'finish_reason': 'stop', 'logprobs': None}, id='run-14780504-3062-4105-8d50-9c73f7fbf83c-0', usage_metadata={'input_tokens': 10, 'output_tokens': 502, 'total_tokens': 512, 'input_token_details': {'audio': 0, 'cache_read': 0}, 'output_token_details': {'audio': 0, 'reasoning': 0}}),
  AIMessage(content="[(1, 1, '毕加索的艺术生涯', '毕加索(Pablo Picasso)是20世纪最具影响力的艺术家之一，他的作品跨越了多种风格。他与乔治·布拉克一起创立了立体主义，彻底改变了现代艺术的面貌。他的代表作包括《亚维农少女》、《格尔尼卡》等。毕加索不仅是一位画家，还是一位雕塑家、陶艺家和舞台设计师。', '2023-01-15'), (2, 1, '解析《格尔尼卡》', '《格尔尼卡》是毕加索最著名的作品之一，创作于1937年，是对西班牙内战中格尔尼卡轰炸事件的强烈控诉。这幅巨作通过扭曲的形象和强烈的黑白对比，表达了战争的恐怖和无辜者的苦难。作品中充满了象征性的图像，如哭泣的女人、死去的战士和受伤的马。', '2023-02-20'), (3, 2, '梵高的艺术与生活', '文森特·梵高(Vincent van Gogh)虽然生前默默无闻，但死后成为了艺术史上最重要的画家之一。他短暂而动荡的一生创作了约2100幅作品。梵高的画作以鲜艳的色彩和富有情感的笔触著称，《星空》、《向日葵》和《夜间咖啡馆》等都是他的代表作。', '2023-03-10'), (4, 3, '莫奈与印象派', '克劳德·莫奈(Claude Monet)是印象派运动的创始人之一，他的作品《日出·印象》给这一艺术流派命名。莫奈专注于捕捉光线和时间流逝对景物的影响，他的睡莲系列画作尤为著名。莫奈晚年在吉维尼的家园创造了著名的水上花园，成为了他创作的重要灵感来源。', '2023-04-05')]", additional_kwargs={}, response_metadata={})]}

可以看到调用 query_database 的时候，最后会追加到我们工作流返回的答案中。

执行逻辑

整体的执行逻辑就是，启动工作流: 从 START 开始，进入 model 节点。调用模型: model 节点调用 GPT 模型，生成响应。查询数据库: 将模型生成的响应传递给 query_database 节点，执行数据库查询。结束工作流: 数据库查询完成后，工作流结束。