LLM知识存储ChromaDB向量数据库应用

1、通过Docker快速部署ChromaDB

docker run -d --name chromadb-container -p 8000:8000 chromadb/chroma

2、打开ChromaDB提供的swagger地址 http://localhost:8000/docs

3、构建代码测试向量

以下代码做了基础测试，创建集合、文字转向量、插入集合、检索数据具体的可以查看Swagger文档提供了更多的操作，运行当前代码可以根据语意搜索到对应的内容，文字转向量需要部署应用Xinference

import requests
import hashlib


# 文字转向量
def embedding(data):
    url = "http://10.10.40.102:19997/v1/embeddings"
    headers = {
        "Content-Type": "application/json"
    }
    data = {
        "input": data,
        "model": "bge-large-zh-v1.5"
    }
    response = requests.post(url, headers=headers, json=data)

    vectorData = [item['embedding'] for item in response.json()['data']]
    return vectorData


# 创建集合
def create_collections(json):
    result = requests.post("http://localhost:8000/api/v1/collections?tenant=default_tenant&database=default_database",
                           json=json)
    return result


# 获取集合列表
def get_collections():
    result = requests.get("http://localhost:8000/api/v1/collections?tenant=default_tenant&database=default_database")
    return result


# 插入向量
def add_collections_vectory(collections_id, json):
    result = requests.post(f"http://localhost:8000/api/v1/collections/{collections_id}/add", json=json)
    print(result.json())


# 根据名称获取集合
def name_get_collections(collections_name):
    result = requests.get(
        f"http://localhost:8000/api/v1/collections/{collections_name}?tenant=default_tenant&database=default_database")
    return result


def delete_collections(collections_name):
    result = requests.delete(
        f"http://localhost:8000/api/v1/collections/{collections_name}?tenant=default_tenant&database=default_database")
    print(result.json())


# 数组转md5
def convert_to_md5(data):
    return [hashlib.md5(item.encode('utf-8')).hexdigest() for item in data]


# 添加向量数据
def conversion_data(raw_data):
    # 转唯一ID
    ids = convert_to_md5(raw_data)

    # 文字转向量
    embeddings = embedding(raw_data)

    # 转向量
    return {"embeddings": embeddings, "documents": raw_data, "ids": ids}


def retrieval(collections_name, name):
    collections_id = name_get_collections(collections_name).json().get("id")
    fruit = embedding(name)
    result = requests.post(f"http://localhost:8000/api/v1/collections/{collections_id}/query", json={
        "query_embeddings": fruit,
        "n_results": 5,  # 返回结果数量
        "include": [
            "metadatas",
            "distances",
            "documents"
        ]
    })
    print(result.json().get("distances"))
    print(result.json().get("documents"))


def generate(collections_name, raw_data):
    # 删除之前的集合
    delete_collections(collections_name)

    # 创建新集合
    create_collections({"name": collections_name})

    # 获取集合ID
    collections_id = name_get_collections(collections_name).json().get("id")

    # 把向量插入到chromadb
    add_collections_vectory(collections_id, conversion_data(raw_data))


# 生成数据插入到集合
generate("test", ["苹果", "香蕉", "汽车", "狗", "猫", "大象", "飞机", "坦克"])

# 检索数据
retrieval("test", "交通工具")