在python中使用chromadb

1.简介

Chroma 是 AI 原生开源矢量数据库。Chroma 通过为 LLM 提供知识、事实和技能，使构建 LLM 应用程序变得容易。同时也是实现大模型RAG技术方案的一种有效工具。

• Chrome提供以下能力：
  1. 存储嵌入类型数据（embeddings）和其元数据
  2. 嵌入（embed）文档和查询
  3. 对嵌入类型的检索
• Chrome 的原则：
  1. 对用户的简单性，并保障开发效率
  2. 同时拥有较好的性能

在Python环境中，Chroma可以在Python脚本中作为客户端、服务端或者直接作为嵌入式向量数据库（数据可以在内存中或磁盘）使用。
在其他环境中目前只能作为客户端使用。
从源码中可以看到Chroma作为服务端的时候是fastapi写的接口。

2.与其他向量数据库对比

chroma与商业向量数据库比起来，无论是性能还是存储规模都是比较低，所以一般用来作为产品初期功能验证时候使用，如果你的数据量不大也不需要经常更新的情况，产品中也是可以使用的。这里的数据规模一般指存储的数量条数在100k以下，按每条500~800字算的话，可以承载100k*650=65M字，用来作为企业内部知识库一般也是足够的，如果用来存放用户自己上传的数据，那就不行了，容纳不了太多用户。

3.安装chroma

1. 在python环境中使用pip进行安装即可

pip install chromadb

2. 在docker中安装

docker pull chromadb/chroma 
docker run -p 8000:8000 chromadb/chroma 
或者从GitHub下下载源码，通过源码中的dockerfile安装 
git clone git@github.com:chroma-core/chroma.git 
cd chroma 
docker-compose up -d --build

4.启动

1.作为服务端启动

在python环境中执行下面的命令,--path后面跟的参数是路径参数，是指数据持久化存放的位置，--host是主机地址，一般写127.0.0.1或0.0.0.0，--port是服务端口

chroma run --path xxx --host xxx --port xxx

2. 作为嵌入数据库

import chromadb 
client = chromadb.PersistentClient(path="/path/to/save/to")

3.作为客户端启动

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

4.作为异步客户端

import chromadb
chroma_client = await chromadb.AsyncHttpClient(host='localhost', port=8000)


如果需要让chroma作为服务端并跟随python程序启动/关闭，可以使用下面的代码

class StartChromaServer:

    def __init__(self):
        self.__command = ["chroma", "run", "--port", str(CHROMA_DATABASE_PORT), "--path", CHROMA_DATABASE_PATH]
        self.__process = subprocess.Popen(self.__command, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
        print(f'chroma server start.PID:{self.__process.pid}')
        atexit.register(self.close_external_program)

    def close_external_program(self):
        if self.__process.poll() is None:
            self.__process.terminate()
            print('chroma server close.')


StartChromaServer()

5.collection

chroma使用collection作为管理数据的基本单元，可以创建、查询、删除collection
创建collection

import chromadb
chroma_client = chromadb.Client()
collection = chroma_client.create_collection(name="my_collection")
#或者使用另外一种更安全的方式
collection = chroma_client.get_or_create_collection(name="my_collection")

查询collection

import chromadb
chroma_client = chromadb.Client()
collection = chroma_client.get_collection(name="my_collection")

删除collection

import chromadb
chroma_client = chromadb.Client()
chroma_client.delete_collection(name="my_collection")

6.数据增删改查

在拿到collection后，就可以在该collection下对数据增删改查了。
插入

collection.add(
embeddings=[[1.2, 2.3, 4.5], [6.7, 8.2, 9.2]],
documents=["This is a document", "This is another document"],
metadatas=[{"source": "my_source"}, {"source": "my_source"}],
ids=["id1", "id2"]
)

查询

results = collection.query(
    query_texts=["This is a query document"],
    n_results=2
)

更新

collection.update(
    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", ...],
)

若不存在则新增

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", ...],
)

删除

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

7.设置embedding function

如果我们在代码中实现了文本转向量的工作，我们在使用chromadb的时候可以使用参数embeddings=[[],[],...]即可，chromadb库同时也提供了它内部实现的方式，检测到我们传入的参数只有文本的情况下它会自行调用我们设置的embedding function。

自定义嵌入方法

class OpenaiEmbeddingFunction(EmbeddingFunction):
    def __init__(self):
        key = OPENAI_API_KEY
        url = "http://xxxxxx.com/v1"
        self.key = key
        self.url = url

    def __call__(self, input: Documents) -> Embeddings:
        embeddings = OpenAIEmbeddings(
            model="text-embedding-3-small",
            openai_api_base=self.url,
            openai_api_key=self.key,
        )
        return embeddings.embed_documents(input)

self.__EMBDEDDING_FUNC = OpenaiEmbeddingFunction() 
#在获取collection的时候把自定义的方法作为参数传入
collection = await chroma_client.get_or_create_collection(
            name=user_id,
            embedding_function=self.__EMBDEDDING_FUNC,
            metadata=self.__DEFAULT_METADATA,
        )

更多内容可以查阅chroma的官网、非官网文档
https://docs.trychroma.com/
https://cookbook.chromadb.dev/