基于 FAISS 的 AI 长期记忆系统示例

FAISS（Facebook AI Similarity Search）是 Meta 开源的高性能向量检索库，专为海量高维向量数据的快速近似搜索而设计。基于 FAISS 构建的 AI 长期记忆系统，本质上是为 AI 模型（尤其是大语言模型）提供持久化、可检索的"外部大脑"。

该系统通常采用"向量索引 + 元数据存储"的双存储架构：

• 向量化（Embedding）：利用嵌入模型（如 OpenAI text-embedding-ada-002、Sentence-BERT）将文本、图像等非结构化数据转换为高维向量（例如 1536 维）。
• 索引构建（Indexing）：使用 FAISS 将向量构建成高效的索引结构（如 FlatL2、IVFFlat、HNSW），支持毫秒级近似最近邻（ANN）搜索。
• 元数据关联：为每个向量附加原始文本、时间戳、来源等元数据（通常存储在 JSON 或 SQLite 中）。
• 检索增强生成（RAG）：当用户提问时，系统将查询文本向量化，通过 FAISS 召回最相关的历史记忆，并将其作为上下文注入 LLM 的 Prompt，实现"记忆唤醒"。

核心优势：

• 高性能：针对 GPU 和批量查询高度优化，能在百万级向量中实现亚秒级检索。
• 可扩展性：支持从本地小文件到分布式集群的平滑扩展。
• 灵活性：提供多种索引算法（精确搜索 vs. 近似搜索）和距离度量（L2、内积、余弦相似度），可根据数据规模和精度需求灵活配置。

典型应用场景：

• 对话记忆持久化：跨会话记住用户的偏好、习惯和重要对话历史。
• 知识库问答（KBQA）：将企业文档、技术手册向量化，构建私有知识库。
• 智能助理：实现基于历史交互的个性化推荐和上下文感知服务。

需求分析

本文示例代码实现了一个AI长期记忆系统，其核心需求源于当前AI应用对上下文记忆和历史交互保持能力的迫切需求。随着大型语言模型的普及，人们发现单次会话的AI虽然能提供即时响应，但缺乏持续的学习和记忆能力，导致每次对话都像是"重新开始"。这个记忆系统旨在解决这一痛点，为AI提供类似人类长期记忆的功能，使其能够跨会话保持知识、学习用户偏好、积累经验，从而提供更加个性化、连贯和智能的服务。

具体需求包括：向量化存储和检索技术文档、会议记录、用户交互历史等非结构化信息；支持高效的相似性搜索，快速找到相关历史信息；提供数据持久化机制，确保记忆不因系统重启而丢失；具备可扩展的架构，能够适应不同规模的数据集和多样的应用场景。

设计架构

该系统采用经典的知识库检索架构，以FAISS向量数据库为核心，结合元数据管理实现完整的记忆生命周期管理。系统设计分为三层：存储层使用FAISS索引进行高维向量的高效相似性搜索，这是系统的性能基础；业务层管理记忆的增删改查操作，包括向量的添加、相似性检索、元数据更新和软删除；持久化层负责将内存中的索引和元数据序列化到磁盘。系统采用模块化设计，索引类型可配置（FlatL2用于精确搜索，HNSW用于近似大规模搜索），维度可调节以适配不同的嵌入模型。元数据与向量通过索引位置隐式关联的设计简化了数据一致性维护，但需要注意删除操作只能软删除的局限性。系统还包含一个简单的文本向量化示例，展示了如何将文本转化为向量，尽管实际应用中需要替换为更专业的嵌入模型。

代码实现

# -*- coding: utf-8 -*-
"""
Created on Wed Jul  9 09:11:42 2025

@author: liguo
"""
import faiss
import numpy as np
from typing import List, Dict, Any
import json
import os

