rag_project/qa_chain/get_vectordb.py

import os
from database.create_db import create_db, load_knowledge_db
from embedding.embedding import get_embedding

# 定义默认路径
# DEFAULT_DB_PATH = os.path.join("..", "knowledge_db")
# DEFAULT_PERSIST_PATH = os.path.join("..", "vector_db", "chroma")

def get_vectordb(file_path, persist_path, embedding="m3e"):
    """
    返回向量数据库对象
    输入参数：
    question：
    llm:
    vectordb:向量数据库(必要参数),一个对象
    embedding：qwen
    """
    embedding = get_embedding(embedding=embedding)
    if os.path.exists(persist_path):  # 持久化目录存在
        contents = os.listdir(persist_path)
        if len(contents) == 0:  # 但是下面为空
            print("目录为空")
            create_db(file_path, persist_path, embedding)
            # presit_knowledge_db(vectordb)
            vectordb = load_knowledge_db(persist_path, embedding)
        else:
            # print("目录不为空")
            vectordb = load_knowledge_db(persist_path, embedding)
    else:  # 目录不存在，从头开始创建向量数据库
        create_db(file_path, persist_path, embedding)
        # presit_knowledge_db(vectordb)
        vectordb = load_knowledge_db(persist_path, embedding)

    return vectordb