跳到主要内容

Chroma 向量数据库

Chroma 是一个专为 AI 应用设计的开源向量数据库，以简单易用为核心设计理念。它是 Python 生态中最友好的向量数据库，特别适合快速原型开发和小型项目。

概述​

为什么选择 Chroma​

特性说明
极简 API几行代码即可启动，学习成本极低
Python 原生与 Python 生态无缝集成
本地优先无需外部依赖，本地文件存储
嵌入式可嵌入到应用程序中运行
自动嵌入内置嵌入模型支持，无需手动编码
轻量级资源占用少，适合边缘设备

适用场景​

  • 快速原型：验证想法，快速搭建 Demo
  • 本地开发：无需配置服务器
  • 小型应用：数据量在百万级以下
  • 边缘部署：资源受限环境
  • 学习实验：理解向量数据库概念

快速开始​

1. 安装​

pip install chromadb

# 如果需要内置嵌入模型
pip install chromadb[sentence-transformers]

2. 第一个示例​

import chromadb

# 创建客户端（本地模式）
client = chromadb.Client()

# 创建集合（类似数据库表）
collection = client.create_collection(name="my_collection")

# 添加文档
collection.add(
    documents=["这是第一篇文档", "这是第二篇文档", "这是第三篇文档"],
    metadatas=[{"category": "tech"}, {"category": "life"}, {"category": "tech"}],
    ids=["doc1", "doc2", "doc3"]
)

# 查询
results = collection.query(
    query_texts=["技术相关文档"],
    n_results=2
)

print(results)

3. 使用持久化存储​

import chromadb

# 创建持久化客户端
client = chromadb.PersistentClient(path="./chroma_db")

# 后续操作与内存模式相同
collection = client.create_collection(name="my_collection")

核心概念​

Client（客户端）​

Chroma 提供三种客户端模式：

import chromadb

# 1. 内存客户端（数据不保存）
client = chromadb.Client()

# 2. 持久化客户端（数据保存到磁盘）
client = chromadb.PersistentClient(path="./chroma_data")

# 3. HTTP 客户端（连接远程服务器）
client = chromadb.HttpClient(host="localhost", port=8000)

Collection（集合）​

集合是 Chroma 中的数据组织单位，包含文档、嵌入和元数据。

# 创建集合
collection = client.create_collection(
    name="documents",
    metadata={"description": "文档集合"}
)

# 获取或创建（如果不存在则创建）
collection = client.get_or_create_collection(name="documents")

# 获取现有集合
collection = client.get_collection(name="documents")

# 删除集合
client.delete_collection(name="documents")

# 列出所有集合
print(client.list_collections())

嵌入函数（Embedding Function）​

Chroma 支持自动嵌入，无需手动调用嵌入模型。

from chromadb.utils import embedding_functions

# 默认嵌入（all-MiniLM-L6-v2）
default_ef = embedding_functions.DefaultEmbeddingFunction()

# Sentence Transformers
sentence_transformer_ef = embedding_functions.SentenceTransformerEmbeddingFunction(
    model_name="all-MiniLM-L6-v2"
)

# OpenAI
openai_ef = embedding_functions.OpenAIEmbeddingFunction(
    api_key="your-api-key",
    model_name="text-embedding-3-small"
)

# 创建集合时指定嵌入函数
collection = client.create_collection(
    name="my_collection",
    embedding_function=sentence_transformer_ef
)

数据操作​

添加数据（Add）​

# 基本添加
collection.add(
    documents=["文档1", "文档2", "文档3"],
    ids=["id1", "id2", "id3"]
)

# 带元数据
collection.add(
    documents=["iPhone 15 发布", "MacBook Pro 评测"],
    metadatas=[
        {"category": "electronics", "brand": "Apple", "price": 999},
        {"category": "electronics", "brand": "Apple", "price": 1999}
    ],
    ids=["product1", "product2"]
)

# 直接提供嵌入（跳过自动嵌入）
import numpy as np
embeddings = np.random.rand(2, 384).tolist()

collection.add(
    embeddings=embeddings,
    documents=["文档1", "文档2"],
    ids=["id1", "id2"]
)

# 批量添加（推荐）
batch_size = 100
for i in range(0, len(documents), batch_size):
    batch_docs = documents[i:i+batch_size]
    batch_ids = [f"id_{j}" for j in range(i, i+len(batch_docs))]
    collection.add(documents=batch_docs, ids=batch_ids)

查询（Query）​

文本查询（自动嵌入）​

# 基本查询
results = collection.query(
    query_texts=["苹果手机"],
    n_results=3
)

# 带过滤条件的查询
results = collection.query(
    query_texts=["高端电子产品"],
    n_results=5,
    where={"price": {"$gt": 1000}}
)

# 指定返回字段
results = collection.query(
    query_texts=["查询文本"],
    n_results=3,
    include=["documents", "metadatas", "distances"]
)

