一、什么是 Chroma?

Chroma 是一个开源的嵌入式向量数据库,专为 AI 应用设计,特别适合存储和检索向量嵌入(embeddings)。它简单易用,支持内存存储和持久化存储。

主要特性:

  • 轻量级:易于安装和部署
  • Python/JavaScript 支持:提供完整的 API
  • 内置嵌入函数:支持多种嵌入模型
  • 元数据过滤:支持高效的元数据查询
  • 持久化存储:支持本地文件和云端存储

二、安装 Chroma

# 基础安装
pip install chromadb

# 包含 CLI 工具
pip install chromadb[cli]

# 完整安装(包含所有依赖)
pip install chromadb[all]

三、核心概念

1. 集合(Collection)

类似于传统数据库的表,用于存储相关文档和向量。

2. 文档(Document)

要存储的文本内容,可以包含:

  • id:唯一标识符
  • text:文本内容
  • metadata:附加信息(标签、来源等)
  • embedding:向量表示(可选)

3. 查询(Query)

基于语义相似度的搜索。

四、基础使用示例

1. 创建客户端和集合

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"
    )
)

2. 添加文档

# 准备数据
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
)

3. 查询文档

# 基本查询
results = collection.query(
    query_texts=["Chroma 有哪些功能?"],
    n_results=3
)

print("查询结果:")
for i, doc in enumerate(results['documents'][0]):
    print(f"{i+1}. {doc}")
    print(f"   元数据: {results['metadatas'][0][i]}")
    print(f"   距离: {results['distances'][0][i]}")
    print()

4. 使用元数据过滤

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

五、高级功能

1. 自定义嵌入函数

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
)

2. 持久化存储

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

# 在指定路径创建集合
collection = persistent_client.create_collection("persistent_collection")

# 数据会自动保存到磁盘

3. 更新和删除文档

# 更新文档
collection.update(
    ids=["doc1"],
    documents=["更新后的文档内容"],
    metadatas=[{"source": "更新", "version": "2.0"}]
)

# 删除文档
collection.delete(ids=["doc2"])

# 按条件删除
collection.delete(where={"source": "教程"})

六、实际应用案例

1. 构建简单的问答系统

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"
)


class SimpleQASystem:
    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
        )
    
    def add_knowledge(self, documents, ids=None, metadatas=None):
        """添加知识库文档"""
        if ids is None:
            ids = [f"doc_{i}" for i in range(len(documents))]
        
        self.collection.add(
            documents=documents,
            ids=ids,
            metadatas=metadatas
        )
    
    def query(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]
        }
    
    def get_context(self, question, max_length=500):
        """获取相关上下文"""
        results = self.query(question)
        context = ""
        
        for answer in results['answers']:
            if len(context) + len(answer) < max_length:
                context += answer + "\n\n"
            else:
                break
        
        return context

# 使用示例
qa_system = SimpleQASystem()
qa_system.add_knowledge([
    "Chroma 支持元数据过滤,可以按条件查询文档",
    "向量数据库主要用于存储和检索高维向量数据",
    "语义搜索基于文本的相似度匹配"
])

context = qa_system.get_context("Chroma 有什么功能?")
print(f"相关上下文:\n{context}")

2. 文档检索系统

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"
)

class DocumentRetriever:
    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
        )
    
    def index_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)
                
                with open(filepath, 'r', encoding='utf-8') as f:
                    content = f.read()
                
                # 按段落分割
                paragraphs = self.split_into_paragraphs(content)
                
                for i, para in enumerate(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 = 0
        for i in range(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) else len(documents)}/{len(documents)} 个段落")
        
        return total_added
    
    def split_into_paragraphs(self, text: str, min_length: int = 50) -> List[str]:
        """将文本分割成段落"""
        paragraphs = text.split('\n\n')
        return [p.strip() for p in paragraphs if len(p.strip()) >= min_length]
    
    def search(self, query: str, n_results: int = 5, **filters):
        """搜索文档"""
        return self.collection.query(
            query_texts=[query],
            n_results=n_results,
            where=filters if filters else None
        )
    
    def search_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) in enumerate(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
    
    def get_collection_info(self):
        """获取集合信息"""
        count = self.collection.count()
        return {
            "total_documents": count,
            "collection_name": self.collection.name
        }
    
    def clear_collection(self):
        """清空集合中的所有文档"""
        self.collection.delete(where={})
        print("集合已清空")
    
    def delete_by_filename(self, filename: str):
        """根据文件名删除文档"""
        self.collection.delete(where={"filename": filename})
        print(f"已删除文件 '{filename}' 的所有段落")


# ================ 使用示例 ================

def main():
    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"
    if not 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():
            with open(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")
    
    # 示例查询 1
    print("查询 1: 'Python 有什么特点?'")
    results1 = retriever.search_with_scores("Python 有什么特点?", n_results=3)
    
    # 示例查询 2
    print("\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 in enumerate(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()

七、最佳实践

1. 数据预处理

  • 清理和标准化文本
  • 合理分割文档(段落或句子级别)
  • 添加有意义的元数据

2. 批量操作

# 使用批量添加提高性能
batch_size = 100
for i in range(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]
    )

3. 错误处理

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}")
    # 创建新集合或使用回退方案

4. 性能优化

  • 使用合适的嵌入维度
  • 定期清理不需要的数据
  • 使用索引优化查询性能

八、部署选项

1. 本地部署(服务器模式)

# 启动 Chroma 服务器
chroma run --host localhost --port 8000

# Python 客户端连接
client = chromadb.HttpClient(host="localhost", port=8000)

2. Docker 部署

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

3. 云部署

Chroma 支持部署到各种云平台,也可以使用 Chroma Cloud 托管服务。

九、与其他工具集成

1. 与 LangChain 集成

https://blog.csdn.net/u013172930/article/details/147719570

推荐参考这个博客学习

2. 与 FastAPI 集成

from fastapi import FastAPI
import chromadb

app = FastAPI()
client = chromadb.PersistentClient()

@app.post("/add_document")
async def add_document(document: dict):
    collection = client.get_collection("api_collection")
    collection.add(**document)
    return {"status": "success"}

@app.get("/search")
async def search(query: str, n_results: int = 5):
    collection = client.get_collection("api_collection")
    results = collection.query(query_texts=[query], n_results=n_results)
    return results

十、常见问题解决

1. 内存不足

  • 使用持久化存储
  • 分批处理大数据集
  • 考虑使用 Chroma 服务器模式

2. 查询速度慢

  • 优化嵌入维度
  • 使用元数据过滤缩小范围
  • 考虑升级硬件配置

十一、资源链接

Logo

有“AI”的1024 = 2048,欢迎大家加入2048 AI社区

更多推荐