RAG 系列 6:使用 Chroma 实现持久化向量检索


上一篇使用 FAISS 完成了本地向量检索,但 FAISS 索引只负责向量和位置编号,原始文本与元数据需要单独保存。

这一篇改用 Chroma 的原生 API。Chroma 可以在一个 Collection 中同时管理 ID、Embedding、Document 和 Metadata,还能直接完成元数据过滤、更新、插入和删除。

本文使用 PersistentClient 把数据保存在项目的 data/p06_chroma_native/,不启动独立服务,也不使用 LangChain 的 Chroma 包装类。

1. Chroma 的整体读写流程

Chroma Collection 写入与查询流程

Chroma 把 FAISS 示例中分开管理的内容组织成 Collection 记录。写入时,应用把 ID、文档向量、原文和元数据一起交给 collection.add();查询时,应用传入查询向量、Top-K 数量和可选的元数据过滤条件。

这里仍然要区分两条 Embedding 路径:

  • 文档在建立 Collection 时生成文档向量,然后写入 Chroma。
  • 查询文本使用 Qwen3 的 query prompt 生成查询向量,只参与 collection.query(),不会成为新记录。

查询结果可以直接返回 ID、Document、Metadata 和 Distance,不再需要像 FAISS 示例那样使用位置编号访问另一份 JSON 列表。

2. Chroma 的本地持久化结构

Chroma 本地持久化目录结构

本篇使用的逻辑结构是:

PersistentClient
└── Collection: life_notes
    ├── ID
    ├── Embedding
    ├── Document
    └── Metadata

PersistentClient 接收一个目录路径。创建 Collection、添加数据后,Chroma 会自动把数据写入这个目录,不需要额外调用 persist()。重新创建指向同一路径的客户端,就可以读取已经保存的 Collection。

实测生成的主要文件包括:

chroma.sqlite3
<collection-id>/data_level0.bin
<collection-id>/header.bin
<collection-id>/length.bin
<collection-id>/link_lists.bin

chroma.sqlite3 保存系统信息、Collection、文档和元数据等内容,UUID 目录保存 HNSW 向量索引及其相关数据。具体内部表和文件可能随 Chroma 版本变化,业务代码不应直接依赖这些内部文件名。

不要只复制其中一个文件。迁移或备份时应该在停止写入后保存整个持久化目录,避免 SQLite 与 HNSW 文件来自不同状态。

3. 创建运行环境

Chroma 和后面的 PyMilvus 共用 .venv_vector_db,与 FAISS 的 Conda 环境分开。

依赖文件为:

chromadb==1.5.0
numpy==2.2.6
pymilvus==2.6.9
sentence-transformers==5.1.2
torch==2.9.1
transformers==4.57.6

创建 Python 3.12 虚拟环境:

cd source/_posts/llm_learning

python3.12 -m venv .venv_vector_db
source .venv_vector_db/bin/activate

python -m pip install -U pip
python -m pip install \
  -r rag/p06_chroma_native/requirements.txt
python -m pip check

本机检查结果为:

No broken requirements found.

4. 创建持久化 Collection

代码文件为 01_create_persistent_collection.py:

# 这个文件使用 Chroma 原生 API 创建持久化 Collection,并写入文本、元数据和向量。

from pathlib import Path
import shutil

import chromadb
import torch
from sentence_transformers import SentenceTransformer


PROJECT_ROOT = Path(__file__).resolve().parents[2]
MODEL_PATH = PROJECT_ROOT / "embedding_model"
DATA_DIR = PROJECT_ROOT / "data" / "p06_chroma_native"
COLLECTION_NAME = "life_notes"

documents = [
    {"id": "doc-1", "text": "西湖边有步道和树荫,很适合周末散步。", "category": "travel"},
    {"id": "doc-2", "text": "苹果富含膳食纤维,是常见的健康水果。", "category": "food"},
    {"id": "doc-3", "text": "汽车需要定期更换机油并检查轮胎。", "category": "car"},
    {"id": "doc-4", "text": "杭州植物园环境安静,适合慢慢游览。", "category": "travel"},
]

# 本脚本负责重建教学数据库,先清理自己的数据目录,避免旧 HNSW 文件残留。
if DATA_DIR.exists():
    shutil.rmtree(DATA_DIR)
DATA_DIR.mkdir(parents=True, exist_ok=True)
client = chromadb.PersistentClient(path=str(DATA_DIR))

collection = client.create_collection(
    name=COLLECTION_NAME,
    configuration={"hnsw": {"space": "cosine"}},
)

device = "mps" if torch.backends.mps.is_available() else "cpu"
model = SentenceTransformer(str(MODEL_PATH), device=device)
vectors = model.encode(
    [item["text"] for item in documents],
    normalize_embeddings=True,
).tolist()

# Chroma 同时保存 ID、向量、原始文本和标量元数据。
collection.add(
    ids=[item["id"] for item in documents],
    embeddings=vectors,
    documents=[item["text"] for item in documents],
    metadatas=[{"category": item["category"]} for item in documents],
)

