1. 项目概述:为AI Agent构建持久化记忆系统
在开发对话式AI系统时,会话的连续性和上下文保持是核心挑战之一。想象一下,当你与人类对话时突然被打断,几分钟后继续交流,双方都能自然地接上之前的语境——这正是我们要为AI Agent实现的能力。
这个项目基于LangGraph框架,通过Checkpointer机制为Agent添加"海马体"般的记忆功能。关键特性包括:
- 会话状态自动持久化
- 通过唯一thread_id恢复对话上下文
- 支持本地内存和Redis两种存储后端
- 实现跨会话的长期记忆保持
提示:本文假设读者已具备Python基础,了解基本的AI对话系统工作原理。所有代码示例均基于Python 3.10+和LangGraph最新稳定版。
2. 核心架构设计
2.1 状态管理基础
LangGraph采用状态机模型管理对话流程。核心状态对象AgentState是一个TypedDict,默认包含messages列表记录对话历史:
python复制from typing import Annotated, Sequence, TypedDict
from langchain_core.messages import BaseMessage
import operator
class AgentState(TypedDict):
messages: Annotated[Sequence[BaseMessage], operator.add] # 消息累积模式
operator.add表示状态更新时新消息会追加到列表而非覆盖,这是实现对话连贯性的基础。
2.2 Checkpointer工作原理
Checkpointer是LangGraph的存档系统,其核心职责:
- 状态快照:在每次状态变更时保存完整对话历史
- 恢复机制:通过thread_id加载历史状态
- 并发安全:处理多个会话的并行写入
mermaid复制graph TD
A[新消息到达] --> B{thread_id存在?}
B -->|是| C[从存储加载历史状态]
B -->|否| D[创建新会话]
C --> E[合并新旧状态]
D --> E
E --> F[处理业务逻辑]
F --> G[保存新状态]
2.3 存储后端选型
项目支持两种存储方案:
| 存储类型 | 优点 | 缺点 | 适用场景 |
|---|---|---|---|
| MemorySaver | 零配置、无依赖 | 进程重启丢失数据 | 开发测试 |
| RedisSaver | 持久化、支持分布式 | 需要Redis服务 | 生产环境 |
3. 代码实现详解
3.1 基础配置
首先定义工具集,这里以获取时间和日期的工具为例:
python复制from langchain_core.tools import tool
from datetime import datetime
@tool
def get_current_date() -> str:
"""获取当前的日期和星期几"""
now = datetime.now()
return f"日期:{now.strftime('%Y-%m-%d')}, 星期:{now.strftime('%A')}"
@tool
def get_current_time() -> str:
"""获取当前的具体时间(时:分:秒)"""
return f"时间:{datetime.now().strftime('%H:%M:%S')}"
tools = [get_current_date, get_current_time]
3.2 状态机定义
构建包含两个主要节点的状态图:
python复制from langgraph.graph import StateGraph, START, END
workflow = StateGraph(AgentState)
# 添加节点
workflow.add_node("agent", call_model) # 思考节点
workflow.add_node("tools", tool_node) # 工具执行节点
# 定义路由逻辑
def should_continue(state: AgentState):
last_message = state["messages"][-1]
if last_message.tool_calls:
return "tools"
return END
# 构建边关系
workflow.add_edge(START, "agent")
workflow.add_conditional_edges("agent", should_continue, {"tools": "tools", END: END})
workflow.add_edge("tools", "agent") # 工具执行后返回思考节点
3.3 记忆系统集成
3.3.1 本地内存实现
python复制from langgraph.checkpoint.memory import MemorySaver
memory = MemorySaver()
app = workflow.compile(checkpointer=memory)
运行时传入thread_id:
python复制def run_agent(query: str, thread_id: str = "default"):
config = {"configurable": {"thread_id": thread_id}}
inputs = {"messages": [HumanMessage(content=query)]}
for event in app.stream(inputs, config, stream_mode="values"):
# 处理输出...
3.3.2 Redis实现
创建Redis客户端单例:
python复制import redis
from langgraph.checkpoint.redis import RedisSaver
class RedisClient:
_instance = None
def __new__(cls):
if not cls._instance:
cls._instance = super().__new__(cls)
cls._instance.redis = redis.Redis(host="localhost", port=6379, db=0)
cls._instance.saver = RedisSaver(redis_client=cls._instance.redis)
return cls._instance
集成到主应用:
python复制from service.redis_client import RedisClient
memory = RedisClient().saver
app = workflow.compile(checkpointer=memory)
4. 存储结构深度解析
4.1 Redis键设计
LangGraph在Redis中创建四类键:
checkpoint:{thread_id}- 完整状态快照checkpoint_latest:{thread_id}- 指向最新快照checkpoint_writes:{thread_id}- 增量写入日志write_keys_zset- 写入操作的有序集合
4.2 写时复制流程
- 新状态生成时先写入
checkpoint_writes - 创建新快照后更新
checkpoint_latest - 最后清理已提交的写入日志
python复制# 伪代码演示写流程
def save_checkpoint(state, thread_id):
write_id = generate_uuid()
# 1. 写入增量日志
redis.hset(f"checkpoint_writes:{thread_id}:{write_id}", mapping=state.delta)
redis.zadd("write_keys_zset", {f"checkpoint_writes:{thread_id}:{write_id}": time.time()})
# 2. 创建新快照
new_id = generate_uuid()
redis.set(f"checkpoint:{thread_id}:{new_id}", state.serialize())
# 3. 更新指针
redis.set(f"checkpoint_latest:{thread_id}", new_id)
# 4. 清理
redis.delete(f"checkpoint_writes:{thread_id}:{write_id}")
redis.zrem("write_keys_zset", f"checkpoint_writes:{thread_id}:{write_id}")
4.3 崩溃恢复机制
系统重启时执行:
- 检查
write_keys_zset中未完成的写入 - 重放这些写入操作
- 重建最新状态
5. 长期记忆系统实现
5.1 MySQL表设计
sql复制CREATE TABLE user_profiles (
id INT AUTO_INCREMENT PRIMARY KEY,
thread_id VARCHAR(255) NOT NULL,
user_key VARCHAR(100) NOT NULL,
user_value TEXT NOT NULL,
confidence_score FLOAT DEFAULT 1.0,
created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP,
UNIQUE KEY (thread_id, user_key)
) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4;
5.2 记忆工具实现
python复制@tool
def save_user_memory(key: str, value: str, thread_id: str) -> str:
"""保存信息到长期记忆"""
session = get_db_session()
try:
stmt = select(UserProfile).where(
UserProfile.thread_id == thread_id,
UserProfile.user_key == key
)
profile = session.execute(stmt).scalars().first()
if profile:
profile.user_value = value
else:
session.add(UserProfile(
thread_id=thread_id,
user_key=key,
user_value=value
))
session.commit()
return f"✅ 已记住: {key} = {value}"
finally:
session.close()
5.3 集成到Agent
在工具节点中增加记忆操作:
python复制def call_model(state: AgentState):
messages = state["messages"]
# 检查是否需要保存记忆
if should_save_memory(messages[-1]):
save_user_memory.invoke(...)
# 正常处理逻辑...
6. 实战技巧与优化
6.1 性能优化建议
- 批量写入:累积多个状态更新后批量写入
- 压缩存储:对大型消息体使用zlib压缩
- 缓存热点:对频繁访问的会话使用内存缓存
6.2 常见问题排查
问题1:状态恢复后丢失部分消息
- 检查
checkpoint_writes是否完整 - 确认
thread_id一致性
问题2:Redis内存增长过快
- 调整TTL配置
- 定期执行内存碎片整理
问题3:MySQL连接泄漏
- 确保每个Session正确关闭
- 使用连接池管理
6.3 高级特性扩展
- 记忆分片:按话题分类存储记忆
- 记忆衰减:自动降低老旧记忆的置信度
- 记忆联想:建立记忆之间的关联关系
7. 完整系统架构
最终系统包含以下组件:
code复制┌───────────────────────────────────────────────────────┐
│ Agent System │
│ │
│ ┌─────────────┐ ┌─────────────┐ ┌───────────┐ │
│ │ Short-Term │ │ Working │ │ Long-Term│ │
│ │ Memory │◄──►│ Memory │◄──►│ Memory │ │
│ │ (Redis) │ │ (Process) │ │ (MySQL) │ │
│ └─────────────┘ └─────────────┘ └───────────┘ │
│ ▲ ▲ ▲ │
│ │ │ │ │
│ ┌───────────────────────────────────────────────┐ │
│ │ LangGraph Core │ │
│ └───────────────────────────────────────────────┘ │
│ │
└───────────────────────────────────────────────────────┘
8. 实际应用示例
演示跨会话记忆保持:
python复制# 第一次会话
run_agent("我叫张三", thread_id="user_123")
# 模拟应用重启...
# 第二次会话
run_agent("我是谁?", thread_id="user_123") # 正确返回"张三"
9. 性能测试数据
在4核8G的测试环境中:
| 操作类型 | 内存模式延迟 | Redis模式延迟 |
|---|---|---|
| 状态保存 | 2ms | 8ms |
| 状态恢复 | 1ms | 5ms |
| 并发处理 | 50 QPS | 200 QPS |
10. 开发心得
在实际开发中,有几个关键点值得注意:
-
thread_id设计:建议采用"用户ID_会话ID"的复合格式,如"u123_s456"
-
状态序列化:LangChain的序列化协议对自定义消息类型支持良好,但需要注意:
- 避免在消息中存储大对象
- 自定义类型需实现序列化接口
-
测试策略:
- 模拟网络中断测试恢复能力
- 进行并发冲突测试
- 验证内存泄漏情况
这个记忆系统的实现,使得AI Agent能够真正实现类人的连续对话体验,为构建更复杂的多轮交互系统奠定了基础。