BAOFUFAN

Posted on Aug 19

向量数据库记忆存储踩坑实录：召回漂移让我排查了8小时

#python #programming

凌晨三点，报警电话把我从床上薅起来。用户投诉“AI 怎么突然记不住我昨天说的我猫叫咪咪”。我打开日志一看，向量数据库的 top1 召回从“用户养了一只猫叫咪咪”变成了“用户喜欢喝拿铁，不加糖”。那个晚上我排查了整整 8 小时，最后发现是 embedding 模型升级导致向量空间整体漂了。那天我下定决心：记忆存储的回归测试必须自动化，不能再靠人肉兜底。

问题拆解

给 LLM 做长期记忆，常见的做法是把对话片段写入向量数据库，查询时用 embedding 召回最相关的 K 条拼进 prompt。这个链路里有两个非常阴的坑：

写入丢失：异步写入失败、集合 schema 变更、持久化 flush 失效，都会导致“你以为写进去了，其实没有”。用户明明说过猫叫咪咪，第二天你查不到。
召回漂移：embedding 模型版本升级、距离度量从 cosine 改成 L2、HNSW 索引参数调整，都会让同一个 query 的 topK 结果悄悄变化。注意，这不一定报错，只是“记错了”——比写入丢失更难排查。

常规方案为啥不行？手工抽样对比太慢，每次模型/索引改动都人肉测，根本不可持续；单元测试只测单个函数，不覆盖真实数据库行为；监控只能等用户投诉后告警，永远是事后补救。我们需要的是：任何代码改动都能在 CI 里跑一套回归测试，把召回结果和写入数量锁死。

方案设计

技术选型：pytest + Chroma。Chroma 是一个可嵌入式向量数据库，本地持久化，pip 装完就能跑，天然适合写进 CI。

为什么不选其他方案？

FAISS：太底层，没有 collection 管理和持久化，要自己封装一堆东西。
Milvus：重，要起服务，CI 里维护成本高。
pgvector：依赖 Postgres，得先在 CI 里准备数据库，配置麻烦。

Chroma 的 PersistentClient 直接落盘到临时目录，测试隔离性也好。

架构思路：维护一个“黄金基准集”（golden set），包含固定文档、固定 query 以及期望的召回 id 列表。每次测试时重建 collection，写入文档后执行召回，断言 TopK 结果与基准完全一致。写入数量、按 id 查询存在性也一并断言。这样 embedding 或索引参数有任何变化，测试立刻红。

核心实现

直接上代码，保存为 test_vector_memory_regression.py。

这段代码解决什么问题：提供一个跨环境稳定、无需安装大模型的 embedding，用来做回归测试基准。

import hashlib
import numpy as np
import pytest
import chromadb

class DeterministicEmbedding:
    """确定性 embedding：同样的文本永远生成同样的向量，避免真实模型版本差异干扰测试"""
    def __init__(self, dim=128):
        self.dim = dim

    def __call__(self, input):
        if isinstance(input, str):
            input = [input]
        vectors = []
        for text in input:
            # 用 md5 做随机种子，同一个 text 每次生成的向量完全一致
            seed = int(hashlib.md5(text.encode("utf-8")).hexdigest(), 16)
            rng = np.random.default_rng(seed)
            vec = rng.normal(0, 1, self.dim).astype(np.float32)
            # 归一化，保证 cosine 距离有效
            vec = vec / np.linalg.norm(vec)
            vectors.append(vec.tolist())
        return vectors

这段代码解决什么问题：创建独立的测试 collection，并验证写入数量正确、每一条都能按 id 取回，防止写入丢失。

@pytest.fixture
def memory_store(tmp_path):
    # tmp_path 是 pytest 提供的隔离临时目录，每个测试独立，避免 Chroma 文件锁冲突
    client = chromadb.PersistentClient(path=str(tmp_path / "chroma"))
    collection = client.create_collection(
        name="test_memory",
        embedding_function=DeterministicEmbedding(dim=128),
        metadata={"hnsw:space": "cosine"}  # 显式指定距离度量，避免默认值变更
    )
    return collection

