LangGraph实战：用状态图构建智能算术助手，解锁AI的时间旅行能力

大家好，今天我要和大家分享一个有趣的LangGraph实战项目。我们将构建一个能够执行算术运算的智能助手，并深入探索LangGraph最强大的特性之一——时间旅行。

为什么选择LangGraph？

在构建复杂的AI应用时，我们常常需要管理多轮对话、控制流程、保存状态。传统的LangChain链式调用虽然简单，但难以处理复杂的交互逻辑。而LangGraph的出现，完美解决了这个问题：

状态管理：像Redux一样管理AI应用的状态
可控流程：清晰地定义节点和边，让AI的思考过程透明化
时间旅行：可以回放、分支和修改历史状态

项目概览：算术智能助手

我们要构建的助手具备以下能力：

支持加法、乘法、除法运算
通过工具调用执行具体计算
具备完整的对话记忆能力
支持时间旅行调试

1. 环境配置与工具定义

首先，让我们设置必要的环境变量和工具函数：

import os, getpass
from langchain_openai import ChatOpenAI

def _set_env(var: str):
    """安全地设置环境变量"""
    if not os.environ.get(var):
        os.environ[var] = getpass.getpass(f"{var}: ")

# 配置API密钥和基础URL
_set_env("OPENAI_API_KEY")
_set_env("OPENAI_BASE_URL")

# 定义三个算术工具函数
def multiply(a: int, b: int) -> int:
    """将两个整数相乘"""
    return a * b

def add(a: int, b: int) -> int:
    """将两个整数相加"""
    return a + b

def divide(a: int, b: int) -> float:
    """将整数a除以整数b"""
    return a / b

# 注册工具
tools = [add, multiply, divide]

关键点解析：

函数签名和文档字符串会被自动解析为工具规范
bind_tools() 让大模型知道可以调用这些函数

2. 构建状态图

LangGraph的核心是状态图。我们使用 MessagesState 作为状态容器，它会自动管理消息历史：

from langgraph.graph import MessagesState, StateGraph, START, END
from langgraph.prebuilt import ToolNode, tools_condition
from langchain_core.messages import SystemMessage

# 系统提示词
sys_msg = SystemMessage(content="You are a helpful assistant tasked with performing arithmetic on a set of inputs.")

# 定义助手节点
def assistant(state: MessagesState):
    return {"messages": [llm_with_tools.invoke([sys_msg] + state["messages"])]}

# 构建图
builder = StateGraph(MessagesState)

# 添加节点
builder.add_node("assistant", assistant)
builder.add_node("tools", ToolNode(tools))

# 定义边的流转
builder.add_edge(START, "assistant")
builder.add_conditional_edges("assistant", tools_condition)
builder.add_edge("tools", "assistant")

# 编译图
memory = MemorySaver()
graph = builder.compile(checkpointer=memory)

图的流转逻辑：

START → assistant：开始对话
assistant：模型生成回复（可能是文本或工具调用）
条件判断：如果有工具调用 → tools节点执行；否则 → END
tools执行完成后 → 回到assistant继续对话

3. 可视化图结构

为了让架构更直观，我们添加了多种可视化方案：

try:
    # 尝试生成PNG图片
    img_data = graph.get_graph(xray=True).draw_mermaid_png()
    with open("arithmetic_graph.png", "wb") as f:
        f.write(img_data)
    print("✅ 图渲染成功！")
except Exception as e:
    # 备选方案：输出Mermaid文本
    print("📝 图结构（Mermaid格式）：")
    print(graph.get_graph().draw_mermaid())

这样的设计确保了在任何环境下都能看到图的结构。

4. 运行第一个对话

让我们测试一下基本功能：

# 输入：计算2乘以3
initial_input = {"messages": HumanMessage(content="Multiply 2 and 3")}
thread = {"configurable": {"thread_id": "1"}}

# 流式执行
for event in graph.stream(initial_input, thread, stream_mode="values"):
    event['messages'][-1].pretty_print()

执行过程：

用户输入 → assistant节点
模型识别需要调用 multiply 工具
进入tools节点执行乘法
返回结果 → assistant生成最终回复

5. 时间旅行的魔力

现在到了最精彩的部分！LangGraph的 MemorySaver 让我们能够：

查看历史状态：

1
2
3

# 获取所有历史状态
all_states = [s for s in graph.get_state_history(thread)]
print(f"共有 {len(all_states)} 个历史状态")

每个状态包含：

values: 当前的消息列表
next: 下一个要执行的节点
config: 状态配置信息

回放特定状态：

1
2
3

# 选择倒数第二个状态进行回放
to_replay = all_states[-2]
print("要回放的状态：", to_replay.values)

分叉创建新分支：

# 在某个历史节点分叉，修改输入
to_fork = all_states[-2]
fork_config = graph.update_state(
    to_fork.config,
    {"messages": [HumanMessage(
        content='Multiply 5 and 3',  # 改为计算5×3
        id=to_fork.values["messages"][0].id  # 保持相同ID
    )]},
)

# 从分叉点继续执行
for event in graph.stream(None, fork_config, stream_mode="values"):
    event['messages'][-1].pretty_print()