Chroma：开源嵌入式向量数据库

bash pip install chromadb

如果希望使用 Cohere、HuggingFace 等提供的嵌入模型，可额外安装：

```bash
pip install chromadb[all]

核心概念速览

在使用 Chroma 之前，需要理解几个关键组件：

• Collection（集合）：类似关系数据库中的“表”，用于组织和隔离向量数据。
• Document（文档）：原始文本或数据内容，会被自动转换为向量。
• Embedding（嵌入）：将文档映射到高维向量空间的表示，Chroma 内置多种嵌入函数。
• Metadata（元数据）：附加在文档上的结构化信息，用于过滤搜索。
• Query（查询）：输入文本或向量，返回语义最相似的文档。

五分钟上手：基本操作

下面的代码演示了 Chroma 的核心工作流。

创建客户端并获取或新建集合

import chromadb

# 创建持久化客户端（数据保存到磁盘）
client = chromadb.PersistentClient(path="./chroma_db")
# 或使用内存模式（数据在程序结束后丢失）
# client = chromadb.EphemeralClient()

# 获取或创建集合
collection = client.get_or_create_collection(name="my_docs")

添加文档

add 方法接受文档、元数据和唯一标识。未提供嵌入时，Chroma 会使用默认的 all-MiniLM-L6-v2 模型自动生成。

collection.add(
    documents=[
        "Redis 是一种高性能键值数据库",
        "Chroma 是轻量级向量数据库",
        "PostgreSQL 支持 JSON 查询",
    ],
    metadatas=[
        {"source": "redis"},
        {"source": "chroma"},
        {"source": "postgres"},
    ],
    ids=["doc1", "doc2", "doc3"],
)

查询语义相似文档

results = collection.query(
    query_texts=["适合缓存和会话存储的数据库"],
    n_results=2,
)

print(results["documents"])
# 可能输出：[['Redis 是一种高性能键值数据库', 'Chroma 是轻量级向量数据库']]

更新与删除

collection.update(
    ids=["doc1"],
    documents=["Redis 是超快的内存数据库，常用于缓存"],
)

collection.delete(ids=["doc3"])

深入掌握：高级功能详解

自定义嵌入函数

你可以指定自己的嵌入模型，例如使用 huggingface 或 openai。

from chromadb.utils import embedding_functions

# 使用 HuggingFace 模型
huggingface_ef = embedding_functions.HuggingFaceEmbeddingFunction(
    model_name="sentence-transformers/all-mpnet-base-v2"
)

collection = client.get_or_create_collection(
    name="custom_embeddings",
    embedding_function=huggingface_ef
)

对于 OpenAI 嵌入，需要先设置 API 密钥：

openai_ef = embedding_functions.OpenAIEmbeddingFunction(
    api_key="sk-...",
    model_name="text-embedding-ada-002"
)

元数据过滤

查询时可以通过 where 条件精确筛选。

results = collection.query(
    query_texts=["数据库"],
    n_results=5,
    where={"source": "chroma"}  # 仅返回 source 为 chroma 的文档
)

更复杂的过滤，如数值比较：

collection.add(
    documents=["产品 A", "产品 B"],
    metadatas=[{"price": 99}, {"price": 150}],
    ids=["p1", "p2"],
)

results = collection.query(
    query_texts=["便宜商品"],
    where={"price": {"$lt": 120}}  # 价格小于 120
)

直接使用向量

除了通过文本自动生成向量，也可以直接添加预先计算好的向量。

import numpy as np

embeddings = np.random.rand(3, 384).tolist()  # 384 维向量

collection.add(
    embeddings=embeddings,
    documents=["doc A", "doc B", "doc C"],
    ids=["va1", "va2", "va3"]
)

获取集合信息与数据

# 查看集合中的文档数量
print(collection.count())

# 获取指定 ID 的文档
data = collection.get(ids=["doc1", "doc2"])
print(data["documents"])

实战场景：构建文档问答系统

下面演示如何用 Chroma 作为语义记忆，让 LLM 基于私有文档回答问题。

from langchain.embeddings import HuggingFaceEmbeddings
from langchain.vectorstores import Chroma
from langchain.text_splitter import CharacterTextSplitter

# 假设 documents 是读取的长文本列表
text_splitter = CharacterTextSplitter(chunk_size=500, chunk_overlap=50)
docs = text_splitter.create_documents(documents)

# 创建 embedding 模型
embeddings = HuggingFaceEmbeddings(model_name="all-MiniLM-L6-v2")

# 将文档存入 Chroma
db = Chroma.from_documents(docs, embeddings, persist_directory="./chroma_db")

# 搜索相关上下文
retriever = db.as_retriever(search_kwargs={"k": 3})
context_docs = retriever.get_relevant_documents("如何使用 Chroma 过滤元数据？")

# 将 context_docs 作为提示词的一部分交给 LLM 生成答案