向量数据库：ChromaDB 实战

向量数据库：ChromaDB 实战#

在上一章，我们用 Python 列表实现了一个简单的向量索引。它功能正确，但有明显局限：

  • 每次启动都要重新计算所有 embedding

  • 无法按元数据（如来源、分类）过滤

  • 文档量大时检索变慢

向量数据库专为 embedding 检索而生，解决了上述问题：

能力

描述

持久化存储

embedding 写入磁盘，重启后无需重算

元数据过滤

先过滤再检索，精准锁定文档子集

大规模检索

用 HNSW 等索引结构加速，百万级文档毫秒响应

批量操作

高效批量插入、更新、删除

本章使用 ChromaDB——最流行的开源向量数据库之一，可本地运行，无需外部服务。

文档 → embedding → ChromaDB collection → 相似度检索 → 结果
                         ↑
                    元数据过滤

import chromadb
import json
import os
import time
import litellm
from dotenv import load_dotenv

load_dotenv()

MODEL = os.getenv("LLM_MODEL")
EMBED_MODEL = os.getenv("EMBED_MODEL", "openai/text-embedding-3-small")

print(f"LLM 模型: {MODEL}")
print(f"Embedding 模型: {EMBED_MODEL}")
print(f"ChromaDB 版本: {chromadb.__version__}")

LLM 模型: openai/gpt-5-mini
Embedding 模型: openai/text-embedding-3-small
ChromaDB 版本: 1.5.5

辅助函数：获取 Embedding#

与上一章相同，我们封装 litellm.embedding 调用。 这里额外提供批量版本 get_embeddings，一次 API 调用处理多条文本，效率更高。

def get_embedding(text: str) -> list[float]:
    """获取单条文本的 embedding 向量。"""
    response = litellm.embedding(model=EMBED_MODEL, input=[text])
    return response.data[0]["embedding"]


def get_embeddings(texts: list[str]) -> list[list[float]]:
    """批量获取多条文本的 embedding 向量。
    
    一次 API 调用处理所有文本，比逐条调用高效得多。
    """
    response = litellm.embedding(model=EMBED_MODEL, input=texts)
    # 按 index 排序，确保顺序与输入一致
    items = sorted(response.data, key=lambda x: x["index"])
    return [item["embedding"] for item in items]


# 快速验证
test_emb = get_embedding("向量数据库")
print(f"单条 embedding 维度: {len(test_emb)}")

test_embs = get_embeddings(["人工智能", "深度学习"])
print(f"批量 embedding 数量: {len(test_embs)}，每条维度: {len(test_embs[0])}")

单条 embedding 维度: 1536

批量 embedding 数量: 2，每条维度: 1536

第一节：ChromaDB 基础用法#

ChromaDB 的核心概念：

  • Client：数据库连接（内存版 / 持久化版）

  • Collection：类似关系型数据库的「表」，存储 embedding + 元数据 + 原始文本

  • Document：一条记录，包含 id、embedding、document（原文）、metadata

我们先用内存版 Client 快速上手，数据存在内存中，进程退出后自动清除。

# 创建内存版 ChromaDB 客户端
client = chromadb.Client()

# 创建 collection，指定余弦相似度（embedding 检索推荐）
collection = client.create_collection(
    name="demo",
    metadata={"hnsw:space": "cosine"}  # 距离越小 = 越相似
)

print(f"Collection 创建成功: {collection.name}")

# 准备 10 条文档，涵盖 AI、烹饪、体育三个主题
documents = [
    # AI 主题
    "神经网络是深度学习的基础，由多层感知机组成，通过反向传播算法训练。",
    "Transformer 架构引入自注意力机制，彻底改变了自然语言处理领域。",
    "大语言模型通过海量文本预训练，具备强大的语言理解和生成能力。",
    "强化学习让智能体在环境中通过试错学习最优策略，AlphaGo 是典型应用。",
    # 烹饪主题
    "红烧肉是中国传统名菜，以五花肉为主料，加入酱油、糖、料酒慢炖而成。",
    "意大利面的种类繁多，最常见的有意大利面、蝴蝶结面和通心粉。",
    "寿司是日本料理的代表，以醋饭为基础，搭配生鱼片、虾、蟹等食材。",
    # 体育主题
    "足球是世界上最受欢迎的运动，每四年举办一次的世界杯是最大赛事。",
    "篮球由詹姆斯·奈史密斯于1891年发明，NBA 是全球最高水平的篮球联赛。",
    "马拉松全程为42.195公里，考验运动员的耐力和意志力。",
]

