深入理解 AI Agent 记忆系统设计

  • King
  • 发布于 15小时前
  • 阅读 52

如何让你的AI助手拥有"长期记忆"?本文将从架构设计到代码实现,带你深入理解一个生产级的Agent记忆系统。前言你有没有发现,和AI聊久了,它会"忘记"之前说过的内容?这是因为LLM本质上是无状态的——每次对话都是全新的开始。要让AI真正成为智能助手,我们需要为它构建

如何让你的 AI 助手拥有"长期记忆"?本文将从架构设计到代码实现,带你深入理解一个生产级的 Agent 记忆系统。

前言

你有没有发现,和 AI 聊久了,它会"忘记"之前说过的内容?这是因为 LLM 本质上是无状态的 —— 每次对话都是全新的开始。要让 AI 真正成为智能助手,我们需要为它构建一套记忆系统

本文将分享我设计的记忆系统架构,探讨如何为 AI Agent 赋予持久化记忆能力。


为什么 AI Agent 需要记忆系统?

LLM 的局限性

大语言模型虽然强大,但存在一个根本性的限制:上下文窗口有限。即使最新的模型支持百万级 tokens,也无法记住所有历史对话。更重要的是:

  • 成本问题:每次都传递完整历史,API 调用费用会快速累积
  • 注意力稀释:上下文过长会导致模型"注意力分散",回复质量下降
  • 无法持久化:新会话无法访问旧会话的信息

记忆系统的价值

一个设计良好的记忆系统可以:

  • 跨会话记忆:记住用户的偏好、历史交互
  • 智能检索:只召回与当前对话相关的内容
  • 知识沉淀:将交互中提取的知识长期保存
  • 成本优化:通过检索替代完整上下文传递

系统架构概览

agent-io 的记忆系统采用分层架构设计:

核心模块职责

模块 职责 文件
MemoryManager 统一入口,协调各组件 manager.rs
MemoryEntry 记忆条目数据结构 entry.rs
MemoryStore 存储抽象 trait store.rs
RingBuffer 短期记忆的环形缓冲区 buffer.rs
EmbeddingProvider 文本向量化接口 embeddings.rs
MemoryRanker 记忆相关性排序 ranker.rs

核心组件详解

记忆条目(MemoryEntry)

每条记忆都是一个结构化的条目:

pub struct MemoryEntry {
    pub id: String,                    // 唯一标识
    pub content: String,               // 记忆内容
    pub embedding: Option<Vec<f32>>,   // 向量表示(用于相似度搜索)
    pub memory_type: MemoryType,       // 记忆类型
    pub metadata: HashMap<String, Value>, // 元数据
    pub created_at: DateTime<Utc>,     // 创建时间
    pub last_accessed: Option<DateTime<Utc>>, // 最后访问时间
    pub importance: f32,               // 重要性评分 (0.0 - 1.0)
    pub access_count: u32,             // 访问次数
}

记忆类型分类

pub enum MemoryType {
    ShortTerm,   // 短期记忆:最近的对话
    LongTerm,    // 长期记忆:持久化知识
    Episodic,    // 情景记忆:特定事件/经历
    Semantic,    // 语义记忆:事实和概念
}

这种分类借鉴了认知心理学中的记忆模型:

  • 短期记忆:容量有限,快速存取,类似工作记忆
  • 长期记忆:容量无限,需要检索才能访问
  • 情景记忆:个人经历,如"昨天我去了咖啡店"
  • 语义记忆:通用知识,如"巴黎是法国首都"

相关性评分算法

系统会根据多个因素动态计算记忆的相关性:

pub fn relevance_score(&self) -> f32 {
    let age_hours = (Utc::now() - self.created_at).num_hours() as f32;
    let recency_factor = (-age_hours / 24.0 / 7.0).exp(); // 一周内指数衰减

    let access_factor = 1.0 + (self.access_count as f32).ln().max(0.0) * 0.1;

    self.importance * recency_factor * access_factor
}

评分公式score = 重要性 × 时间衰减因子 × 访问频率因子


存储抽象(MemoryStore)

采用 Trait 抽象存储层,支持多种后端实现:

#[async_trait]
pub trait MemoryStore: Send + Sync {
    async fn add(&self, entry: MemoryEntry) -> Result<String>;
    async fn search(&self, query: &str, limit: usize) -> Result<Vec<MemoryEntry>>;
    async fn search_by_embedding(&self, embedding: &[f32], limit: usize, threshold: f32) 
        -> Result<Vec<MemoryEntry>>;
    async fn get(&self, id: &str) -> Result<Option<MemoryEntry>>;
    async fn update(&self, entry: MemoryEntry) -> Result<()>;
    async fn delete(&self, id: &str) -> Result<()>;
    async fn clear(&self) -> Result<()>;
    async fn count(&self) -> Result<usize>;
}