class AIMemory:
    def __init__(self, dim: int = 1536, index_type: str = "FlatL2", storage_path: str = "ai_memory"):
        """
        初始化AI长期记忆系统
        
        Args:
            dim: 向量维度
            index_type: 索引类型，支持FlatL2、HNSW等
            storage_path: 数据存储路径
        """
        self.dim = dim
        self.index_type = index_type
        self.storage_path = storage_path
        self.metadata = []  # 存储向量对应的元数据
        
        # 创建存储目录
        if not os.path.exists(storage_path):
            os.makedirs(storage_path)
            
        # 初始化索引
        self._init_index()
        
        # 加载已保存的数据
        self._load_memory()
    
    def _init_index(self):
        """根据指定类型初始化FAISS索引"""
        if self.index_type == "FlatL2":
            self.index = faiss.IndexFlatL2(self.dim)
        elif self.index_type == "HNSW":
            self.index = faiss.IndexHNSWFlat(self.dim, 32)
            self.index.hnsw.efConstruction = 40
        else:
            raise ValueError(f"不支持的索引类型: {self.index_type}")
    
    def _load_memory(self):
        """从磁盘加载已保存的记忆"""
        index_path = os.path.join(self.storage_path, "index.faiss")
        metadata_path = os.path.join(self.storage_path, "metadata.json")
        
        if os.path.exists(index_path) and os.path.exists(metadata_path):
            try:
                # 加载索引
                self.index = faiss.read_index(index_path)
                
                # 加载元数据
                with open(metadata_path, "r", encoding="utf-8") as f:
                    self.metadata = json.load(f)
                
                print(f"已从 {self.storage_path} 加载记忆: {len(self.metadata)} 条记录")
            except Exception as e:
                print(f"加载记忆失败: {e}")
                # 重新初始化索引
                self._init_index()
    
    def save_memory(self):
        """保存记忆到磁盘"""
        index_path = os.path.join(self.storage_path, "index.faiss")
        metadata_path = os.path.join(self.storage_path, "metadata.json")
        
        try:
            # 保存索引
            faiss.write_index(self.index, index_path)
            
            # 保存元数据
            with open(metadata_path, "w", encoding="utf-8") as f:
                json.dump(self.metadata, f, ensure_ascii=False, indent=2)
            
            print(f"已保存记忆到 {self.storage_path}")
        except Exception as e:
            print(f"保存记忆失败: {e}")
    
    def add_memory(self, vector: np.ndarray, data: Dict[str, Any]):
        """
        添加记忆条目
        
        Args:
            vector: 特征向量，numpy数组
            data: 相关元数据，字典格式
        """
        # 确保向量维度正确
        vector = vector.reshape(1, -1)
        if vector.shape[1] != self.dim:
            raise ValueError(f"向量维度不匹配，期望 {self.dim}，实际 {vector.shape[1]}")
        
        # 添加到索引
        self.index.add(vector)
        
        # 保存元数据
        self.metadata.append(data)
        
        # 自动保存（可优化为定期保存）
        self.save_memory()
        
        return len(self.metadata) - 1  # 返回添加的记忆ID
    
    def search_memory(self, query_vector: np.ndarray, k: int = 5) -> List[Dict[str, Any]]:
        """
        搜索相似记忆
        
        Args:
            query_vector: 查询向量
            k: 返回结果数量
            
        Returns:
            包含相似度和元数据的列表
        """
        # 确保向量维度正确
        query_vector = query_vector.reshape(1, -1)
        if query_vector.shape[1] != self.dim:
            raise ValueError(f"向量维度不匹配，期望 {self.dim}，实际 {query_vector.shape[1]}")
        
        # 搜索
        distances, indices = self.index.search(query_vector, k)
        
        # 构建结果
        results = []
        for i, idx in enumerate(indices[0]):
            if idx != -1:  # -1表示未找到
                results.append({
                    "similarity": float(distances[0][i]),
                    "metadata": self.metadata[idx],
                    "memory_id": idx
                })
        
        return results
    
    def update_memory(self, memory_id: int, new_data: Dict[str, Any]):
        """
        更新记忆元数据
        
        Args:
            memory_id: 记忆ID
            new_data: 新的元数据
        """
        if 0 <= memory_id < len(self.metadata):
            self.metadata[memory_id].update(new_data)
            self.save_memory()
            return True
        return False
    
    def delete_memory(self, memory_id: int):
        """
        删除记忆
        
        Args:
            memory_id: 记忆ID
        """
        if 0 <= memory_id < len(self.metadata):
            # 注意：FAISS不支持直接删除索引项，这里采用标记删除
            self.metadata[memory_id]["deleted"] = True
            self.save_memory()
            return True
        return False

# 简单的文本向量化器示例（实际应用中应使用更强大的模型）
def simple_text_embedding(text: str) -> np.ndarray:
    """简单的文本向量化函数，实际应用中应替换为如OpenAI Embeddings等"""
    # 这里仅作示例，返回随机向量
    # 实际应用中应使用如sentence-transformers等模型
    return np.random.random(1536).astype('float32')

# 使用示例
if __name__ == "__main__":
    # 初始化记忆系统
    memory = AIMemory(dim=1536, index_type="FlatL2")
    
    # 添加记忆
    memory.add_memory(
        simple_text_embedding("机器学习是人工智能的一个分支"),
        {"text": "机器学习是人工智能的一个分支", "timestamp": "2023-05-15", "source": "教科书"}
    )
    
    memory.add_memory(
        simple_text_embedding("深度学习是机器学习的一个子领域"),
        {"text": "深度学习是机器学习的一个子领域", "timestamp": "2023-05-16", "source": "网络文章"}
    )
    
    # 搜索记忆
    query_vector = simple_text_embedding("人工智能的分支有哪些")
    results = memory.search_memory(query_vector, k=2)
    
    print("\n搜索结果:")
    for result in results:
        print(f"相似度: {result['similarity']:.4f}")
        print(f"内容: {result['metadata']['text']}")
        print(f"来源: {result['metadata']['source']}")
        print("-" * 40)
    
    # 更新记忆
    memory.update_memory(0, {"importance": "high"})
    
    # 删除记忆
    # memory.delete_memory(1)    

代码执行结果

C:\Users\xiayu\miniconda3\envs\langchain03\python.exe "C:\Users\xiayu\PyCharmMiscProject\AI-Agent-Dev-Practices-Code\第3章代码\3.18-实现了一个基于 FAISS 的 AI 长期记忆系统.py"

已从 ai_memory 加载记忆: 12 条记录

已保存记忆到 ai_memory

已保存记忆到 ai_memory

搜索结果:

相似度: 259.3894

内容: 机器学习是人工智能的一个分支

来源: 教科书

---

相似度: 260.8898

内容: 机器学习是人工智能的一个分支

来源: 教科书

---

已保存记忆到 ai_memory

Process finished with exit code 0

代码解析

代码实现展示了清晰的面向对象设计，AIMemory类封装了所有记忆管理功能。构造函数初始化向量维度、索引类型和存储路径，并调用_init_index()根据配置创建FAISS索引。add_memory()方法首先验证输入向量维度，然后将其添加到索引并关联元数据。search_memory()方法接收查询向量，通过FAISS的search方法找到最相似的k个记忆，返回包含相似度分数、元数据和内存ID的结果列表。持久化通过save_memory()和_load_memory()方法实现，分别将FAISS索引和元数据JSON文件保存到磁盘。示例中使用的simple_text_embedding()函数只是随机向量生成器，实际部署时需要替换为真正的嵌入模型如Sentence-BERT或OpenAI Embeddings。主程序示例演示了完整的工作流程：初始化系统、添加记忆条目、执行搜索、更新和删除记忆，展示了系统的基本使用模式。