metadatas = [
    {"topic": "AI", "source": "教材"},
    {"topic": "AI", "source": "论文"},
    {"topic": "AI", "source": "教材"},
    {"topic": "AI", "source": "百科"},
    {"topic": "烹饪", "source": "菜谱"},
    {"topic": "烹饪", "source": "菜谱"},
    {"topic": "烹饪", "source": "百科"},
    {"topic": "体育", "source": "百科"},
    {"topic": "体育", "source": "百科"},
    {"topic": "体育", "source": "教材"},
]

ids = [f"doc_{i}" for i in range(len(documents))]

# 批量获取 embedding
print("正在计算 embedding...")
embeddings = get_embeddings(documents)

# 添加到 collection
collection.add(
    ids=ids,
    embeddings=embeddings,
    documents=documents,
    metadatas=metadatas,
)

print(f"已添加 {collection.count()} 条文档到 collection")

Collection 创建成功: demo
正在计算 embedding...

已添加 10 条文档到 collection

# 查询：搜索与「神经网络训练」最相关的文档
query_text = "神经网络训练"
query_embedding = get_embedding(query_text)

results = collection.query(
    query_embeddings=[query_embedding],
    n_results=3,
    include=["documents", "metadatas", "distances"],
)

print(f"查询: \"{query_text}\"")
print("=" * 60)
print(f"{'排名':<4} {'距离':<8} {'主题':<6} {'文档内容'}")
print("-" * 60)
for rank, (doc, meta, dist) in enumerate(
    zip(results["documents"][0], results["metadatas"][0], results["distances"][0]),
    start=1,
):
    print(f"{rank:<4} {dist:<8.4f} {meta['topic']:<6} {doc[:35]}...")

print()
print("注意：余弦距离范围 [0, 2]，0 = 完全相同，2 = 完全相反")

查询: "神经网络训练"
============================================================
排名   距离       主题     文档内容
------------------------------------------------------------
1    0.3496   AI     神经网络是深度学习的基础，由多层感知机组成，通过反向传播算法训练。...
2    0.6391   AI     大语言模型通过海量文本预训练，具备强大的语言理解和生成能力。...
3    0.6867   AI     强化学习让智能体在环境中通过试错学习最优策略，AlphaGo 是典型应...

注意：余弦距离范围 [0, 2]，0 = 完全相同，2 = 完全相反

第二节：元数据过滤#

向量数据库的强大之处在于语义检索 + 精确过滤的结合。

场景：用户只想在「AI」主题的文档中搜索，不希望烹饪或体育内容混入结果。

ChromaDB 使用 where 参数实现元数据过滤，支持 $eq、$ne、$in、$gt 等操作符。

全量文档 → [where 过滤] → 候选子集 → [embedding 相似度] → Top-K 结果

query_text = "学习算法"
query_embedding = get_embedding(query_text)

print(f"查询: \"{query_text}\"")
print()

# --- 不过滤：全库检索 ---
results_all = collection.query(
    query_embeddings=[query_embedding],
    n_results=3,
    include=["documents", "metadatas", "distances"],
)

print("【无过滤】全库 Top-3:")
for doc, meta, dist in zip(
    results_all["documents"][0],
    results_all["metadatas"][0],
    results_all["distances"][0],
):
    print(f"  [{meta['topic']:4s}] 距离={dist:.4f}  {doc[:40]}...")

print()

# --- 过滤 topic=AI ---
results_ai = collection.query(
    query_embeddings=[query_embedding],
    n_results=3,
    where={"topic": {"$eq": "AI"}},  # 只在 AI 文档中搜索
    include=["documents", "metadatas", "distances"],
)

print("【过滤 topic=AI】Top-3:")
for doc, meta, dist in zip(
    results_ai["documents"][0],
    results_ai["metadatas"][0],
    results_ai["distances"][0],
):
    print(f"  [{meta['topic']:4s}] 距离={dist:.4f}  {doc[:40]}...")

print()

# --- 组合过滤：source=百科 且 topic 在 [AI, 体育] 中 ---
results_multi = collection.query(
    query_embeddings=[query_embedding],
    n_results=3,
    where={
        "$and": [
            {"source": {"$eq": "百科"}},
            {"topic": {"$in": ["AI", "体育"]}},
        ]
    },
    include=["documents", "metadatas", "distances"],
)