print(f"数据库目录:{DATA_DIR.relative_to(PROJECT_ROOT)}")
print(f"Collection:{collection.name}")
print(f"记录数:{collection.count()}")
print("数据库文件:")
for path in sorted(DATA_DIR.rglob("*")):
    if path.is_file():
        print(f"- {path.relative_to(DATA_DIR)}")

4.1 为什么直接传 embeddings

Chroma 可以配置内部 Embedding Function,但本系列已经部署了 Qwen3-Embedding-0.6B。为了保证文档和查询都使用同一个本地模型,代码先调用 SentenceTransformer,再把向量直接传给 Chroma:

collection.add(
    embeddings=vectors,
    ...
)

这样 Chroma 只负责存储和检索,不会在后台下载其他默认模型。

4.2 设置 cosine HNSW

创建 Collection 时明确设置:

configuration={"hnsw": {"space": "cosine"}}

本篇查询返回的是 cosine distance,计算方式为 1 - cosine_similarity,数值越小越接近。它与 FAISS 中“余弦相似度越大越接近”的方向相反。

IndexFlatIP 会精确比较全部向量,而 Chroma 单机 Collection 使用 HNSW 完成近似最近邻检索。HNSW 用召回率换取更低的查询开销,因此不能把小样例中排序一致理解为所有数据规模下都必然返回完全相同的结果。

运行代码:

python rag/p06_chroma_native/01_create_persistent_collection.py

UUID 目录由 Chroma 每次创建 Collection 时生成。把动态目录名统一写成 后,真实输出为:

数据库目录:data/p06_chroma_native
Collection:life_notes
记录数:4
数据库文件:
- <collection-id>/data_level0.bin
- <collection-id>/header.bin
- <collection-id>/length.bin
- <collection-id>/link_lists.bin
- chroma.sqlite3

UUID 目录名由 Chroma 生成,不应该在代码中写死。

5. Top-K 查询和元数据过滤

代码文件为 02_query_and_filter.py:

# 这个文件查询已经持久化的 Chroma Collection,并演示元数据过滤。

from pathlib import Path

import chromadb
import torch
from sentence_transformers import SentenceTransformer


PROJECT_ROOT = Path(__file__).resolve().parents[2]
MODEL_PATH = PROJECT_ROOT / "embedding_model"
DATA_DIR = PROJECT_ROOT / "data" / "p06_chroma_native"

client = chromadb.PersistentClient(path=str(DATA_DIR))
collection = client.get_collection("life_notes")

device = "mps" if torch.backends.mps.is_available() else "cpu"
model = SentenceTransformer(str(MODEL_PATH), device=device)
query = "周末想找一个适合散步的地方"
query_vector = model.encode(
    [query],
    prompt_name="query",
    normalize_embeddings=True,
).tolist()

# cosine distance 越小,表示两段文本的方向越接近。
results = collection.query(
    query_embeddings=query_vector,
    n_results=3,
    include=["documents", "metadatas", "distances"],
)

print(f"查询:{query}")
print("全部分类的 Top 3:")
for item_id, document, metadata, distance in zip(
    results["ids"][0],
    results["documents"][0],
    results["metadatas"][0],
    results["distances"][0],
):
    print(
        f"distance={distance:.4f} id={item_id} "
        f"category={metadata['category']} text={document}"
    )

filtered_results = collection.query(
    query_embeddings=query_vector,
    n_results=2,
    where={"category": "travel"},
    include=["documents", "metadatas", "distances"],
)

print("\n只查询 category=travel:")
for item_id, document, distance in zip(
    filtered_results["ids"][0],
    filtered_results["documents"][0],
    filtered_results["distances"][0],
):
    print(f"distance={distance:.4f} id={item_id} text={document}")

查询端继续使用 Qwen3 的 query prompt,文档端不使用这个 prompt。执行:

python rag/p06_chroma_native/02_query_and_filter.py

真实输出为:

查询:周末想找一个适合散步的地方
全部分类的 Top 3:
distance=0.3502 id=doc-1 category=travel text=西湖边有步道和树荫,很适合周末散步。
distance=0.5340 id=doc-4 category=travel text=杭州植物园环境安静,适合慢慢游览。
distance=0.7817 id=doc-2 category=food text=苹果富含膳食纤维,是常见的健康水果。

只查询 category=travel:
distance=0.3502 id=doc-1 text=西湖边有步道和树荫,很适合周末散步。
distance=0.5340 id=doc-4 text=杭州植物园环境安静,适合慢慢游览。

FAISS 第一名的余弦相似度是 0.6498,Chroma 第一名的 cosine distance 是 0.3502:

1 - 0.6498 = 0.3502

本次四条数据实测得到相同排序,只是返回值的表达方式不同。FAISS 示例使用精确索引,Chroma 使用 HNSW;数据规模和索引参数变化后,两者的 Top-K 不保证始终完全一致。

where={“category”: “travel”} 会先限定业务分类,再返回符合条件的相似结果。相比 FAISS 需要在应用层自己维护元数据映射,Chroma 的过滤接口更直接。

6. 更新、插入和删除记录

为了不破坏前面的 life_notes,CRUD 示例使用独立的 life_notes_crud Collection。

代码文件为 03_update_and_delete.py:

