database/load_database.py

from langchain_chroma import Chroma
from utils.load_llm import TEXT_SPLITTER, EMBEDDING_MODEL
from config.nodes import PERSIST_DIRECTORY


def load_retriever(collection_name: str, score_threshold=0.6):
    load_database_chroma = Chroma(
        collection_name=collection_name,
        persist_directory=PERSIST_DIRECTORY,
        embedding_function=EMBEDDING_MODEL
    )
    return load_database_chroma.as_retriever(
        search_type="similarity_score_threshold",
        search_kwargs={"score_threshold": score_threshold}
    )