Merge pull request #8 from cyber-evangelists/dev-branch

Removed index with qdrant

Author: Saqlain-ce

client.py

@@ -98,7 +98,7 @@ async def record_feedback(feedback, msg ) -> gr.Info:
 
 
 with gr.Blocks(
-    title="CAPEC RAG Chatbot",
+    title="EASM RAG Chatbot",
     theme=gr.themes.Soft(),
     css="""
         .gradio-container {
@@ -171,7 +171,7 @@ async def record_feedback(feedback, msg ) -> gr.Info:
 
     # Header
     gr.Markdown(
-        "<div id='header'>CAPEC RAG Application</div>"
+        "<div id='header'>EASM RAG Application</div>"
     )
 
     # Chatbot Component

requirements.txt

@@ -7,10 +7,7 @@ langchain-core==0.3.18
 gradio
 loguru==0.7.2
 python-dotenv==1.0.1
-llama-index==0.11.21
-llama-index-vector-stores-qdrant==0.3.3
 langchain-groq==0.2.1
-llama-index-embeddings-huggingface==0.3.1
 langchain==0.3.7
 transformers==4.46.2
 torch==2.5.1

server.py

@@ -1,16 +1,13 @@
 from fastapi import FastAPI, WebSocket, WebSocketDisconnect
 from loguru import logger
 from src.utils.utils import find_file_names
-from llama_index.core.vector_stores.types import MetadataFilters, ExactMatchFilter
 
 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 src.embedder.embedder import EmbeddingWrapper
 from src.parser.csv_parser import CsvParser
-from llama_index.core import Settings
-Settings.llm = None
 
 from src.utils.connections_manager import ConnectionManager
 from src.chatbot.rag_chat_bot import RAGChatBot
@@ -28,6 +25,10 @@
 
 try:
 
+    qdrant_client.delete_collection(collection_name=collection_name)
+    logger.info("collection deleted...")
+    qdrant_client._create_collection_if_not_exists()
+    logger.info("Collection created....")
     processed_chunks = file_processor.process_directory()
     qdrant_client.ingest_embeddings(processed_chunks)
 
@@ -63,10 +64,12 @@ async def handle_search(websocket: WebSocket, query: str) -> None:
     try:
         logger.info(f"Processing search query")
 
-        filename = find_file_names(query, database_files)
+        # filename = find_file_names(query, database_files)
 
         query_embeddings = embedding_client.generate_embeddings(query)
 
+
+        logger.info("Searching for top 5 results....")
         top_5_results = qdrant_client.search(query_embeddings, 5)
         logger.info("Retrieved top 5 results")
 

src/chatbot/rag_chat_bot.py

@@ -35,6 +35,7 @@ def __init__(self):
             model_name="llama-3.1-8b-instant",
             temperature=0,
             max_tokens=4096,
+            frequency_penalty=0.9
         )
 
         # Initialize memory

src/docker-files/Dockerfile.client

@@ -9,7 +9,7 @@ COPY client-requirements.txt .
 
 # Update pip and install dependencies
 RUN pip install --upgrade pip && \
-    pip install -r client-requirements.txt
+    pip install --default-timeout=5000  -r client-requirements.txt
 
 # Copy only the required files for the application
 COPY client.py .

src/docker-files/Dockerfile.server

@@ -7,8 +7,9 @@ COPY requirements.txt .
 
 # Update pip and install dependencies
 RUN pip install --upgrade pip && \
-    pip install -r requirements.txt
+    pip install --default-timeout=5000  -r requirements.txt
 
+
 COPY server.py .
 
 # Set Python to run in unbuffered mode

src/embedder/embedder.py

@@ -0,0 +1,21 @@
+import numpy as np
+from sentence_transformers import SentenceTransformer
+from src.config.config import Config
+
+
+class EmbeddingWrapper:
+    def __init__(self, model_name='all-MiniLM-L6-v2'):
+        self.model = SentenceTransformer(Config.EMBEDDING_MODEL_PATH)
+    
+    def generate_embeddings(self, texts):
+        """
+        Generate embeddings for a list of texts.
+        
+        Args:
+            texts (list): A list of strings to generate embeddings for.
+        
+        Returns:
+            numpy.ndarray: A 2D array of embeddings, where each row corresponds to a text input.
+        """
+        embeddings = self.model.encode(texts)
+        return np.array(embeddings)

src/embedder/embedder_llama_index.py

@@ -1,37 +1,44 @@
 
 import pandas as pd
-from typing import List, Dict, Any, Optional
+from typing import List, Dict, Any, Optional, TypedDict
 from pathlib import Path
 import numpy as np
-from llama_index.core.schema import Document
-from llama_index.core.node_parser import SentenceSplitter
+from src.embedder.embedder import EmbeddingWrapper
 
 from datetime import datetime
 from dataclasses import dataclass
-
 from loguru import logger
-
 from src.config.config import Config
 
+
 @dataclass
-class DocumentMetadata:
+class DocumentMetadata(TypedDict):
     """Class to store document metadata"""
     source_file: str
     ingestion_timestamp: str
     last_updated_timestamp: str
     embedding_version: str
-    embedding_model: str
+    embedding_model_name: str
     processing_status: str
 
 
+
+class ProcessedChunk(TypedDict):
+    """Type definition for processed file chunks."""
+    embeddings: List[float]
+    text: str
+    metadata: str
+
+
+
 class CsvParser:
 
-    def __init__(self, data_dir: str, embedding_version: str =  Config.EMBEDDING_VERSION_NUMBER, embedding_model: str = Config.EMBEDDING_MODEL) -> None:
+    def __init__(self, data_dir: str, embedding_version: str =  Config.EMBEDDING_VERSION_NUMBER, embedding_model_name: str = Config.EMBEDDING_MODEL) -> None:
         self.data_dir = Path(data_dir)
         self.embedding_version = embedding_version
-        self.embedding_model = embedding_model
-        self.node_parser = SentenceSplitter(chunk_size=1200, chunk_overlap=200)
-        
+        self.embedding_model_name = embedding_model_name
+        self.embedder = EmbeddingWrapper()
+        self.chunks: List[ProcessedChunk] = []
 
     def create_document_metadata(self, row: pd.Series, file_name: str,) -> DocumentMetadata:
         """Create comprehensive document metadata"""
@@ -42,7 +49,7 @@ def create_document_metadata(self, row: pd.Series, file_name: str,) -> DocumentM
             ingestion_timestamp=current_time,
             last_updated_timestamp=current_time,
             embedding_version=self.embedding_version,
-            embedding_model=self.embedding_model,  # In practice, this might be different
+            embedding_model_name=self.embedding_model_name,  # In practice, this might be different
             processing_status="processed",
         )
 
@@ -68,35 +75,35 @@ def read_file(self, file_path: Path) -> pd.DataFrame:
         return df
 
 
-    def process_file(self, file_path: Path) -> List[Document]:
+    def process_file(self, file_path: Path) -> None:
         """Process a single CSV file with enhanced metadata and version control"""
         try:
             logger.info(f"Processing file: {file_path}")
 
             # Read CSV file
             df = self.read_file(file_path)
 
-            documents = []
             for _, row in df.iterrows():
                 # Combine text fields
                 text_content = self.get_text(row)
 
                 # Create comprehensive metadata
                 metadata = self.create_document_metadata(row, file_path.name)
-                
+                embeddings = self.embedder.generate_embeddings(text_content)
+
+
                 # Create Document object with enhanced metadata
-                doc = Document(
-                    text=text_content,
-                    metadata=metadata.__dict__
-                )
-
-                nodes = self.node_parser.get_nodes_from_documents([doc])
-                documents.extend(
-                    [Document(text=node.text, metadata=node.metadata) for node in nodes]
-                )
+                doc : ProcessedChunk = {
+                    "embeddings": embeddings,
+                    "text":text_content,
+                    "metadata":"metadata"
+                }
+
+
+                self.chunks.append(doc)
 
-            logger.info(f"Successfully processed {len(documents)} documents from {file_path}")
-            return documents
+            logger.info(f"Successfully processed all documents from {file_path}")
+
 
         except Exception as e:
             logger.error(f"Error processing file {file_path}: {str(e)}")
@@ -126,18 +133,18 @@ def get_text(self, row: pd.Series) -> str:
         return " | ".join(text_parts)
 
 
-    def process_directory(self) -> List[Document]:
+    def process_directory(self) -> List[Dict[str, Any]]:
         """Process all CSV files in directory"""
         all_documents = []
 
         logger.info("Attempting to read all .csv files and indexing....")
         for file_path in self.data_dir.glob('*.csv'):
             try:
-                documents = self.process_file(file_path)
-                all_documents.extend(documents)
+                self.process_file(file_path)
             except Exception as e:
                 logger.error(f"Skipping file {file_path} due to error: {str(e)}")
                 continue
+            
 
         logger.info("All .csv files processed. Returning chunks...")
-        return all_documents
+        return self.chunks