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

husnain-ce · web-flow · commit 2bf773880196 · 2024-11-21T13:42:15.000+05:00
PR for Language Issue Fix and Input Validations
diff --git a/client-requirements.txt b/client-requirements.txt
@@ -3,4 +3,5 @@ loguru==0.7.2
 websockets
 python-dotenv==1.0.1
 transformers==4.46.2
-torch==2.5.1
+torch==2.5.1
+python-bidi
diff --git a/client.py b/client.py
@@ -1,8 +1,5 @@
 import gradio as gr
 import websockets
-import json
-import asyncio
-import logging
 from typing import Tuple, List, Optional, Dict, Any
 from loguru import logger
 
@@ -15,24 +12,41 @@
 guardrails_model = GuardRails()
 
 
-async def search_click(msg, history):
+async def search_click(msg: str, history: List[Tuple[str, str]]) -> Tuple[str, List[Tuple[str, str]], gr.Info]:
+
+    if not msg.strip():
+        logger.error(f"No input provided")
+        return "", history,  gr.Warning("Please enter a query.")
 
     response = int(guardrails_model.classify_prompt(msg))
 
     if response == 0:
-        return await ws_client.handle_request(
+        result =  await ws_client.handle_request(
             "search",
             {"query": msg, "history": history if history else []}
         )
+        if result[2] == "right":
+
+            styled_response = (f"<div style='direction: rtl; text-align: right; direction: right;'>{result[1]}</div>")
+        else:
+            styled_response = f"<div style='direction: ltr; text-align: left; direction: left;'>{result[1]}</div>"
+        
+        # Append the styled response to the chat history
+        updated_history = history + [(msg, styled_response)]
+
+
+        return result[0], updated_history, gr.Info("Query Processed")
+
     else:
         return await return_protection_message(msg, history)
 
 
 async def return_protection_message(msg, history):
 
-    new_message = (msg, "Your query appears a prompt injection. I would prefer Not to answer it.")
+    new_message = (msg, "Your query appears inappropriate. Do you have any other question?I am here to help.. ")
     updated_history = history + [new_message]
-    return "", updated_history
+    return "", updated_history, gr.Warning("Query is Inapproprite..")
+
                     
 
 async def handle_ingest() -> gr.Info:
@@ -74,6 +88,11 @@ async def record_feedback(feedback, msg ) -> gr.Info:
     logger.info(feedback)
     logger.info(msg)
 
+
+    if not msg.strip():
+        logger.error(f"No Comments provided")
+        return gr.Info("Please Enter Some Feed back First"), ""
+
     message, _ = await ws_client.handle_request(feedback, {"comment": msg})
     return gr.Info(message) if "success" in message.lower() else gr.Warning(message), ""
 
@@ -107,7 +126,7 @@ async def record_feedback(feedback, msg ) -> gr.Info:
             margin-top: 0.25rem;
             flex: 0 0 auto;
         }
-        #chatbot {
+        #chatbot-left {
             border: 1px solid #E5E7EB;
             border-radius: 8px;
             background-color: #FFFFFF;
@@ -118,6 +137,24 @@ async def record_feedback(feedback, msg ) -> gr.Info:
             flex-direction: column;
             overflow-y: auto; /* To allow scrolling if content overflows */
             min-height: 62vh; 
+            text-direction: left;
+            direction: left;
+            text-align: left;
+        }
+        #chatbot-right {
+            border: 1px solid #E5E7EB;
+            border-radius: 8px;
+            background-color: #FFFFFF;
+            box-shadow: 0 2px 8px rgba(0, 0, 0, 0.1);
+            flex: 1 1 auto;
+            min-height: 0;
+            display: flex;
+            flex-direction: column;
+            overflow-y: auto; /* To allow scrolling if content overflows */
+            min-height: 62vh; 
+            text-direction: right;
+            direction: right;
+            text-align: right;
         }
         #feedback-button {
             max-width: 0.25vh;
@@ -141,7 +178,7 @@ async def record_feedback(feedback, msg ) -> gr.Info:
     chatbot = gr.Chatbot(
         show_label=False,
         container=True,
-        elem_id="chatbot"
+        elem_id="chatbot-left"
     )
 
     with gr.Row(elem_id="feedback-container"):