print("【过滤 source=百科 AND topic in [AI,体育]】Top-3:")
for doc, meta, dist in zip(
    results_multi["documents"][0],
    results_multi["metadatas"][0],
    results_multi["distances"][0],
):
    print(f"  [{meta['topic']:4s}/{meta['source']}] 距离={dist:.4f}  {doc[:35]}...")

查询: "学习算法"

【无过滤】全库 Top-3:
  [AI  ] 距离=0.6448  强化学习让智能体在环境中通过试错学习最优策略，AlphaGo 是典型应用。...
  [AI  ] 距离=0.7395  神经网络是深度学习的基础，由多层感知机组成，通过反向传播算法训练。...
  [AI  ] 距离=0.7572  大语言模型通过海量文本预训练，具备强大的语言理解和生成能力。...

【过滤 topic=AI】Top-3:
  [AI  ] 距离=0.6448  强化学习让智能体在环境中通过试错学习最优策略，AlphaGo 是典型应用。...
  [AI  ] 距离=0.7395  神经网络是深度学习的基础，由多层感知机组成，通过反向传播算法训练。...
  [AI  ] 距离=0.7572  大语言模型通过海量文本预训练，具备强大的语言理解和生成能力。...

【过滤 source=百科 AND topic in [AI,体育]】Top-3:
  [AI  /百科] 距离=0.6448  强化学习让智能体在环境中通过试错学习最优策略，AlphaGo 是典型应...
  [体育  /百科] 距离=0.9172  篮球由詹姆斯·奈史密斯于1891年发明，NBA 是全球最高水平的篮球联...
  [体育  /百科] 距离=0.9573  足球是世界上最受欢迎的运动，每四年举办一次的世界杯是最大赛事。...

第三节：持久化存储#

内存版 Client 适合测试，生产环境需要持久化存储：

  • 数据写入磁盘，进程重启后数据依然存在

  • 无需重新计算 embedding，节省 API 费用

  • 多次运行程序共享同一份数据

只需将 chromadb.Client() 替换为 chromadb.PersistentClient(path=...)。

import shutil

PERSIST_PATH = "/tmp/chroma_demo"

# 清理可能存在的旧数据
if os.path.exists(PERSIST_PATH):
    shutil.rmtree(PERSIST_PATH)

# ── 第一步：写入数据 ──────────────────────────────────────
print("[步骤 1] 创建持久化 Client，写入数据...")
persist_client = chromadb.PersistentClient(path=PERSIST_PATH)
persist_col = persist_client.create_collection(
    name="knowledge_base",
    metadata={"hnsw:space": "cosine"},
)

docs_to_persist = [
    "RAG（检索增强生成）将知识库检索与大语言模型生成结合，解决知识截止问题。",
    "向量数据库是 RAG 系统的核心组件，负责高效存储和检索 embedding。",
    "文本切块策略直接影响 RAG 效果，块太大丢失精度，块太小丢失上下文。",
]
embs = get_embeddings(docs_to_persist)
persist_col.add(
    ids=["p0", "p1", "p2"],
    embeddings=embs,
    documents=docs_to_persist,
    metadatas=[{"chapter": "4"}, {"chapter": "4"}, {"chapter": "4"}],
)
print(f"  写入 {persist_col.count()} 条文档")
print(f"  数据目录: {PERSIST_PATH}")

# 删除客户端对象（模拟进程退出）
del persist_client, persist_col
print("  客户端已关闭（模拟进程退出）")

# ── 第二步：重新打开，验证数据 ────────────────────────────
print()
print("[步骤 2] 重新打开持久化 Client，验证数据...")
persist_client2 = chromadb.PersistentClient(path=PERSIST_PATH)
persist_col2 = persist_client2.get_collection("knowledge_base")
print(f"  重新打开后文档数量: {persist_col2.count()}  ✓ 数据持久化成功！")

# 验证查询仍然正常
q_emb = get_embedding("RAG 系统")
res = persist_col2.query(query_embeddings=[q_emb], n_results=1, include=["documents"])
print(f"  查询 'RAG 系统' 最相关: {res['documents'][0][0][:40]}...")

# 清理
del persist_client2, persist_col2
shutil.rmtree(PERSIST_PATH)
print()
print("[清理] 已删除临时目录")

[步骤 1] 创建持久化 Client，写入数据...

  写入 3 条文档
  数据目录: /tmp/chroma_demo
  客户端已关闭（模拟进程退出）

