auditqa/engine/vectorstore.py

from langchain.embeddings import OpenAIEmbeddings, HuggingFaceEmbeddings, HuggingFaceInferenceAPIEmbeddings
from langchain.vectorstores import Chroma, Qdrant
from qdrant_client import QdrantClient
from qdrant_client.models import Distance, VectorParams
from dotenv import load_dotenv
import os

provider_retrieval_model = "HF"
embeddingmodel = "BAAI/bge-small-en-v1.5"
load_dotenv()
HF_Token = os.environ.get("HF_TOKEN")
client_path = f"./vectorstore"
collection_name = f"collection"
provider_retrieval_model = "HF"

def create_vectorstore(docs):
    
    if provider_retrieval_model == "HF":
      qdrantClient = QdrantClient(path=client_path, prefer_grpc=True)
    
      embeddings = HuggingFaceInferenceAPIEmbeddings(
          api_key=HF_Token, model_name=embeddingmodel
      )
    
      dim = 1024
    

    
    qdrantClient.create_collection(
        collection_name=collection_name,
        vectors_config=VectorParams(size=dim, distance=Distance.COSINE),
    )
    
    vectorstore = Qdrant(
        client=qdrantClient,
        collection_name=collection_name,
        embeddings=embeddings,
    )
    
    vectorstore.add_documents(docs)