AI Agent的记忆系统实现：从短期对话到长期知识
在上一篇文章中，我们搭建了 AI Agent 的基础框架。今天，我想深入讲讲 AI Agent 最核心的部分之一：记忆系统。说实话，我在实现记忆系统时走了不少弯路，希望通过这篇文章，能帮大家少走一些弯路。
从一个bug说起
还记得在开发知识助手的过程中，我遇到了一个很有意思的问题。一天我正在测试多轮对话功能：
我：Python的装饰器是什么？
助手：装饰器是Python中用于修改函数或类行为的一种设计模式...（省略具体解释）
我：能给个例子吗？
助手：抱歉，你想要什么的例子？
我：...（扶额）这个对话看起来很搞笑，但实际上反映了一个严重的问题：AI Agent 没有"记忆"，它不记得上一轮谈论的是什么。这让我开始思考：如何给 AI Agent 实现一个真正好用的记忆系统？
记忆的层次
经过研究和实践，我把 AI Agent 的记忆分成了三个层次：
- 短期记忆：对话上下文
- 工作记忆：当前任务相关的信息
- 长期记忆：持久化的知识库
就像人类大脑一样，这三种记忆各有特点和用途。
短期记忆：对话上下文
先从最基础的短期记忆开始。实现起来大概是这样的：
class ConversationMemory:
def __init__(
self,
max_turns: int = 5,
max_tokens: int = 2000
):
self.messages = []
self.max_turns = max_turns
self.max_tokens = max_tokens
def add_message(
self,
role: str,
content: str
):
self.messages.append({
"role": role,
"content": content,
"timestamp": datetime.now()
})
# 保持对话长度在限制内
self._truncate_if_needed()
def get_context(self) -> List[dict]:
return [
{
"role": msg["role"],
"content": msg["content"]
}
for msg in self.messages
]
def _truncate_if_needed(self):
# 1. 按对话轮数截断
if len(self.messages) > self.max_turns * 2:
self.messages = self.messages[-(self.max_turns * 2):]
# 2. 按token数截断
total_tokens = sum(
len(msg["content"].split())
for msg in self.messages
)
while (
total_tokens > self.max_tokens and
len(self.messages) > 2
):
self.messages.pop(0)
total_tokens = sum(
len(msg["content"].split())
for msg in self.messages
)使用示例：
memory = ConversationMemory(max_turns=5)
# 记录对话
memory.add_message("user", "Python的装饰器是什么？")
memory.add_message("assistant", "装饰器是Python中用于修改函数或类行为的一种设计模式...")
memory.add_message("user", "能给个例子吗？")
# 获取上下文
context = memory.get_context()这样实现有几个关键点：
容量限制
- 控制对话轮数
- 限制总token数
- 防止上下文爆炸
时间衰减
- 优先保留最近的对话
- 重要信息可以标记保留
- 定期清理过期内容
上下文组织
- 保持对话的连贯性
- 突出重要信息
- 压缩冗余内容
工作记忆：任务状态
工作记忆比短期记忆要复杂一些，因为它需要维护当前任务的状态：
class WorkingMemory:
def __init__(self):
self.current_task = None
self.task_stack = []
self.context = {}
def start_task(
self,
task: Task,
context: dict = None
):
# 如果有正在进行的任务，压入栈
if self.current_task:
self.task_stack.append({
"task": self.current_task,
"context": self.context
})
# 开始新任务
self.current_task = task
self.context = context or {}
def end_task(self) -> Optional[Task]:
# 结束当前任务
if not self.task_stack:
self.current_task = None
self.context = {}
return None
# 恢复上一个任务
previous = self.task_stack.pop()
self.current_task = previous["task"]
self.context = previous["context"]
return self.current_task
def update_context(
self,
key: str,
value: Any
):
self.context[key] = {
"value": value,
"updated_at": datetime.now()
}
def get_context(
self,
key: str
) -> Optional[Any]:
if key in self.context:
return self.context[key]["value"]
return None使用示例：
working_memory = WorkingMemory()
# 开始一个复杂任务
working_memory.start_task(
task=Task(
name="重构代码",
steps=["分析代码", "设计方案", "实现重构"]
),
context={
"file": "main.py",
"changes_required": ["优化性能", "改善可读性"]
}
)
# 更新任务进度
working_memory.update_context(
"current_step",
"正在分析代码结构..."
)
# 临时切换到子任务
working_memory.start_task(
task=Task(
name="检查依赖",
steps=["列出依赖", "检查版本"]
)
)
# 完成子任务，返回主任务
working_memory.end_task()工作记忆的关键特点：
任务栈
- 支持任务嵌套
- 维护任务层级
- 管理任务切换
上下文管理
- 存储任务相关信息
- 追踪状态变化
- 支持上下文恢复
资源管理
- 及时释放无用内容
- 避免内存泄漏
- 优化资源使用
长期记忆：知识库
最后是最复杂的长期记忆系统：
class LongTermMemory:
def __init__(
self,
vector_store: VectorStore,
ttl_store: TTLStore
):
self.vector_store = vector_store
self.ttl_store = ttl_store
self.importance_threshold = 0.7
async def remember(
self,
content: str,
metadata: dict = None,
importance: float = None
):
# 1. 评估重要性
if importance is None:
importance = await self._evaluate_importance(
content
)
# 2. 决定存储策略
if importance > self.importance_threshold:
# 重要信息，存入向量数据库
await self.vector_store.add(
content=content,
metadata={
**(metadata or {}),
"importance": importance,
"created_at": datetime.now()
}
)
else:
# 临时信息，存入TTL存储
ttl = self._calculate_ttl(importance)
await self.ttl_store.set(
key=self._generate_key(content),
value={
"content": content,
"metadata": metadata,
"importance": importance
},
ttl=ttl
)
async def recall(
self,
query: str,
limit: int = 5
) -> List[Memory]:
# 1. 搜索向量数据库
vector_results = await self.vector_store.search(
query=query,
limit=limit
)
# 2. 搜索TTL存储
ttl_results = await self.ttl_store.search(
query=query,
limit=limit
)
# 3. 合并结果
all_results = [
*vector_results,
*ttl_results
]
# 4. 按相关性排序
return sorted(
all_results,
key=lambda x: x.relevance,
reverse=True
)[:limit]
async def forget(
self,
query: str = None,
before: datetime = None
):
# 删除旧的或不相关的记忆
if query:
await self.vector_store.delete(query)
if before:
await self.vector_store.delete_before(before)
async def _evaluate_importance(
self,
content: str
) -> float:
# 使用 LLM 评估内容的重要性
response = await self.llm.evaluate(
content=content,
criteria=[
"信息的独特性",
"内容的时效性",
"知识的通用性"
]
)
return float(response.importance)
def _calculate_ttl(
self,
importance: float
) -> int:
# 根据重要性计算过期时间
base_ttl = 60 * 60 * 24 # 1天
return int(base_ttl * importance)使用示例：
memory = LongTermMemory(
vector_store=MilvusStore(),
ttl_store=RedisStore()
)
# 存储新知识
await memory.remember(
content="Python 3.12 引入了新的类型语法...",
metadata={
"category": "Python",
"source": "PEP-xxx"
}
)
# 回忆相关内容
results = await memory.recall(
query="Python 3.12 新特性",
limit=5
)
# 清理旧记忆
await memory.forget(
before=datetime.now() - timedelta(days=30)
)长期记忆的核心特点：
分级存储
- 重要信息持久化
- 临时信息用TTL
- 自动清理过期内容
智能归档
- 评估信息重要性
- 选择存储策略
- 管理存储周期
高效检索
- 向量相似度搜索
- 关键词匹配
- 实时更新索引
实践心得
在实现这个记忆系统的过程中，我总结了几点经验：
分层设计很重要
- 不同类型的记忆用不同的存储
- 根据使用频率优化访问
- 合理设置过期策略
要平衡性能和复杂度
- 简单任务用简单方案
- 复杂功能按需加载
- 注意资源消耗
别忘了可维护性
- 做好监控和日志
- 支持手动干预
- 保留调试接口
写在最后
一个好的记忆系统是 AI Agent 智能行为的基础。它不仅要能存储和检索信息，还要像人类大脑一样，能够区分信息的重要性，合理管理记忆的生命周期。
在下一篇文章中，我会讲解如何实现 AI Agent 的工具调用系统。如果你对记忆系统的实现有什么想法，欢迎在评论区交流。