RAG进阶：开源的AI原生向量数据库Chroma

沐雪架构师

已于 2025-04-24 22:52:01 修改

阅读量341

点赞数 4

分类专栏： AI检索增强生成RAG 文章标签：人工智能 llama-index rag embedding

于 2025-04-22 21:19:48 首次发布

本文链接：https://blog.csdn.net/puzi0315/article/details/147431226

版权

AI检索增强生成RAG 专栏收录该内容

10 篇文章

订阅专栏

一、Chroma 核心概念与优势

1. 什么是 Chroma？

Chroma 是一款开源的向量数据库，专为高效存储和检索高维向量数据设计。其核心能力在于语义相似性搜索，支持文本、图像等嵌入向量的快速匹配，广泛应用于大模型上下文增强（RAG）、推荐系统、多模态检索等场景。与传统数据库不同，Chroma 基于向量距离（如余弦相似度、欧氏距离）衡量数据关联性，而非关键词匹配。

GitHub地址：

https://github.com/chroma-core/chroma

官方文档：

https://docs.trychroma.com/

2. 核心优势

轻量易用：以 Python/JS 包形式嵌入代码，无需独立部署，适合快速原型开发。
灵活集成：支持自定义嵌入模型（如 OpenAI、HuggingFace），兼容 LangChain 等框架。
高性能检索：采用 HNSW 算法优化索引，支持百万级向量毫秒级响应。
多模式存储：内存模式用于开发调试，持久化模式支持生产环境数据落地。

二、安装和基础配置

1、安装Chroma

支持windows、ubuntu等操作系统，Python>=3.9 。

创建虚拟环境以及安装

#创建虚拟环境
conda create -n chromadb python==3.10

#激活
conda activate chromadb

#安装chromadb
pip install chromadb

注意：Chroma 默认是本地嵌入式数据库，并不原生支持远程访问像传统数据库那样（比如 PostgreSQL 那种 client-server 模式）。

当然官方也提供了客户端-服务器端模式（Client-Server Mode），服务器端的启动方式如下：

#服务器端启动，默认端口号8000
chroma run --path /db_path

2、初始化客户端

内存模式（调试，实验的场景）：

import chromadb
client = chromadb.Client()

持久化模式（生产环境）：

在创建的时候，可以配置本地的存储路径

client = chromadb.PersistentClient(path="/path/to/save") # 数据保存至本地目录

Client-Server模式的客户端：

前两种，都是本地模式，chroma的服务端和客户端需要位于同一台机器。CS模式可以独立部署，通过httpclient进行访问。

import chromadb

chroma_client = chromadb.HttpClient(host='localhost', port=8000)

三、增删改查操作

1. 创建集合（Collection）

集合是 Chroma 中管理数据的基本单元，类似传统数据库的表。集合的name名称有以下约束：

名称的长度必须介于 3 到 63 个字符之间。
名称必须以小写字母或数字开头和结尾，中间可以包含点、破折号和下划线。
名称不得包含两个连续的点。
该名称不能是有效的 IP 地址。

Chroma 集合是用一个名称和一个可选的嵌入函数创建的。

如果您提供嵌入函数，则每次获取集合时都必须提供它。

# 创建
collection = client.create_collection(name="my_collection", embedding_function=emb_fn)


# 获取
collection = client.get_collection(name="my_collection", embedding_function=emb_fn)


# 若没有则创建，若有则获取
collection = chroma_client.get_or_create_collection(name="my_collection2")

如果不提供嵌入函数，则使用默认的嵌入函数 sentence transformer 使用的的是一个小型的模型all-MiniLM-L6-v2，该模型主要是针对英语场景。一般我们都需要自定义一个嵌入函数：

import chromadb
from sentence_transformers import SentenceTransformer


class SentenceTransformerEmbeddingFunction:
    def __init__(self, model_path: str, device: str = "cuda"):
        self.model = SentenceTransformer(model_path, device=device)

    def __call__(self, input: list[str]) -> list[list[float]]:
        if isinstance(input, str):
            input = [input]
        return self.model.encode(input, convert_to_numpy=True).tolist()


