Chroma 向量数据库完全指南 - AI创想

# 基础安装
pip install chromadb
# 包含 CLI 工具
pip install chromadb[cli]# 完整安装（包含所有依赖）
pip install chromadb[all]

复制代码

import chromadb
from chromadb.utils import embedding_functions
# 创建客户端（内存模式）
client = chromadb.Client()# 创建或获取集合
collection = client.create_collection(
name="my_collection",
embedding_function=embedding_functions.SentenceTransformerEmbeddingFunction(
model_name="all-MiniLM-L6-v2"))

复制代码

# 准备数据
documents =["Chroma 是一个向量数据库","它专为 AI 应用设计","支持语义搜索和相似度匹配","可以用于构建问答系统","支持多种嵌入模型"]
ids =["doc1","doc2","doc3","doc4","doc5"]
metadatas =[{"source":"官方文档","type":"介绍"},{"source":"官方文档","type":"特性"},{"source":"官方文档","type":"功能"},{"source":"教程","type":"应用"},{"source":"教程","type":"特性"}]# 添加文档到集合
collection.add(
documents=documents,
metadatas=metadatas,
ids=ids
)

复制代码

# 基本查询
results = collection.query(
query_texts=["Chroma 有哪些功能？"],
n_results=3)print("查询结果:")for i, doc inenumerate(results['documents'][0]):print(f"{i+1}. {doc}")print(f" 元数据: {results['metadatas'][0][i]}")print(f" 距离: {results['distances'][0][i]}")print()

复制代码