[步骤 2] 重新打开持久化 Client，验证数据...
  重新打开后文档数量: 3  ✓ 数据持久化成功！

  查询 'RAG 系统' 最相关: 向量数据库是 RAG 系统的核心组件，负责高效存储和检索 embedding。...

[清理] 已删除临时目录

第四节：Collection 管理#

ChromaDB 提供完整的 Collection 生命周期管理接口，类似关系型数据库的 DDL 操作。

# 用同一个内存 client 演示管理操作
mgmt_client = chromadb.Client()

# 创建两个 collection
col_a = mgmt_client.create_collection("collection_a")
col_b = mgmt_client.create_collection("collection_b")

# 向 col_a 添加一些数据
sample_docs = ["Python 是一种高级编程语言。", "机器学习是 AI 的子领域。", "深度学习使用多层神经网络。"]
sample_embs = get_embeddings(sample_docs)
col_a.add(ids=["s0", "s1", "s2"], embeddings=sample_embs, documents=sample_docs)

print("── list_collections() ──")
collections = mgmt_client.list_collections()
print(f"  现有 collection: {[c.name for c in collections]}")

print()
print("── get_collection() ──")
fetched = mgmt_client.get_collection("collection_a")
print(f"  获取到: {fetched.name}")

print()
print("── count() ──")
print(f"  collection_a 文档数: {col_a.count()}")
print(f"  collection_b 文档数: {col_b.count()}")

print()
print("── peek() ── 查看前 2 条记录")
peeked = col_a.peek(limit=2)
for pid, pdoc in zip(peeked["ids"], peeked["documents"]):
    print(f"  id={pid}  doc={pdoc}")

print()
print("── get() ── 按 id 精确获取")
got = col_a.get(ids=["s1"], include=["documents", "metadatas"])
print(f"  id=s1: {got['documents'][0]}")

print()
print("── delete_collection() ──")
mgmt_client.delete_collection("collection_b")
remaining = mgmt_client.list_collections()
print(f"  删除后剩余: {[c.name for c in remaining]}")

print()
print("── update() ── 更新已有文档")
new_emb = get_embedding("Python 是一种解释型高级编程语言，以简洁著称。")
col_a.update(
    ids=["s0"],
    embeddings=[new_emb],
    documents=["Python 是一种解释型高级编程语言，以简洁著称。"],
)
updated = col_a.get(ids=["s0"], include=["documents"])
print(f"  更新后 s0: {updated['documents'][0]}")

print()
print("── delete() ── 删除指定文档")
col_a.delete(ids=["s2"])
print(f"  删除 s2 后文档数: {col_a.count()}")

── list_collections() ──
  现有 collection: ['collection_a', 'collection_b', 'demo']

── get_collection() ──
  获取到: collection_a

── count() ──
  collection_a 文档数: 3
  collection_b 文档数: 0

── peek() ── 查看前 2 条记录
  id=s0  doc=Python 是一种高级编程语言。
  id=s1  doc=机器学习是 AI 的子领域。

── get() ── 按 id 精确获取
  id=s1: 机器学习是 AI 的子领域。

── delete_collection() ──
  删除后剩余: ['collection_a', 'demo']

── update() ── 更新已有文档
  更新后 s0: Python 是一种解释型高级编程语言，以简洁著称。

── delete() ── 删除指定文档
  删除 s2 后文档数: 2

第五节：批量操作性能对比#

真实场景中，知识库往往有数千甚至数万条文档。
一条条插入 vs 批量插入，性能差距显著。

此节演示两种方式的耗时对比，帮助你选择正确的插入策略。

注意：embedding API 调用是最大瓶颈。本实验使用预先批量获取的 embedding， 专门衡量 ChromaDB 插入层面的性能差异。

import random