# 创建/加载集合（含自定义嵌入函数）
embed_model = SentenceTransformerEmbeddingFunction(
    model_path=r"D:\Test\LLMTrain\testllm\llm\BAAI\bge-m3",
    device="cuda"  # 无 GPU 改为 "cpu"
)

# 创建客户端和集合
client = chromadb.Client()
collection = client.create_collection("my_knowledge_base", 
                                      metadata={"hnsw:space": "cosine"},
                                      embedding_function=embed_model)

创建collect时，可以配置如下参数。

name标识collect的名称，是必填项；
embedding_function，指定嵌入函数，不填为默认的嵌入模型。
metadata，元数据，比如索引方式等，非必填。

from datetime import datetime

collection = client.create_collection(
    name="my_collection", 
    embedding_function=emb_fn,
    metadata={
        "description": "my first Chroma collection",
        "created": str(datetime.now())
    }  
)

集合有一些常用方法：

peek() - 返回集合中前 10 个项目的列表。
count() -返回集合中的项目数。
modify() -重命名集合

collection.peek() 
collection.count() 
collection.modify(name="new_name")

2、写入数据

写入数据时，配置以下参数:

document，原始的文本块。
metadatas，描述文本块的元数据，kv键值对。
ids，文本块的唯一标识，每个文档必须具有唯一关联的ID。若添加两次相同的 ID 将导致仅存储初始值。
embeddings，对于已经向量化的文本块，可以直接写入结果。如果不填，则在写入时，使用指定或者默认的嵌入函数对documents进行向量化。

collection.add(
    documents=["lorem ipsum...", "doc2", "doc3", ...],
    metadatas=[{"chapter": "3", "verse": "16"}, {"chapter": "3", "verse": "5"}, {"chapter": "29", "verse": "11"}, ...],
    ids=["id1", "id2", "id3", ...]
)

或者

collection.add(
    embeddings=[[1.1, 2.3, 3.2], [4.5, 6.9, 4.4], [1.1, 2.3, 3.2], ...],
    metadatas=[{"chapter": "3", "verse": "16"}, {"chapter": "3", "verse": "5"}, {"chapter": "29", "verse": "11"}, ...],
    ids=["id1", "id2", "id3", ...]
)

3、修改数据

提供ids（文本唯一标识）。

collection.update(
    ids=["doc1"],  # 使用已存在的ID
    documents=["RAG是一种检索增强生成技术222"]
)

4、更新插入方法

Chroma 还支持更新插入操作，更新现有项目，如果项目尚不存在则添加它们。

collection.upsert(
    ids=["id1", "id2", "id3", ...],
    embeddings=[[1.1, 2.3, 3.2], [4.5, 6.9, 4.4], [1.1, 2.3, 3.2], ...],
    metadatas=[{"chapter": "3", "verse": "16"}, {"chapter": "3", "verse": "5"}, {"chapter": "29", "verse": "11"}, ...],
    documents=["doc1", "doc2", "doc3", ...],
)

5、删除数据

Chroma 支持通过以下方式从集合中删除项目ID使用delete。与每个项目相关的嵌入、文档和元数据将被删除。

还支持where过滤器。如果没有ID提供，它将删除集合中与where筛选。

# 提供ids
collection.delete(ids=["doc1"])

collection.delete(
    ids=["id1", "id2", "id3",...],
	where={"chapter": "20"}
)

6、查询数据

（1）查询所有数据

all_docs = collection.get()
print("集合中所有文档：", all_docs)

（2）根据ids查询

可以通过以下方式从集合中检索项目ID使用get。

collection.get(
	ids=["id1", "id2", "id3", ...],
	where={"style": "style1"}
)

（3）查询嵌入

可以通过多种方式查询 Chroma 集合，使用query方法。比如使用query_embedding。

collection.query(
    query_embeddings=[[11.1, 12.1, 13.1],[1.1, 2.3, 3.2], ...],
    n_results=10,
    where={"metadata_field": "is_equal_to_this"},
    where_document={"$contains":"search_string"}
)

查询将返回n_result每个最接近的匹配查询嵌入，按顺序排列。可选where过滤字典可以通过metadata与每个文档关联。此外，where document可以提供过滤字典来根据文档内容进行过滤。

（4）查询相似文档