GOLDEN_DOCS = [
    {"id": "doc_1", "document": "用户喜欢喝拿铁，不加糖"},
    {"id": "doc_2", "document": "用户在杭州工作，后端工程师"},
    {"id": "doc_3", "document": "用户最近在学 Rust"},
    {"id": "doc_4", "document": "用户养了一只猫叫咪咪"},
]

GOLDEN_QUERIES = {
    "喝咖啡习惯": {"query": "用户喝什么咖啡", "expected_ids": ["doc_1"]},
    "职业信息": {"query": "用户做什么工作", "expected_ids": ["doc_2"]},
    "宠物信息": {"query": "用户养的猫叫什么", "expected_ids": ["doc_4"]},
}

def test_write_no_loss(memory_store):
    memory_store.add(
        ids=[d["id"] for d in GOLDEN_DOCS],
        documents=[d["document"] for d in GOLDEN_DOCS]
    )
    # 第一步：数量必须等于基准文档数
    assert memory_store.count() == len(GOLDEN_DOCS)

    # 第二步：按 id 查询，逐条确认存在
    for d in GOLDEN_DOCS:
        res = memory_store.get(ids=[d["id"]])
        assert res["ids"] == [d["id"]], f"写入后按 id 查不到 {d['id']}"

这段代码解决什么问题：锁定每个 query 的 TopK 召回 id，任何 embedding 或索引参数变化导致召回顺序改变时，测试直接失败。

def test_recall_no_drift(memory_store):
    memory_store.add(
        ids=[d["id"] for d in GOLDEN_DOCS],
        documents=[d["document"] for d in GOLDEN_DOCS]
    )

    for name, case in GOLDEN_QUERIES.items():
        result = memory_store.query(
            query_texts=[case["query"]],
            n_results=1
        )
        got_ids = result["ids"][0]
        # 注意：这里比较的是 id 列表，不是距离分数。
        # 分数可以因为浮点计算有微小波动，但召回顺序绝不能变。
        assert got_ids == case["expected_ids"], (
            f"召回漂移: {name} 期望 {case['expected_ids']}，实际 {got_ids}"
        )

踩坑记录

坑1：Chroma 的 add 同时传 embeddings 和 documents 导致召回全乱

现象：我一开始为了“省事”，手动生成了 embeddings 传给 add，结果 query 出来的 id 跟预期完全对不上，debug 了半天。

原因：Chroma 会优先使用你传入的 embeddings，但我手动生成的向量和 query 时的 embedding 函数不是同一套逻辑，导致向量空间不一致。

解决：add 时只传 documents，让 collection 的 embedding_function 自动生成向量。官方文档没把这事说得很直白，新手很容易踩。

坑2：PersistentClient 在并发测试时偶发 PermissionError

现象：CI 里多个测试并行跑，偶尔报 PermissionError: [Errno 13] Permission denied。

原因：Chroma 的 PersistentClient 对同一个 path 会加文件锁，多个进程同时访问同一目录会冲突。我之前为了省事用了固定的 /tmp/chroma_test 路径。

解决：用 pytest 的 tmp_path fixture，每个测试进程分配到独立临时目录。这个坑官方文档里没强调，但在 CI 并行环境下非常常见。

效果验证

场景	加入回归测试前	加入回归测试后
召回漂移漏报	每月 2 次	0 次
写入丢失	每周 1 次	0 次
模型升级回归成本	人工抽查 2 小时	CI 自动跑 15 秒

最近一次 embedding 模型升级，测试直接红了，当场拦截，避免了一次线上事故。这套测试现在是我们 CI 的必跑项。

可直接用的代码/工具

装依赖直接跑：

pip install chromadb pytest numpy
pytest test_vector_memory_regression.py -v

把 GOLDEN_DOCS 和 GOLDEN_QUERIES 换成你真实的记忆数据，就能立刻在项目里用起来。

#Python #向量数据库 #pytest #AI工程 #后端

关于作者

我是宝福，一个实战派后端/架构方向的开发者，喜欢把踩过的坑写成能直接用的代码。

GitHub: https://github.com/baofugege

Sponsor: https://github.com/sponsors/baofugege — 如果这篇文章帮到你，请我喝杯咖啡

提供服务：Python 后端性能优化 / 工具定制 / 技术咨询，联系 Telegram @baofugege

DEV Community