内存存储实现

适合开发测试,所有数据保存在内存中:

pub struct InMemoryStore {
    memories: RwLock<HashMap<String, MemoryEntry>>,
}

支持向量相似度搜索(余弦相似度):

fn cosine_similarity(a: &[f32], b: &[f32]) -> f32 {
    let dot: f32 = a.iter().zip(b.iter()).map(|(x, y)| x * y).sum();
    let norm_a: f32 = a.iter().map(|x| x * x).sum::<f32>().sqrt();
    let norm_b: f32 = b.iter().map(|x| x * x).sum::<f32>().sqrt();
    dot / (norm_a * norm_b)
}

SQLite 存储实现

适合生产环境,支持持久化和全文搜索:

pub struct SqliteStore {
    conn: Arc<Mutex<Connection>>,
}

数据库设计亮点

  1. FTS5 全文搜索:利用 SQLite 的 FTS5 扩展实现高效文本搜索
  2. 向量存储:将向量序列化为 BLOB 存储
  3. 自动触发器:插入/更新/删除时自动同步 FTS 索引
CREATE VIRTUAL TABLE memories_fts USING fts5(
    id UNINDEXED,
    content,
    content='memories',
    content_rowid='rowid'
);

-- 自动同步触发器
CREATE TRIGGER memories_ai AFTER INSERT ON memories BEGIN
    INSERT INTO memories_fts(rowid, id, content) 
    VALUES (new.rowid, new.id, new.content);
END;

短期记忆缓冲区(RingBuffer)

短期记忆采用环形缓冲区实现,固定容量,自动淘汰最旧的数据:

pub struct RingBuffer<T> {
    buffer: VecDeque<T>,
    capacity: usize,
}

impl<T> RingBuffer<T> {
    pub fn push(&mut self, item: T) {
        if self.buffer.len() == self.capacity {
            self.buffer.pop_front(); // 淘汰最旧的
        }
        self.buffer.push_back(item);
    }
}

设计优势

  • O(1) 时间复杂度的插入和删除
  • 内存占用可控
  • 保留最近 N 条对话上下文

向量嵌入(Embeddings)

记忆系统依赖向量嵌入实现语义搜索:

#[async_trait]
pub trait EmbeddingProvider: Send + Sync {
    async fn embed(&self, text: &str) -> Result<Vec<f32>>;
    async fn embed_batch(&self, texts: &[&str]) -> Result<Vec<Vec<f32>>>;
    fn dimension(&self) -> usize;
}

内置实现

Provider 模型 维度
OpenAIEmbedding text-embedding-3-small 1536
OpenAIEmbedding text-embedding-3-large 3072
MockEmbedding 测试用 Mock 可配置

记忆排序器(Ranker)

当检索到多条相关记忆时,需要按相关性排序:

pub struct RankingWeights {
    pub similarity: f32,   // 向量相似度权重 (0.4)
    pub importance: f32,   // 重要性权重 (0.25)
    pub recency: f32,      // 时效性权重 (0.2)
    pub frequency: f32,    // 访问频率权重 (0.15)
}

综合评分公式

score = 0.4 × similarity + 0.25 × importance + 0.2 × recency + 0.15 × frequency

时间衰减机制

记忆的重要性会随时间衰减:

pub struct DecayConfig {
    pub daily_rate: f32,        // 每日衰减率 (默认 1%)
    pub min_threshold: f32,     // 最小阈值 (低于此值的记忆可清理)
    pub grace_period_days: u32, // 宽限期 (新记忆不衰减)
}

衰减公式importance × (1 - daily_rate)^days


核心工作流程

记忆存储流程(Remember)

pub async fn remember(&mut self, content: &str, memory_type: MemoryType) -> Result<String> {
    // 1. 生成向量嵌入
    let embedding = self.embedder.embed(content).await?;

    // 2. 创建记忆条目
    let entry = MemoryEntry::new(content)
        .with_type(memory_type)
        .with_embedding(embedding);

    // 3. 根据类型选择存储位置
    match memory_type {
        MemoryType::ShortTerm => {
            self.short_term.push(entry.clone());
            Ok(entry.id)
        }
        _ => {
            if self.config.enable_long_term {
                self.store.add(entry).await  // 持久化存储
            } else {
                self.short_term.push(entry.clone());
                Ok(entry.id)
            }
        }
    }
}

流程图

记忆检索流程(Recall)