向量查询​

# 使用向量直接查询
query_embedding = [0.1, 0.2, 0.3, ...]  # 384维向量

results = collection.query(
    query_embeddings=[query_embedding],
    n_results=3
)

获取数据（Get）​

# 根据 ID 获取
results = collection.get(ids=["id1", "id2"])

# 获取所有数据
results = collection.get()

# 带过滤条件获取
results = collection.get(
    where={"category": "electronics"}
)

# 分页获取
results = collection.get(
    limit=10,
    offset=20
)

更新数据​

# 使用相同的 ID 更新
collection.update(
    ids=["id1"],
    documents=["更新后的文档内容"],
    metadatas=[{"updated": True}]
)

# upsert（存在则更新，不存在则插入）
collection.upsert(
    ids=["id1", "id_new"],
    documents=["更新内容", "新内容"],
    metadatas=[{"status": "updated"}, {"status": "new"}]
)

删除数据​

# 根据 ID 删除
collection.delete(ids=["id1", "id2"])

# 根据过滤条件删除
collection.delete(where={"category": "obsolete"})

# 删除所有数据
collection.delete()

过滤条件​

支持的运算符​

# 等于
{"category": "electronics"}

# 不等于
{"category": {"$ne": "clothing"}}

# 大于/大于等于
{"price": {"$gt": 100}}
{"price": {"$gte": 100}}

# 小于/小于等于
{"price": {"$lt": 1000}}
{"price": {"$lte": 1000}}

# 包含在列表中
{"category": {"$in": ["electronics", "computers"]}}

# 不包含在列表中
{"category": {"$nin": ["food", "book"]}}

组合条件​

# AND 条件
{
    "$and": [
        {"category": "electronics"},
        {"price": {"$lte": 1000}}
    ]
}

# OR 条件
{
    "$or": [
        {"brand": "Apple"},
        {"brand": "Samsung"}
    ]
}

# 复杂组合
{
    "$and": [
        {"category": "electronics"},
        {"$or": [
            {"brand": "Apple"},
            {"brand": "Samsung"}
        ]},
        {"price": {"$gte": 500, "$lte": 2000}}
    ]
}

完整 RAG 示例​

import chromadb
from openai import OpenAI
import os

class ChromaRAG:
    def __init__(self, collection_name="knowledge_base"):
        # 创建持久化客户端
        self.client = chromadb.PersistentClient(path="./chroma_db")
        
        # 获取或创建集合
        self.collection = self.client.get_or_create_collection(
            name=collection_name,
            metadata={"description": "RAG 知识库"}
        )
        
        # OpenAI 客户端
        self.openai = OpenAI(api_key=os.environ.get("OPENAI_API_KEY"))
    
    def add_documents(self, documents):
        """
        添加文档到知识库
        documents: [{"content": "...", "metadata": {...}}, ...]
        """
        texts = [doc["content"] for doc in documents]
        metadatas = [doc.get("metadata", {}) for doc in documents]
        ids = [f"doc_{i}" for i in range(len(documents))]
        
        self.collection.add(
            documents=texts,
            metadatas=metadatas,
            ids=ids
        )
        print(f"已添加 {len(documents)} 篇文档")
    
    def search(self, query, top_k=5, filters=None):
        """搜索相关文档"""
        where_clause = filters if filters else {}
        
        results = self.collection.query(
            query_texts=[query],
            n_results=top_k,
            where=where_clause if where_clause else None
        )
        
        documents = []
        for i in range(len(results["ids"][0])):
            documents.append({
                "id": results["ids"][0][i],
                "content": results["documents"][0][i],
                "metadata": results["metadatas"][0][i],
                "distance": results["distances"][0][i]
            })
        
        return documents
    
    def answer(self, question, top_k=3):
        """生成回答"""
        # 检索相关文档
        docs = self.search(question, top_k)
        
        # 构建上下文
        context = "\n\n".join([
            f"[文档 {i+1}] {doc['content']}"
            for i, doc in enumerate(docs)
        ])
        
        # 调用 LLM
        prompt = f"""基于以下文档回答问题。如果文档中没有相关信息，请说明。

文档：
{context}

问题：{question}

回答："""
        
        response = self.openai.chat.completions.create(
            model="gpt-3.5-turbo",
            messages=[
                {"role": "system", "content": "你是一个基于文档回答问题的助手。"},
                {"role": "user", "content": prompt}
            ]
        )
        
        return {
            "answer": response.choices[0].message.content,
            "sources": docs
        }

# 使用示例
rag = ChromaRAG()

# 添加文档
docs = [
    {"content": "Chroma 是一个开源的嵌入式向量数据库...", "metadata": {"source": "docs"}},
    {"content": "向量数据库用于 AI 应用的语义搜索...", "metadata": {"source": "blog"}}
]
rag.add_documents(docs)

