AI Agent的记忆系统实现：从短期对话到长期知识

在上一篇文章中，我们搭建了 AI Agent 的基础框架。今天，我想深入讲讲 AI Agent 最核心的部分之一：记忆系统。说实话，我在实现记忆系统时走了不少弯路，希望通过这篇文章，能帮大家少走一些弯路。

从一个bug说起

还记得在开发知识助手的过程中，我遇到了一个很有意思的问题。一天我正在测试多轮对话功能：

我：Python的装饰器是什么？
助手：装饰器是Python中用于修改函数或类行为的一种设计模式...（省略具体解释）

我：能给个例子吗？
助手：抱歉，你想要什么的例子？

我：...（扶额）

这个对话看起来很搞笑，但实际上反映了一个严重的问题：AI Agent 没有"记忆"，它不记得上一轮谈论的是什么。这让我开始思考：如何给 AI Agent 实现一个真正好用的记忆系统？

记忆的层次

经过研究和实践，我把 AI Agent 的记忆分成了三个层次：

1. 短期记忆：对话上下文
2. 工作记忆：当前任务相关的信息
3. 长期记忆：持久化的知识库

就像人类大脑一样，这三种记忆各有特点和用途。

短期记忆：对话上下文

先从最基础的短期记忆开始。实现起来大概是这样的：

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()

这样实现有几个关键点：

1. 容量限制

   • 控制对话轮数
   • 限制总token数
   • 防止上下文爆炸

2. 时间衰减

   • 优先保留最近的对话
   • 重要信息可以标记保留
   • 定期清理过期内容

3. 上下文组织

   • 保持对话的连贯性
   • 突出重要信息
   • 压缩冗余内容

工作记忆：任务状态

工作记忆比短期记忆要复杂一些，因为它需要维护当前任务的状态：

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()

工作记忆的关键特点：

1. 任务栈

   • 支持任务嵌套
   • 维护任务层级
   • 管理任务切换

2. 上下文管理

   • 存储任务相关信息
   • 追踪状态变化
   • 支持上下文恢复

3. 资源管理

   • 及时释放无用内容
   • 避免内存泄漏
   • 优化资源使用

长期记忆：知识库

最后是最复杂的长期记忆系统：

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)
)

长期记忆的核心特点：

1. 分级存储

   • 重要信息持久化
   • 临时信息用TTL
   • 自动清理过期内容

2. 智能归档

   • 评估信息重要性
   • 选择存储策略
   • 管理存储周期

3. 高效检索

   • 向量相似度搜索
   • 关键词匹配
   • 实时更新索引

实践心得

在实现这个记忆系统的过程中，我总结了几点经验：

1. 分层设计很重要

   • 不同类型的记忆用不同的存储
   • 根据使用频率优化访问
   • 合理设置过期策略

2. 要平衡性能和复杂度

   • 简单任务用简单方案
   • 复杂功能按需加载
   • 注意资源消耗

3. 别忘了可维护性

   • 做好监控和日志
   • 支持手动干预
   • 保留调试接口

写在最后

一个好的记忆系统是 AI Agent 智能行为的基础。它不仅要能存储和检索信息，还要像人类大脑一样，能够区分信息的重要性，合理管理记忆的生命周期。

在下一篇文章中，我会讲解如何实现 AI Agent 的工具调用系统。如果你对记忆系统的实现有什么想法，欢迎在评论区交流。