Create vectorstore.py

auditqa/engine/vectorstore.py

@@ -0,0 +1,45 @@
+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)