ChromaDB：轻量级向量数据库实战
--- title: "ChromaDB：轻量级向量数据库实战" description: "掌握ChromaDB的核心用法，从安装部署到构建完整的RAG应用，适合快速原型开发" tags: ["ChromaDB", "向量数据库", "RAG", "嵌入"] category: "llm" icon: "🧠"
ChromaDB：轻量级向量数据库实战
ChromaDB简介
ChromaDB是一个开源的向量数据库，专为AI应用设计。它提供了简洁的Python API，支持多种嵌入模型，是构建RAG应用的理想选择。ChromaDB的核心优势在于易用性和轻量级部署，无需复杂的配置即可快速开始。
与传统向量数据库相比，ChromaDB具有以下特点：
- 零配置启动：支持内存模式，无需外部服务
- Python原生：与Python生态系统深度集成
- 持久化存储：支持数据持久化到本地磁盘
- 多嵌入模型：内置OpenAI、Sentence Transformers等模型支持
安装与配置
安装ChromaDB
pip install chromadb
# 可选：安装OpenAI嵌入模型支持
pip install chromadb[openai]
基本使用
import chromadb
# 创建客户端（内存模式）
client = chromadb.Client()
# 创建集合
collection = client.create_collection(
name="documents",
metadata={"hnsw:space": "cosine"}
)
# 添加文档
collection.add(
documents=["机器学习是AI的核心技术", "深度学习是机器学习的子集"],
metadatas=[{"source": "tech"}, {"source": "tech"}],
ids=["doc1", "doc2"]
)
# 查询
results = collection.query(
query_texts=["什么是人工智能"],
n_results=2
)
print(results)
核心功能详解
文档管理
ChromaDB支持批量添加、更新和删除文档：
# 批量添加
collection.add(
documents=["文档1", "文档2", "文档3"],
metadatas=[
{"category": "tech", "year": 2024},
{"category": "science", "year": 2024},
{"category": "tech", "year": 2023}
],
ids=["id1", "id2", "id3"]
)
# 更新文档
collection.update(
ids=["id1"],
documents=["更新后的文档内容"]
)
# 删除文档
collection.delete(ids=["id2"])
元数据过滤
支持基于元数据的精确过滤和范围查询：
# 精确过滤
results = collection.query(
query_texts=["查询"],
where={"category": "tech"}
)
# 范围过滤
results = collection.query(
query_texts=["查询"],
where={"year": {"$gte": 2023}}
)
# 组合过滤
results = collection.query(
query_texts=["查询"],
where={
"$and": [
{"category": "tech"},
{"year": {"$gte": 2023}}
]
}
)
自定义嵌入函数
ChromaDB支持自定义嵌入模型，可集成任意Embedding模型：
from chromadb.utils import embedding_functions
# 使用Sentence Transformers
ef = embedding_functions.SentenceTransformerEmbeddingFunction(
model_name="all-MiniLM-L6-v2"
)
collection = client.create_collection(
name="custom_embeddings",
embedding_function=ef
)
# 使用OpenAI
openai_ef = embedding_functions.OpenAIEmbeddingFunction(
api_key="YOUR_API_KEY",
model_name="text-embedding-ada-002"
)
持久化存储
ChromaDB支持数据持久化，确保数据不丢失：
import chromadb
# 持久化到指定目录
client = chromadb.PersistentClient(path="./chroma_db")
# 创建或获取集合
collection = client.get_or_create_collection("persistent_docs")
# 添加的数据会自动持久化
collection.add(
documents=["持久化文档"],
ids=["persistent_doc1"]
)
构建RAG应用
以下是使用ChromaDB构建简单RAG应用的完整示例：
import chromadb
from chromadb.utils import embedding_functions
from openai import OpenAI
# 初始化组件
client = chromadb.PersistentClient(path="./rag_db")
ef = embedding_functions.OpenAIEmbeddingFunction(
api_key="YOUR_API_KEY"
)
# 创建知识库
collection = client.get_or_create_collection(
name="knowledge_base",
embedding_function=ef
)
# 添加知识文档
documents = [
"Python是一种解释型编程语言",
"机器学习需要大量训练数据",
"深度学习使用神经网络模型"
]
collection.add(
documents=documents,
ids=[f"doc_{i}" for i in range(len(documents))]
)
# RAG查询函数
def rag_query(question, top_k=3):
# 检索相关文档
results = collection.query(
query_texts=[question],
n_results=top_k
)
context = "\n".join(results["documents"][0])
# 使用LLM生成回答
openai_client = OpenAI(api_key="YOUR_API_KEY")
response = openai_client.chat.completions.create(
model="gpt-3.5-turbo",
messages=[
{"role": "system", "content": f"基于以下上下文回答问题：\n{context}"},
{"role": "user", "content": question}
]
)
return response.choices[0].message.content
# 使用示例
answer = rag_query("Python是什么？")
print(answer)
性能优化
批量操作
# 批量添加比逐条添加快10倍以上
batch_size = 100
for i in range(0, len(documents), batch_size):
batch = documents[i:i+batch_size]
collection.add(
documents=batch,
ids=[f"doc_{i+j}" for j in range(len(batch))]
)
索引优化
# 自定义HNSW参数
collection = client.create_collection(
name="optimized",
metadata={
"hnsw:space": "cosine",
"hnsw:M": 16, # 每个节点的连接数
"hnsw:construction_ef": 200, # 构建时的搜索范围
"hnsw:search_ef": 100 # 查询时的搜索范围
}
)
ChromaDB以其简洁的设计和强大的功能，成为快速构建AI应用的首选向量数据库。