auditqa/engine/vectorstore.py

from langchain.embeddings import OpenAIEmbeddings, HuggingFaceEmbeddings, HuggingFaceInferenceAPIEmbeddings
from dotenv import load_dotenv

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"

if provider_retrieval_model == "HF":
  qdrantClient = QdrantClient(path=client_path, prefer_grpc=True)

  embeddings = HuggingFaceInferenceAPIEmbeddings(
      api_key=HF_Token, model_name=embeddingmodel
  )

  dim = 1024

elif provider_retrieval_model == "OAI":

  qdrantClient = QdrantClient(path=client_path, prefer_grpc=True)

  embeddings = OpenAIEmbeddings(
          model="text-embedding-ada-002",
          openai_api_key=os.getenv("OPENAI_API_KEY"),
  )

  dim = 1536


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_samp)