pub async fn recall(&self, query: &str) -> Result<Vec<MemoryEntry>> {
    // 1. 查询向量化
    let query_embedding = self.embedder.embed(query).await?;

    // 2. 从长期记忆检索
    let mut memories = self.store
        .search_by_embedding(&query_embedding, limit, threshold)
        .await?;

    // 3. 从短期记忆检索
    for entry in self.short_term.iter_recent() {
        if let Some(ref embedding) = entry.embedding {
            let similarity = cosine_similarity(&query_embedding, embedding);
            if similarity >= threshold {
                memories.push(entry.clone());
            }
        }
    }

    // 4. 按相关性排序
    memories.sort_by(|a, b| b.relevance_score().partial_cmp(&a.relevance_score()).unwrap());

    // 5. 限制返回数量
    memories.truncate(self.config.retrieval_limit);

    Ok(memories)
}

流程图


实战示例

基本使用

use agent_io::memory::{
    MemoryManager, MemoryConfig, MemoryType,
    InMemoryStore, MockEmbedding
};

#[tokio::main]
async fn main() -> Result<()> {
    // 1. 创建存储和向量化服务
    let store = Arc::new(InMemoryStore::new());
    let embedder = Arc::new(MockEmbedding::new(384));

    // 2. 创建记忆管理器
    let config = MemoryConfig {
        short_term_size: 20,
        enable_long_term: true,
        retrieval_limit: 5,
        relevance_threshold: 0.7,
        ..Default::default()
    };
    let mut manager = MemoryManager::new(config, store, embedder);

    // 3. 存储记忆
    manager.remember("用户喜欢 Rust 编程语言", MemoryType::LongTerm).await?;
    manager.remember("用户是软件工程师", MemoryType::Semantic).await?;
    manager.remember("上次讨论了异步编程", MemoryType::Episodic).await?;

    // 4. 检索相关记忆
    let memories = manager.recall("编程").await?;
    for memory in memories {
        println!("相关记忆: {}", memory.content);
    }

    // 5. 构建上下文
    let context = manager.recall_context("用户的技术背景").await?;
    println!("上下文: {}", context);

    Ok(())
}

使用 SQLite 持久化

use agent_io::memory::SqliteStore;

// 创建文件数据库
let store = Arc::new(SqliteStore::open("./data/memories.db")?);

// 或使用内存数据库(测试用)
let store = Arc::new(SqliteStore::new()?);

自定义排序权重

use agent_io::memory::{MemoryRanker, RankingWeights};

let weights = RankingWeights {
    similarity: 0.5,   // 更重视语义相似
    importance: 0.3,
    recency: 0.1,
    frequency: 0.1,
};

let ranker = MemoryRanker::with_weights(weights)
    .with_recency_half_life(24.0 * 3.0); // 3天半衰期

设计亮点总结

架构设计

特性 实现方式
存储抽象 Trait + 多后端实现
短期/长期分离 RingBuffer + Store 双层架构
语义搜索 Embedding 向量化 + 余弦相似度
混合检索 FTS 全文 + 向量相似度
相关性排序 多因子加权评分

性能优化

  • 异步设计:所有 I/O 操作都是 async,支持高并发
  • 批量嵌入embed_batch 支持批量向量化,减少 API 调用
  • 连接池:SQLite 使用 Arc<Mutex<Connection>> 安全共享
  • 索引优化:SQLite 创建了 memory_type、importance、created_at 索引

可扩展性

// 自定义存储后端
impl MemoryStore for MyCustomStore { ... }

// 自定义嵌入服务
impl EmbeddingProvider for MyEmbeddingService { ... }

未来展望

当前的实现已经相当完善,但仍有一些可以改进的方向:

  1. 向量数据库集成:对于大规模记忆,可集成 Qdrant、Milvus 等专业向量数据库
  2. 记忆压缩:使用 LLM 总结压缩旧记忆,减少存储空间
  3. 遗忘机制:基于 DecayConfig 实现自动清理低重要性记忆
  4. 多模态记忆:支持图片、音频等多模态记忆存储
  5. 记忆推理:基于记忆进行推理,生成新的知识

结语

这套记忆系统设计体现了认知科学原理工程实践的结合。通过分层架构、存储抽象、多因子排序等设计,实现了既灵活又高效的 AI 记忆能力。

希望这些设计思路能为你的 AI Agent 开发提供参考。

完整代码请访问GitHub 仓库:https://github.com/lispking/agent-io


参考资料

点赞 0
收藏 0
分享
本文参与登链社区写作激励计划 ,好文好收益,欢迎正在阅读的你也加入。

0 条评论

请先 登录 后评论