# 带过滤条件的查询
results = collection.query(
query_texts=["向量数据库"],
n_results=5,
where={"source":"官方文档"},# 过滤条件
where_document={"$contains":"支持"}# 文档内容过滤)

复制代码

import openai
from chromadb.utils import embedding_functions
# 使用 OpenAI格式嵌入
openai_ef = embedding_functions.OpenAIEmbeddingFunction(
api_key="sk-*********************************",
api_base="https://api.siliconflow.cn/v1/",
model_name="BAAI/bge-m3")# 创建使用 OpenAI 嵌入的集合
collection = client.create_collection(
name="openai_collection",
embedding_function=openai_ef
)

复制代码

# 持久化客户端
persistent_client = chromadb.PersistentClient(path="./chroma_db")# 在指定路径创建集合
collection = persistent_client.create_collection("persistent_collection")# 数据会自动保存到磁盘

复制代码

# 更新文档
collection.update(
ids=["doc1"],
documents=["更新后的文档内容"],
metadatas=[{"source":"更新","version":"2.0"}])# 删除文档
collection.delete(ids=["doc2"])# 按条件删除
collection.delete(where={"source":"教程"})

复制代码

import os
import chromadb
from chromadb.utils import embedding_functions
openai_ef = embedding_functions.OpenAIEmbeddingFunction(
api_key="sk-******************************",
api_base="https://api.siliconflow.cn/v1/",
model_name="BAAI/bge-m3")classSimpleQASystem:def__init__(self, collection_name="qa_collection"):
self.client = chromadb.PersistentClient(path="./qa_db")
self.collection = self.client.get_or_create_collection(
name=collection_name,
embedding_function=openai_ef
)defadd_knowledge(self, documents, ids=None, metadatas=None):"""添加知识库文档"""if ids isNone:
ids =[f"doc_{i}"for i inrange(len(documents))]
self.collection.add(
documents=documents,
ids=ids,
metadatas=metadatas
)defquery(self, question, n_results=2):"""查询相关问题"""
results = self.collection.query(
query_texts=[question],
n_results=n_results
)return{'answers': results['documents'][0],'sources': results['metadatas'][0],'scores': results['distances'][0]}defget_context(self, question, max_length=500):"""获取相关上下文"""
results = self.query(question)
context =""for answer in results['answers']:iflen(context)+len(answer)< max_length:
context += answer +"\n\n"else:breakreturn context
# 使用示例
qa_system = SimpleQASystem()
qa_system.add_knowledge(["Chroma 支持元数据过滤，可以按条件查询文档","向量数据库主要用于存储和检索高维向量数据","语义搜索基于文本的相似度匹配"])
context = qa_system.get_context("Chroma 有什么功能？")print(f"相关上下文:\n{context}")

复制代码

import os
import chromadb
from typing import List
from chromadb.utils import embedding_functions
openai_ef = embedding_functions.OpenAIEmbeddingFunction(
api_key="sk-***********************************",
api_base="https://api.siliconflow.cn/v1/",
model_name="BAAI/bge-m3")classDocumentRetriever:def__init__(self, collection_name="documents"):
self.client = chromadb.PersistentClient("./qa_db")
self.collection = self.client.get_or_create_collection(
name=collection_name,
embedding_function=openai_ef
)defindex_documents(self, folder_path:str):"""索引文件夹中的所有文档"""
documents =[]
metadatas =[]
ids =[]for filename in os.listdir(folder_path):if filename.endswith('.txt'):
filepath = os.path.join(folder_path, filename)withopen(filepath,'r', encoding='utf-8')as f:
content = f.read()# 按段落分割
paragraphs = self.split_into_paragraphs(content)for i, para inenumerate(paragraphs):if para.strip():# 跳过空段落
documents.append(para)
metadatas.append({"filename": filename,"paragraph": i,"source":"local_file","filepath": filepath
})
ids.append(f"{filename}_{i}")# 批量添加
batch_size =100
total_added =0for i inrange(0,len(documents), batch_size):
self.collection.add(
documents=documents[i:i+batch_size],
metadatas=metadatas[i:i+batch_size],
ids=ids[i:i+batch_size])
total_added +=len(documents[i:i+batch_size])print(f"已添加 {i+batch_size if i+batch_size <len(documents)elselen(documents)}/{len(documents)} 个段落")return total_added
defsplit_into_paragraphs(self, text:str, min_length:int=50)-> List[str]:"""将文本分割成段落"""
paragraphs = text.split('\n\n')return[p.strip()for p in paragraphs iflen(p.strip())>= min_length]defsearch(self, query:str, n_results:int=5,**filters):"""搜索文档"""return self.collection.query(
query_texts=[query],
n_results=n_results,
where=filters if filters elseNone)defsearch_with_scores(self, query:str, n_results:int=5,**filters):"""搜索文档并显示相似度分数"""
results = self.search(query, n_results,**filters)print(f"查询: '{query}'")print(f"找到 {len(results['documents'][0])} 个结果:")print("-"*50)for i,(doc, metadata, distance)inenumerate(zip(
results['documents'][0],
results['metadatas'][0],
results['distances'][0])):print(f"\n结果 {i+1}:")print(f"相似度: {1- distance:.4f}")# 转换为相似度分数print(f"来源文件: {metadata.get('filename','未知')}")print(f"段落: {metadata.get('paragraph','未知')}")print(f"内容预览: {doc[:150]}...")print("-"*30)return results
defget_collection_info(self):"""获取集合信息"""
count = self.collection.count()return{"total_documents": count,"collection_name": self.collection.name
}defclear_collection(self):"""清空集合中的所有文档"""
self.collection.delete(where={})print("集合已清空")defdelete_by_filename(self, filename:str):"""根据文件名删除文档"""
self.collection.delete(where={"filename": filename})print(f"已删除文件 '{filename}' 的所有段落")# ================ 使用示例 ================defmain():print("=== 文档检索系统演示 ===\n")# 1. 初始化文档检索器
retriever = DocumentRetriever(collection_name="knowledge_base")# 2. 检查集合状态
info = retriever.get_collection_info()print(f"当前集合: {info['collection_name']}")print(f"现有文档数: {info['total_documents']}\n")# 3. 创建示例文档文件夹（如果不存在）
docs_folder ="sample_docs"ifnot os.path.exists(docs_folder):
os.makedirs(docs_folder)# 创建一些示例文档
sample_docs ={"python_tutorial.txt":"""
Python 是一种高级编程语言，由 Guido van Rossum 于 1991 年创建。
它以简洁的语法和强大的功能而闻名，广泛应用于Web开发、数据科学、人工智能等领域。
Python 的主要特性包括：
1. 简单易学：语法清晰，适合初学者
2. 解释型语言：无需编译，可直接运行
3. 动态类型：变量类型在运行时确定
4. 丰富的库生态系统：有大量的第三方库可以使用
在数据科学领域，Python 的 pandas、numpy、matplotlib 等库非常受欢迎。
在人工智能领域，TensorFlow、PyTorch 等框架都支持 Python。
""","machine_learning.txt":"""
机器学习是人工智能的一个分支，它使计算机能够从数据中学习而无需显式编程。
机器学习主要分为三类：监督学习、无监督学习和强化学习。
监督学习：使用标记数据训练模型，用于分类和回归任务。
无监督学习：使用未标记数据发现模式，用于聚类和降维。
强化学习：智能体通过与环境互动学习最优策略。
常见的机器学习算法包括线性回归、决策树、支持向量机、神经网络等。
深度学习是机器学习的一个子领域，使用多层神经网络解决复杂问题。
""","vector_databases.txt":"""
向量数据库是专门用于存储和检索向量嵌入的数据库系统。
它们在高维空间中执行相似性搜索，广泛应用于推荐系统、图像检索和自然语言处理。
Chroma 是一个开源的向量数据库，具有以下特点：
1. 轻量级设计，易于部署
2. 支持多种嵌入模型
3. 提供元数据过滤功能
4. 支持语义搜索
向量数据库的工作原理是将文本、图像等数据转换为向量表示，然后通过计算向量之间的相似度来查找最相关的项目。
余弦相似度是常用的相似度度量方法。
"""}for filename, content in sample_docs.items():withopen(os.path.join(docs_folder, filename),'w', encoding='utf-8')as f:
f.write(content)print(f"已创建示例文档到 '{docs_folder}' 文件夹\n")# 4. 索引文档（可选）if info['total_documents']==0:print("开始索引文档...")
total_added = retriever.index_documents(docs_folder)print(f"成功索引 {total_added} 个段落\n")# 5. 执行搜索查询print("=== 搜索演示 ===\n")# 示例查询 1print("查询 1: 'Python 有什么特点？'")
results1 = retriever.search_with_scores("Python 有什么特点？", n_results=3)# 示例查询 2print("\n\n查询 2: '机器学习有哪些类型？'")
results2 = retriever.search_with_scores("机器学习有哪些类型？", n_results=2)# 示例查询 3：带过滤条件的查询print("\n\n查询 3: '向量数据库'（只搜索特定文件）")
results3 = retriever.search_with_scores("向量数据库",
n_results=3,
filename="vector_databases.txt"# 过滤条件)# 示例查询 4：更具体的查询print("\n\n查询 4: 'Chroma 的特点是什么？'")
results4 = retriever.search_with_scores("Chroma 的特点是什么？", n_results=2)# 6. 高级功能演示print("\n=== 高级功能演示 ===\n")# 查看搜索结果详情print("查看查询 4 的完整结果：")
full_results = retriever.search("Chroma 的特点是什么？", n_results=2)for i, doc inenumerate(full_results['documents'][0]):print(f"\n结果 {i+1} 的完整内容：")print(doc)print("-"*50)# 7. 清理演示（可选）# print("\n=== 清理演示 ===\n")# 删除特定文件的所有段落# retriever.delete_by_filename("python_tutorial.txt")# print("已删除 python_tutorial.txt 的所有段落")# 清空整个集合（谨慎使用）# retriever.clear_collection()# print("集合已清空")if __name__ =="__main__":
main()

复制代码

# 使用批量添加提高性能
batch_size =100for i inrange(0,len(documents), batch_size):
collection.add(
documents=documents[i:i+batch_size],
ids=ids[i:i+batch_size],
metadatas=metadatas[i:i+batch_size])

复制代码

import chromadb
try:
client = chromadb.HttpClient(host="localhost", port=8000)
collection = client.get_collection("my_collection")except chromadb.errors.ChromaError as e:print(f"Chroma 错误: {e}")# 创建新集合或使用回退方案

复制代码

# 启动 Chroma 服务器
chroma run --host localhost --port 8000# Python 客户端连接
client = chromadb.HttpClient(host="localhost", port=8000)

复制代码

# Dockerfile
FROM chromadb/chroma:latest

复制代码

# 运行容器
docker run -p 8000:8000 chromadb/chroma

复制代码

https://blog.csdn.net/u013172930/article/details/147719570
推荐参考这个博客学习

复制代码

from fastapi import FastAPI
import chromadb
app = FastAPI()
client = chromadb.PersistentClient()@app.post("/add_document")asyncdefadd_document(document:dict):
collection = client.get_collection("api_collection")
collection.add(**document)return{"status":"success"}@app.get("/search")asyncdefsearch(query:str, n_results:int=5):
collection = client.get_collection("api_collection")
results = collection.query(query_texts=[query], n_results=n_results)return results

复制代码