Chroma向量数据库使用

from typing import List
import chromadb
import logging

# chroma client 常用方法封装
class ChromaClient:
    def __init__(self, client = None):
        self.client = client or chromadb.Client() # 支持外部传入client
        self.collection = None

    # 支持自定义配置连接客户端
    def connet(
            self,
            settings=None,
            path: str = ".chroma_db1",
            tenant: str = "default_tenant",
            database: str = "default_database",
        ):
        self.client = chromadb.PersistentClient(
            path=path,
            settings=settings,
            tenant=tenant,
            database=database
        )

    # 创建collection
    def create_collection(self, collection_name):
        if not collection_name or not isinstance(collection_name, str):
            logging.error("Invalid collection name")
            return None
        try:
            self.collection = self.client.create_collection(collection_name)
            logging.info(f"Collection {collection_name} created")
            return self.collection
        except Exception as e:
            logging.error(f"Failed to create collection: {e}")
            return None

    # 切换到collection
    def switch_collection(self, collection_name):
        self.collection = self.client.get_or_create_collection(collection_name)
        return self.collection

    # 删除collection
    def delete_collection(self, collection_name):
        collection = self.get_collection(collection_name)
        if collection:
            try:
                self.client.delete_collection(collection_name)
                logging.info(f"Collection {collection_name} deleted")
                return True
            except Exception as e:
                logging.error(f"Failed to delete collection: {e}")
                return False
        logging.warning(f"Collection {collection_name} not found")
        return None

    # 查询collection
    def get_collection(self, collection_name):
        try:
            return self.client.get_collection(collection_name)
        except Exception as e:
            print(f"Failed to get collection")
            return None

    # 列出所有collection
    def get_list_collections(self):
        list_collections = self.client.list_collections()
        return list_collections

    # 相似度检索
    # n_results  默认是返回10个结果
    def query(self, query_embeddings, n_results=2):
        result = self.collection.query(
            query_embeddings=query_embeddings,
            n_results=n_results,
        )
        return result

    # 插入向量
    def add(
            self,
            documents: List[str],
            metadatas: List[dict],
            ids: List[str],
            embeddings:List[List[float]],
        ) -> bool:
        if not all([documents, metadatas, ids, embeddings]):
            raise ValueError("All arguments must be provided")
        self.collection.add(
            documents=documents,
            metadatas=metadatas,
            ids=ids,
            embeddings=embeddings
        )
        return True

    # 添加或更新向量
    def add_or_update(
            self,
            documents: List[str],
            metadatas: List[dict],
            ids: List[str],
            embeddings:List[List[float]],
        ) -> bool:
        self.collection.upsert(
            documents=documents,
            metadatas=metadatas,
            ids=ids,
            embeddings=embeddings
        )
        return True

from pathlib import Path
from langchain_text_splitters import RecursiveCharacterTextSplitter
from openai import OpenAI
from config.settings import settings
from langchain_community.document_loaders import TextLoader
from embedding.chroma_client import ChromaClient
import uuid

API_KEY = settings.api_key

def init():
    # 1、读取txt文件为 documents list
    cur_path = Path(__file__).parent
    doc_path = (cur_path.parent / "docs/deepseek百度百科.txt").resolve()
    documents = get_text(doc_path)
    # print(documents[0].page_content) // 读取出来是一个list 只有一个item
    # 2、切分文档
    chunks = splitter_document(documents[0].page_content)

    # 将chunks list向量化 转化为存储对象
    store_data = vector_documents(chunks)

    # 链接向量数据库 存储数据
    client = ChromaClient()
    collection = client.switch_collection("deepseek_vector")
    bool = client.add(
        documents=store_data["documents"],
        metadatas=store_data["metadatas"],
        ids=store_data["ids"],
        embeddings=store_data["embeddings"],
    )
    print(f"向量化结果存储成功-----------------------{bool}")

    # 查询collection里面的数据 include=["embeddings"]
    data = collection.get()
    print(f"查询存储结果-----------------------{data}")
    print(data)


# 读取文本
def get_text(file_path):
    loader = TextLoader(file_path, encoding="utf-8")
    documents = loader.load()
    return documents

#  切分文档
def splitter_document(document_text):
    splitter = RecursiveCharacterTextSplitter(
        chunk_size=200,  # 分割长度
        chunk_overlap=40,  # 重叠长度 /重叠窗口大小
        separators=[
            "\n\n",
            "\n",
            "。",
            "！",
            "？",
            "；",
            " ",
            ""],  # 分割符 在分割长度内依次使用段落、换行符、中文句号、中文问号、中文分号、空格、无内容来切分
    )
    chunks = splitter.split_text(document_text)
    return chunks

# 文本向量化 并转为要存储的对象
def vector_documents(chunks):
    llm_client = OpenAI(
        api_key=API_KEY,
        base_url="https://dashscope.aliyuncs.com/compatible-mode/v1",
    )
    chunks_pre_5 = chunks[:5]
    result = llm_client.embeddings.create(input=chunks_pre_5, model="text-embedding-v4")
    embeddings_list = [emb.embedding for emb in result.data]
    store_data = {
        "documents": chunks_pre_5,
        "embeddings": embeddings_list,
        "metadatas": [{"source": "deepseek 百科"} for i in range(len(chunks_pre_5))],
        "ids": [str(uuid.uuid4()) for i in range(len(chunks_pre_5))]
    }
    return store_data

if __name__ == "__main__":
    init()