# 生成 50 条模拟文档
TOTAL_DOCS = 50
topics_pool = [
    "机器学习模型训练需要大量标注数据和计算资源。",
    "自然语言处理研究计算机理解和生成人类语言的方法。",
    "计算机视觉让机器能够理解和分析图像与视频。",
    "强化学习通过奖励信号引导智能体学习最优行为。",
    "知识图谱以结构化方式表示实体和关系。",
]
bulk_docs = [
    f"{random.choice(topics_pool)} 文档编号 {i}。" for i in range(TOTAL_DOCS)
]
bulk_ids = [f"bulk_{i}" for i in range(TOTAL_DOCS)]
bulk_metas = [{"batch": i // 10} for i in range(TOTAL_DOCS)]

print(f"准备 {TOTAL_DOCS} 条文档，批量获取 embedding...")
bulk_embeddings = get_embeddings(bulk_docs)
print(f"Embedding 计算完成，开始对比插入性能\n")

# ── 方式一：逐条插入 ──────────────────────────────────────
col_one_by_one = chromadb.Client().create_collection("one_by_one")
t0 = time.perf_counter()
for i in range(TOTAL_DOCS):
    col_one_by_one.add(
        ids=[bulk_ids[i]],
        embeddings=[bulk_embeddings[i]],
        documents=[bulk_docs[i]],
        metadatas=[bulk_metas[i]],
    )
t_one = time.perf_counter() - t0

# ── 方式二：一次性批量插入 ────────────────────────────────
col_batch_all = chromadb.Client().create_collection("batch_all")
t0 = time.perf_counter()
col_batch_all.add(
    ids=bulk_ids,
    embeddings=bulk_embeddings,
    documents=bulk_docs,
    metadatas=bulk_metas,
)
t_batch_all = time.perf_counter() - t0

# ── 方式三：分批插入（每批 10 条）────────────────────────
BATCH_SIZE = 10
col_batched = chromadb.Client().create_collection("batched")
t0 = time.perf_counter()
for start in range(0, TOTAL_DOCS, BATCH_SIZE):
    end = min(start + BATCH_SIZE, TOTAL_DOCS)
    col_batched.add(
        ids=bulk_ids[start:end],
        embeddings=bulk_embeddings[start:end],
        documents=bulk_docs[start:end],
        metadatas=bulk_metas[start:end],
    )
t_batched = time.perf_counter() - t0

print("性能对比结果（ChromaDB 插入层面，不含 embedding API 耗时）")
print("=" * 55)
print(f"{'方式':<20} {'耗时(ms)':<12} {'文档数/ms':<12} {'加速比'}")
print("-" * 55)
print(f"{'逐条插入':<20} {t_one*1000:<12.1f} {TOTAL_DOCS/(t_one*1000):<12.2f} 1.0x")
print(f"{'一次性批量':<20} {t_batch_all*1000:<12.1f} {TOTAL_DOCS/(t_batch_all*1000):<12.2f} {t_one/t_batch_all:.1f}x")
print(f"{'分批(每批10条)':<20} {t_batched*1000:<12.1f} {TOTAL_DOCS/(t_batched*1000):<12.2f} {t_one/t_batched:.1f}x")
print()
print("结论：批量插入比逐条插入快数倍，生产环境务必使用批量模式。")
print(f"实际项目推荐批大小：100-1000 条（视内存和 API 限制调整）。")

准备 50 条文档，批量获取 embedding...

Embedding 计算完成，开始对比插入性能

性能对比结果（ChromaDB 插入层面，不含 embedding API 耗时）
=======================================================
方式                   耗时(ms)       文档数/ms       加速比
-------------------------------------------------------
逐条插入                 48.6         1.03         1.0x
一次性批量                7.4          6.78         6.6x
分批(每批10条)            12.0         4.18         4.1x

结论：批量插入比逐条插入快数倍，生产环境务必使用批量模式。
实际项目推荐批大小：100-1000 条（视内存和 API 限制调整）。

总结#

本章要点#

  1. ChromaDB 基础：Client → Collection → Document，三层结构清晰

  2. 元数据过滤：where 参数支持精确过滤，与语义检索无缝结合

  3. 持久化存储：PersistentClient 一行代码切换，数据跨进程保留

  4. Collection 管理：完整 CRUD 接口，peek()/count()/get() 方便调试

  5. 批量操作：插入时使用批量 API，性能提升显著

向量数据库横向对比#

特性

SimpleVectorIndex（自建）

ChromaDB

Pinecone

Weaviate

Qdrant

部署方式

纯 Python，无依赖

本地 / 服务器

云托管

本地 / 云

本地 / 云

持久化

手动序列化

内置支持

内置支持

内置支持

内置支持

元数据过滤

不支持

支持

支持

支持（GraphQL）

支持

最大规模

万级

百万级

十亿级

亿级

亿级

适合场景

学习 / 原型

开发 / 中小规模

大规模生产

企业级

高性能生产

费用

免费

开源免费

按量付费

开源 / 商业版

开源免费

下一章预告#

我们已经分别学习了：

  • 文档加载与切块

  • Embedding 原理与生成

  • 向量数据库存储与检索

第四章将把这三项技术组合成完整的 RAG 系统，实现端到端的知识库问答！