from langchain_chroma import Chroma
from langchain_huggingface.embeddings import HuggingFaceEmbeddings


def init_embedding(
    model_name: str,
    # cache_folder: str = "src/models/embeddings/tmp",
    model_kwargs: dict,
    test_query: str,
):
    """Initializing the embedding model and return it"""
    # Options: BAAI/bge-m3 (75%) - intfloat/multilingual-e5-large (80%) - sentence-transformers/all-mpnet-base-v2 (50%) - intfloat/e5-mistral-7b-instruct (RAM-Limit) - jinaai/xlm-roberta-flash-implementation (Base) - jinaai/jina-embeddings-v3 (68%)
    embedding_model = HuggingFaceEmbeddings(
        model_name=model_name,  # "intfloat/multilingual-e5-large"
        # cache_folder=cache_folder,
        model_kwargs=model_kwargs,
    )  # {"device":"cpu", "trust_remote_code": True}
    test_result = embedding_model.embed_query(test_query)

    return embedding_model, test_result


def get_vector_store(persist_directory: str, embedding_model, k_arg: int):
    """Get the existed Chroma vector store and return it as retriever"""
    vector_store = Chroma(persist_directory=persist_directory, embedding_function=embedding_model)
    retriever = vector_store.as_retriever(search_kwargs={"k": k_arg})

    return retriever