# 这个文件使用独立的 Chroma Collection 演示 get、update、upsert、delete 和持久化重开。

from pathlib import Path
import shutil

import chromadb


PROJECT_ROOT = Path(__file__).resolve().parents[2]
DATA_DIR = PROJECT_ROOT / "data" / "p06_chroma_native" / "crud"
COLLECTION_NAME = "life_notes_crud"

# CRUD 示例使用独立目录,每次从空数据库开始,保证输出稳定。
if DATA_DIR.exists():
    shutil.rmtree(DATA_DIR)
DATA_DIR.mkdir(parents=True, exist_ok=True)
client = chromadb.PersistentClient(path=str(DATA_DIR))
collection = client.create_collection(name=COLLECTION_NAME)
collection.add(
    ids=["note-1", "note-2", "note-3"],
    documents=["去西湖散步", "买一些苹果", "检查汽车轮胎"],
    metadatas=[
        {"category": "travel"},
        {"category": "food"},
        {"category": "car"},
    ],
    embeddings=[[1.0, 0.0], [0.0, 1.0], [-1.0, 0.0]],
)

initial_ids = sorted(collection.get(include=["documents", "metadatas"])["ids"])
print("初始记录:", initial_ids)

collection.update(
    ids=["note-2"],
    documents=["买苹果和香蕉"],
    metadatas=[{"category": "food"}],
    embeddings=[[0.1, 0.9]],
)
collection.upsert(
    ids=["note-4"],
    documents=["去植物园游览"],
    metadatas=[{"category": "travel"}],
    embeddings=[[0.9, 0.1]],
)
collection.delete(ids=["note-1"])

print("更新并删除后的数量:", collection.count())
remaining_ids = sorted(collection.get(include=["documents", "metadatas"])["ids"])
print("剩余记录:", remaining_ids)

# 创建新的客户端对象,验证数据已经写入磁盘。
reopened_client = chromadb.PersistentClient(path=str(DATA_DIR))
reopened_collection = reopened_client.get_collection(COLLECTION_NAME)
print("重新打开后的数量:", reopened_collection.count())
print("note-2:", reopened_collection.get(ids=["note-2"], include=["documents"])["documents"][0])

这里使用二维手工向量,只为了演示 CRUD,不参与真实语义检索。运行结果为:

初始记录: ['note-1', 'note-2', 'note-3']
更新并删除后的数量: 3
剩余记录: ['note-2', 'note-3', 'note-4']
重新打开后的数量: 3
note-2: 买苹果和香蕉

新的 PersistentClient 能读到修改后的三条记录,说明数据已经持久化。

7. add、update 和 upsert 的区别

方法 作用
add() 添加新 ID,ID 已存在时不能把它当普通更新使用
update() 更新已经存在的记录
upsert() ID 存在就更新,不存在就插入
delete() 按 ID 或过滤条件删除
get() 按 ID 或过滤条件读取
count() 返回 Collection 中的记录数

修改 Document 后,Embedding 也应该同步重新生成。否则数据库中的向量仍代表旧文本,语义查询结果会发生错配。

8. 常见问题

8.1 查询时自动下载了其他模型

如果没有传入 query_embeddings,又使用了默认 Embedding Function,Chroma 可能尝试加载默认模型。本篇始终显式传入 Qwen3 生成的向量。

8.2 Collection not found

查询脚本依赖第一步创建的 life_notes。请按顺序执行:

python rag/p06_chroma_native/01_create_persistent_collection.py
python rag/p06_chroma_native/02_query_and_filter.py

8.3 把 distance 当成 similarity

本文使用 cosine HNSW,返回的是 distance:

distance 越小 -> 越接近

8.4 只复制了 chroma.sqlite3

Chroma 的持久化目录还包含 HNSW 文件。迁移或备份时应停止写入并保存整个 data/p06_chroma_native/,不能把 SQLite 与 HNSW 文件拆开复制。

8.5 同一个 ID 重复添加

需要更新时使用 update(),需要兼容新增和更新时使用 upsert()。不要依赖重复 add() 覆盖旧记录。

8.6 修改了 Document 却没有更新 Embedding

Chroma 不会根据手工传入的新 Document 自动重算应用自己管理的向量。如果文本发生变化,必须使用同一个 Embedding 模型重新生成向量,并在同一次更新中写入 Document 和 Embedding。

9. 小结

检查项 结果
Chroma 版本 1.5.0
客户端 PersistentClient
索引空间 cosine HNSW
Embedding 模型 本地 Qwen3-Embedding-0.6B
向量维度 1024
初始记录数 4
Top-K 查询 成功,本次四条数据的排序与 FAISS 一致
元数据过滤 category=travel 返回 2 条
CRUD get、update、upsert、delete、count 均成功
持久化重开 成功

Chroma 比 FAISS 多承担了文本、元数据和 CRUD 管理,但仍可以很轻量地运行在本地目录中。下一篇进入 Milvus Standalone,把向量数据库作为独立服务部署和访问。


文章作者: hnbian
版权声明: 本博客所有文章除特別声明外,均采用 CC BY 4.0 许可协议。转载请注明来源 hnbian !
评论
  目录