@@ -173,7 +210,7 @@ async def record_feedback(feedback, msg ) -> gr.Info:
     send_button.click(
         fn=search_click,
         inputs=[msg, chatbot],
-        outputs=[msg, chatbot]
+        outputs=[msg, chatbot, status_box]
     )
     clear_button.click(
         fn=clear_chat,
@@ -203,4 +240,3 @@ async def record_feedback(feedback, msg ) -> gr.Info:
         share=False,
         debug=True,
         show_error=True,)
-
diff --git a/docker-compose.yml b/docker-compose.yml
@@ -13,9 +13,22 @@ services:
       - capec-network
     hostname: rag-server
     volumes:
+      - ./src:/app/src
+      - ./capec-dataset:/app/capec-dataset
+      - ./.env:/app/.env
       - ./src/index/index/:/app/src/index/index/
     environment:
       - TOKENIZERS_PARALLELISM=false
+    command:
+      [
+        "uvicorn",
+        "server:app",
+        "--host",
+        "0.0.0.0",
+        "--port",
+        "8000",
+        "--reload"
+      ]
 
   client:
     build:
@@ -28,6 +41,11 @@ services:
     environment:
       - SERVER_HOST=rag-server
       - SERVER_PORT=8000
+    volumes:
+      - ./src:/app/src
+      - ./client.py:/app/client.py
+      - ./client-requirements.txt:/app/client-requirements.txt
+    command: ["python", "client.py"]
 
   qdrant:
     image: qdrant/qdrant:v0.10.1
diff --git a/server.py b/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")
 
diff --git a/src/config/config.py b/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
diff --git a/src/docker-files/Dockerfile.client b/src/docker-files/Dockerfile.client
@@ -13,8 +13,6 @@ RUN pip install --upgrade pip && \
 
 # Copy only the required files for the application
 COPY client.py .
-COPY src/ ./src/
-
 
 # Run the application
 CMD ["python", "client.py"]
diff --git a/src/docker-files/Dockerfile.server b/src/docker-files/Dockerfile.server
@@ -10,10 +10,6 @@ RUN pip install --upgrade pip && \
     pip install -r requirements.txt
     
 COPY server.py .
-COPY src/ ./src/
-COPY .env .
-COPY capec-dataset/ ./capec-dataset/
-
 
 # Set Python to run in unbuffered mode
 ENV PYTHONUNBUFFERED=1
diff --git a/src/parser/csv_parser.py b/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
diff --git a/src/qdrant/qdrant_utils.py b/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)}")
diff --git a/src/reranker/re_ranking.py b/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 = [
diff --git a/src/utils/utils.py b/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.")
+
+
diff --git a/src/websocket/web_socket_client.py b/src/websocket/web_socket_client.py
@@ -1,10 +1,13 @@
-import gradio as gr
 import websockets
 import json
 import asyncio
 from typing import Tuple, List, Optional, Dict, Any
 from loguru import logger
 
+from bidi.algorithm import get_display
+import unicodedata
+
+
 from src.config.config import Config
 
 
@@ -19,6 +22,16 @@ async def connect(self):
             self.websocket = await websockets.connect(self.uri)
             logger.info("Connected to WebSocket server")
         return self.websocket
+
+    async def get_text_direction(self, text: str):
+        # Use Unicode character properties to check if the text is RTL
+        for char in text:
+            # Simulate asynchronous behavior (if needed for consistency)
+            if unicodedata.bidirectional(char) in ['R', 'AL']:  # Right-to-left or Arabic Letter
+                return "right"
+        return "left"
+
+    
         
     async def disconnect(self):
         if self.websocket:
@@ -104,12 +117,6 @@ async def handle_request(
         """
 
         logger.info("Into handle search function..")
-
-        if action ==  "search":
-            query = payload["query"]
-            if not query.strip():
-                logger.error(f"No input provided")
-                return "", [(payload.get("query", ""), "No query Entered")]
             
         try:
             
@@ -159,10 +166,14 @@ async def _handle_websocket_communication(
                 result = response_data.get("result", "No response from server")
                 if result:
                     if action == "search":
+                        direction = await self.get_text_direction(result)
+
+                        logger.info(direction)
+
                         history = payload.get("history", [])
                         new_message = (payload.get("query", ""), result)
                         updated_history = history + [new_message]
-                        return "", updated_history
+                        return "", result, direction
                     elif action == "ingest_data":
                         return result, []
                     elif action == "positive":