还可以通过一组查询文本query_texts. Chroma 将首先嵌入每个查询文本与集合的嵌入函数，然后使用生成的嵌入执行查询。

# 查询相似文档
results = collection.query(
    query_texts=["什么是RAG技术？"],
    n_results=3
)

print("查询的结果",results)

查询结果

当使用 get 或 query 时，您可以使用include参数来指定您想要返回的数据包括：embeddings， documents，metadatas；include为数组，可以传多个值。
对于查询query，默认返回距离distances结果。
embeddings出于性能考虑，默认不返回，直接显示None ，若想返回，则include中包含embeddings即可。
ID始终会返回。
返回值里有included参数，表明本次返回的数据有哪些类型。
embeddings将以二维 NumPy 数组的形式返回。

# Only get documents and ids
collection.get(
    include=["documents"]
)

collection.query(
    query_embeddings=[[11.1, 12.1, 13.1],[1.1, 2.3, 3.2], ...],
    include=["documents"]
)

#查询的结果

{'ids': [['doc1', 'doc3', 'doc2']], 
'embeddings': None, 
'documents': [['RAG是一种检索增强生成技术', '三英战吕布', '向量数据库存储文档的嵌入表示']], 
'uris': None,
'included': ['metadatas', 'documents', 'distances'], 
'data': None, 
'metadatas': [[{'source': 'tech_doc'}, {'source': 'tutorial1'}, {'source': 'tutorial'}]], 
'distances': [[0.2373753786087036, 0.7460092902183533, 0.7651787400245667]]
}

四、实战操作

将一批数据插入向量数据库，再根据一个问题从向量数据库中找出相似数据。

1、安装包

pip install sentence_transformers

pip install modelscope

2、下载Embedding模型到本地

#模型下载
from modelscope import snapshot_download
model_dir = snapshot_download('BAAI/bge-m3',cache_dir=r"D:\Test\LLMTrain\testllm\llm")

3、写入数据和查询相似文本

import chromadb
from sentence_transformers import SentenceTransformer


class SentenceTransformerEmbeddingFunction:
    def __init__(self, model_path: str, device: str = "cuda"):
        self.model = SentenceTransformer(model_path, device=device)

    def __call__(self, input: list[str]) -> list[list[float]]:
        if isinstance(input, str):
            input = [input]
        return self.model.encode(input, convert_to_numpy=True).tolist()


# 创建/加载集合（含自定义嵌入函数）
embed_model = SentenceTransformerEmbeddingFunction(
    model_path=r"D:\Test\LLMTrain\testllm\llm\BAAI\bge-m3",
    device="cuda"  # 无 GPU 改为 "cpu"
)

# 创建客户端和集合
client = chromadb.PersistentClient(path=r"D:\Test\LLMTrain\chromadb_test\chroma_data")
collection = client.get_or_create_collection("my_knowledge_base",
                                      metadata={"hnsw:space": "cosine"},
                                      embedding_function=embed_model)

# 添加文档
collection.add(
    documents=[ "向量数据库存储文档的嵌入表示", "三英战吕布","RAG是一种检索增强生成技术"],
    metadatas=[{"source": "tech_doc"}, {"source": "tutorial"}, {"source": "tutorial1"}],
    ids=["doc1", "doc2", "doc3"]
)

# 查询相似文档
results = collection.query(
    query_texts=["什么是RAG技术？"],
    n_results=3
)

print("查询的结果",results)

执行返回结果：

查询的结果 
{
'ids': [['doc3', 'doc2', 'doc1']], 
'embeddings': None, 
'documents': [['RAG是一种检索增强生成技术', '三英战吕布', '向量数据库存储文档的嵌入表示']], 
'uris': None, 
'included': ['metadatas', 'documents', 'distances'], 
'data': None, 
'metadatas': [[{'source': 'tutorial1'}, {'source': 'tutorial'}, {'source': 'tech_doc'}]], 
'distances': [[0.2373753786087036, 0.7460092902183533, 0.7651787400245667]]
}

查看结果，我们重点看 distances，值是从小到大排序的，所以3条数据与问题“什么是RAG技术？”的相似度情况：distances值越小越相似。因此，第1条数据与问题越相似。