# 提问
result = rag.answer("什么是 Chroma？")
print(f"回答：{result['answer']}")

部署方式​

嵌入式模式​

直接嵌入到 Python 应用中：

import chromadb

# 内存模式
client = chromadb.Client()

# 持久化模式
client = chromadb.PersistentClient(path="./db")

服务器模式​

# 安装 Chroma 服务器
pip install chromadb

# 启动服务器
chroma run --path ./chroma_data --port 8000
# 客户端连接
import chromadb

client = chromadb.HttpClient(host="localhost", port=8000)
collection = client.get_or_create_collection("my_collection")

Docker 部署​

# docker-compose.yml
version: '3.9'

services:
  chroma:
    image: chromadb/chroma:latest
    ports:
      - "8000:8000"
    volumes:
      - ./chroma_data:/chroma/chroma
    environment:
      - IS_PERSISTENT=TRUE
docker-compose up -d

性能优化​

批量操作​

# 批量添加比逐条添加快得多
batch_size = 1000

for i in range(0, len(documents), batch_size):
    batch = documents[i:i+batch_size]
    collection.add(
        documents=[d["text"] for d in batch],
        metadatas=[d["metadata"] for d in batch],
        ids=[f"id_{j}" for j in range(i, i+len(batch))]
    )

索引配置​

# Chroma 使用 HNSW 索引，可以通过参数调优
collection = client.create_collection(
    name="optimized",
    metadata={
        "hnsw:space": "cosine",  # 距离度量：cosine, l2, ip
        "hnsw:construction_ef": 128,  # 构建参数
        "hnsw:search_ef": 64,  # 搜索参数
        "hnsw:M": 16  # 连接数
    }
)

使用预计算嵌入​

from sentence_transformers import SentenceTransformer

model = SentenceTransformer('all-MiniLM-L6-v2')

# 预计算所有嵌入
embeddings = model.encode(documents, show_progress_bar=True).tolist()

# 直接添加嵌入（跳过 Chroma 的嵌入步骤）
collection.add(
    embeddings=embeddings,
    documents=documents,
    ids=ids
)

与其他工具集成​

LangChain 集成​

from langchain.vectorstores import Chroma
from langchain.embeddings import OpenAIEmbeddings
from langchain.text_splitter import CharacterTextSplitter
from langchain.document_loaders import TextLoader

# 加载文档
loader = TextLoader("document.txt")
documents = loader.load()

# 切分文档
text_splitter = CharacterTextSplitter(chunk_size=1000, chunk_overlap=0)
texts = text_splitter.split_documents(documents)

# 创建 Chroma 向量存储
embeddings = OpenAIEmbeddings()
vectorstore = Chroma.from_documents(texts, embeddings, persist_directory="./chroma_db")

# 相似度搜索
docs = vectorstore.similarity_search("查询内容", k=3)

LlamaIndex 集成​

from llama_index import VectorStoreIndex, SimpleDirectoryReader
from llama_index.vector_stores import ChromaVectorStore
from llama_index.storage.storage_context import StorageContext
import chromadb

# 创建 Chroma 客户端
chroma_client = chromadb.PersistentClient(path="./chroma_db")
chroma_collection = chroma_client.get_or_create_collection("llamaindex")

# 创建向量存储
vector_store = ChromaVectorStore(chroma_collection=chroma_collection)
storage_context = StorageContext.from_defaults(vector_store=vector_store)

# 加载文档
documents = SimpleDirectoryReader("./data").load_data()

# 创建索引
index = VectorStoreIndex.from_documents(documents, storage_context=storage_context)

# 查询
query_engine = index.as_query_engine()
response = query_engine.query("你的问题")

常见问题​

Q: Chroma 适合生产环境吗？​

Chroma 目前更适合：

  • 原型开发和概念验证
  • 小到中等规模的应用（百万级向量）
  • 单节点部署

对于大规模生产环境，建议使用 Milvus 或 Pinecone。

Q: 如何备份数据？​

import shutil
import json

# 持久化模式的备份
# 1. 直接复制数据目录
shutil.copytree("./chroma_db", "./chroma_backup")

# 2. 导出为 JSON
collection = client.get_collection("my_collection")
data = collection.get()

with open("backup.json", "w") as f:
    json.dump(data, f)

Q: 如何处理大量数据？​

# 1. 使用批量添加
# 2. 考虑使用服务器模式
# 3. 定期清理旧数据

# 清理示例：删除超过一定时间的数据
collection.delete(
    where={"timestamp": {"$lt": cutoff_time}}
)

Q: 如何迁移到 Milvus？​

# 1. 从 Chroma 导出
data = chroma_collection.get()

# 2. 导入到 Milvus
from pymilvus import Collection

milvus_collection = Collection("new_collection")
milvus_collection.insert([
    data["documents"],
    data["embeddings"]
])

下一步​