Update: Updated the qdrant pipeline

Author: Saqlain-ce

server.py

@@ -6,38 +6,37 @@
 from typing import Dict, Any, List, Optional
 
 from src.config.config import Config
+from src.qdrant.qdrant_utils import QdrantWrapper
 from src.embedder.embedder_llama_index import EmbeddingWrapper
-from llama_index.core.retrievers import VectorIndexRetriever
+from src.parser.csv_parser import CsvParser
 from llama_index.core import Settings
 Settings.llm = None
 
-from src.qdrant.qdrant_manager import QdrantManager
 from src.utils.connections_manager import ConnectionManager
 from src.chatbot.rag_chat_bot import RAGChatBot
 from src.reranker.re_ranking import RerankDocuments
 
-import os
-
 app = FastAPI()
 
 chatbot = RAGChatBot()
+file_processor = CsvParser(data_dir = Config.DATA_DIRECTORY)
 
 collection_name = Config.COLLECTION_NAME
-qdrantManager = QdrantManager(Config.QDRANT_HOST, Config.QDRANT_PORT, collection_name)
-
+qdrant_client = QdrantWrapper()
 embedding_client = EmbeddingWrapper()
 
 
-data_dir = Config.CAPEC_DATA_DIR
+try:
 
-reranker = RerankDocuments()
+    processed_chunks = file_processor.process_directory()
+    qdrant_client.ingest_embeddings(processed_chunks)
 
-index = qdrantManager.load_index(persist_dir=Config.PERSIST_DIR, embed_model=embedding_client)
+    logger.info("Successfully ingested Data")
 
-retriever = VectorIndexRetriever(
-            index=index,
-            similarity_top_k=5
-        )
+except Exception as e:
+    logger.error(f"Error in data ingestion: {str(e)}")
+
+reranker = RerankDocuments()
 
 # Manually added file names of the CAPEC daatset. In production, These files will be fetched from database
 database_files = ["333.csv", "658.csv", "659.csv", "1000.csv", "3000.csv"]
@@ -66,27 +65,26 @@ async def handle_search(websocket: WebSocket, query: str) -> None:
 
         filename = find_file_names(query, database_files)
 
-        if filename:
-            logger.info("Searching for file names...")
+        query_embeddings = embedding_client.generate_embeddings(query)
 
-            filters = MetadataFilters(filters=[ExactMatchFilter(key="source_file", value=filename)])
-            relevant_nodes =  index.as_retriever(filters=filters).retrieve(query)
-            if not relevant_nodes:
-                logger.info("Searching without file name filter....")
-                relevant_nodes = retriever.retrieve(query)
-        else:
-            logger.info("Searching without file names....")
-            relevant_nodes = retriever.retrieve(query)
+        top_5_results = qdrant_client.search(query_embeddings, 5)
+        logger.info("Retrieved top 5 results")
 
-
-        context = [node.text for node in relevant_nodes]
-    
-        reranked_docs =  reranker.rerank_docs(query, context)
+        if not top_5_results:
+            logger.warning("No results found in database")
+            await websocket.send_json({
+                "result": "The database is empty. Please ingest some data first before searching."
+            })
+            return
 
-        # only top 2 documents are passing as a context
-        response, conversation_id  = chatbot.chat(query, reranked_docs[:2])
 
+        reranked_docs = reranker.rerank_docs(query, top_5_results)
+        reranked_top_5_list = [item['content'] for item in reranked_docs]
 
+        context = reranked_top_5_list[:2]
+
+        # only top 2 documents are passing as a context
+        response, conversation_id  = chatbot.chat(query, context)
 
         logger.info("Generating response from Groq")
 

src/config/config.py

@@ -13,7 +13,7 @@ class Config:
     GRADIO_SERVER_NAME = "0.0.0.0" 
     GRADIO_SERVER_PORT = int(7860)
     WEBSOCKET_URI = "ws://rag-server:8000/ws"
-    DATA_DIRECTORY = "data/"
+    DATA_DIRECTORY = "capec-dataset/"
     WEBSOCKET_TIMEOUT = 300  # 5 minutes
     HEARTBEAT_INTERVAL = 30  # 30 seconds
     MAX_CONNECTIONS = 100

src/parser/csv_parser.py

@@ -139,5 +139,5 @@ def process_directory(self) -> List[Document]:
                 logger.error(f"Skipping file {file_path} due to error: {str(e)}")
                 continue
 
-        logger.info("All .csv files indexed....")
+        logger.info("All .csv files processed. Returning chunks...")
         return all_documents

src/qdrant/qdrant_utils.py

@@ -52,6 +52,7 @@ def _connect_with_retry(self) -> None:
                 self.client.get_collections()
                 logger.info("Successfully connected to Qdrant")
                 self._create_collection_if_not_exists()
+                self.clear_collection()
                 break
             except Exception as e:
                 logger.error(f"Connection attempt {attempt + 1} failed: {str(e)}")

src/reranker/re_ranking.py

@@ -27,10 +27,11 @@ def rerank_docs(self,
         """
         # Re-ranking using cross-encoder
         # Prepare pairs for reranking
-        pairs = [[query, doc] for doc in top_5_results]
+         # Prepare pairs for reranking
+        pairs = [[query, doc["content"]] for doc in top_5_results]
 
         # Get relevance scores
-        scores = self.reranker.predict(pairs) 
+        scores = self.reranker.predict(pairs)
 
         # Sort by new scores
         reranked_results = [

src/utils/utils.py

@@ -2,6 +2,7 @@
 import re
 from loguru import logger
 
+
 def match_file_names(filename, database_files):
     if filename in database_files:
         return filename
@@ -25,4 +26,6 @@ def find_file_names(query: str, database_files: List) -> str:
         else:
             return ""
     else:
-        logger.info("No filename found.")
+        logger.info("No filename found.")
+
+