Update auditqa/engine/vectorstore.py

auditqa/engine/vectorstore.py

@@ -1,4 +1,6 @@
 from langchain.embeddings import OpenAIEmbeddings, HuggingFaceEmbeddings, HuggingFaceInferenceAPIEmbeddings
+from langchain.vectorstores import Chroma, Qdrant
+from qdrant_client import QdrantClient
 from dotenv import load_dotenv
 import os
 
@@ -6,6 +8,32 @@ 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
+    
 
-def token():
-    print(HF_Token)
+    
+    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)