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

Dev branch

Author: usman-ce

Dockerfile

@@ -7,175 +7,7 @@
 from loguru import logger
 
 from src.config.config import Config
-
-
-
-class WebSocketClient:
-    def __init__(self, uri: str = "ws://rag-server:8000/ws"):
-        self.uri = uri
-        self.websocket: Optional[websockets.WebSocketClientProtocol] = None
-        self._connection_lock = asyncio.Lock()
-        
-    async def connect(self):
-        if not self.websocket:
-            self.websocket = await websockets.connect(self.uri)
-            logger.info("Connected to WebSocket server")
-        return self.websocket
-        
-    async def disconnect(self):
-        if self.websocket:
-            await self.websocket.close()
-            self.websocket = None
-            logger.info("Disconnected from WebSocket server")
-            
-    async def send_search_query(self, query: str, payload:str) -> str:
-        try:
-            # Ensure connection is established
-            if not self.websocket:
-                await self.connect()
-                
-            # Send search query
-            await self.websocket.send(json.dumps({"query": query}))
-            
-            # Wait for response
-            response = await self.websocket.recv()
-            data = json.loads(response)
-            
-            if data["status"] == "success":
-                # Format results for display
-                results = data["results"]
-                formatted_results = "\n\n".join([
-                    f"Title: {result['title']}\nURL: {result['url']}"
-                    for result in results
-                ])
-                return formatted_results, payload
-            else:
-                return f"Error: {data['message'], payload}"
-                
-        except websockets.exceptions.ConnectionClosedError:
-            logger.error("Connection closed unexpectedly. Attempting to reconnect...")
-            self.websocket = None
-            return "Connection lost. Please try again."
-        except Exception as e:
-            logger.error(f"Error during search: {str(e)}")
-            return f"Error: {str(e)}"
-        
-    async def ensure_connection(self): 
-        if self.websocket is None or self.websocket.closed:
-            await self.connect()
-
-
-    async def connect(self):
-        """Establish WebSocket connection and start heartbeat monitoring"""
-        async with self._connection_lock:
-            try:
-                uri = Config.WEBSOCKET_URI
-                if not uri.startswith(('ws://', 'wss://')):
-                    logger.error("Invalid WebSocket URI format")
-                    return
-
-                self.websocket = await websockets.connect(
-                    uri,
-                    ping_interval=20,
-                    ping_timeout=60,
-                    max_size=10_485_760
-                )
-                logger.info("Connected to server.....")
-
-                # self._heartbeat_task = asyncio.create_task(self._heartbeat())
-
-                return True
-
-            except Exception as e:
-                logger.error(f"Connection error: {e}")
-                return False
-            
-    async def handle_request(
-        self, action: str, payload: dict = {}
-    ) -> Tuple[str, List[Tuple[str, str]]]:
-        """
-        Handle WebSocket requests to the server.
-
-        Args:
-            action (str): The action to perform (e.g., 'search', 'ingest_data').
-            payload (dict): The payload containing request data.
-
-        Returns:
-            Tuple[str, List[Tuple[str, str]]]: A tuple containing the response message
-            and updated chat history.
-        """
-
-        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:
-            
-            logger.info("Ensuring Connection....")
-            await self.ensure_connection()
-
-            result = await self._handle_websocket_communication(action, payload)
-            return  result
-        
-        except Exception as e:
-            logger.error(f"Connection error: {e}")
-            await self.disconnect() 
-            return "", [(payload.get("query", ""), f"Connection error: {str(e)}")]
-        
-        
-    async def _handle_websocket_communication(
-        self, action: str, payload: dict
-    ) -> Tuple[str, List[Tuple[str, str]]]:
-        """
-        Handle the WebSocket communication with the server.
-
-        Args:
-            action (str): The action to perform.
-            payload (dict): The payload containing request data.
-
-        Returns:
-            Tuple[str, List[Tuple[str, str]]]: A tuple containing the response message
-            and updated chat history.
-        """
-        try:
-            await self.websocket.send(json.dumps({"action": action, "payload": payload}))
-
-            while True:
-                response = await self.websocket.recv()
-                response_data = json.loads(response)
-
-                logger.info("Response received...")
-
-                # Handle heartbeat
-                if response_data.get("type") == "ping":
-                    await self.websocket.send(json.dumps({
-                        "action": "pong",
-                        "timestamp": response_data.get("timestamp")
-                    }))
-                    continue
-
-                result = response_data.get("result", "No response from server")
-                if result:
-                    if action == "search":
-                        history = payload.get("history", [])
-                        new_message = (payload.get("query", ""), result)
-                        updated_history = history + [new_message]
-                        return "", updated_history
-                    elif action == "ingest_data":
-                        return result, []
-
-                error = response_data.get("error")
-                if error:
-                    return "", [(payload.get("query", ""), f"Error: {error}")]
-
-        except Exception as e:
-            logger.error(f"Communication error: {e}")
-            return "", [(payload.get("query", ""), f"Communication error: {str(e)}")]
-
+from src.websocket.web_socket_client import WebSocketClient
 
 
 ws_client = WebSocketClient(Config.WEBSOCKET_URI)
@@ -212,46 +44,56 @@ def clear_chat() -> Optional[List[Tuple[str, str]]]:
         return None
 
 
-# Create Gradio interface
 with gr.Blocks(
-        title="Capec Chatbot",
-        theme=gr.themes.Soft(),
-        css="""
-            .gradio-container {
-                max-width: 700px; 
-                margin: auto; 
-                font-family: Arial, sans-serif;
-            }
-            #header {
-                text-align: center; 
-                font-size: 1.5rem; 
-                font-weight: bold; 
-                color: #008080; 
-                padding: 0.125rem;
-            }
-            #input-container {
-                display: flex; 
-                align-items: center;
-                background-color: #f7f7f8;
-                padding: 0.25rem; 
-                border-radius: 8px;
-                margin-top: 0.25rem;
-            }
-            #chatbot {
-                border: 1px solid #E5E7EB;
-                border-radius: 8px;
-                background-color: #FFFFFF;
-                box-shadow: 0 2px 8px rgba(0, 0, 0, 0.1);
-            }
-            .gr-button-primary {
-                background-color: #008080;
-                border-color: #008080;
-            }
-            .gr-button-primary:hover {
-                background-color: #006666;
-            }
-        """
-    ) as demo:
+    title="CAPEC RAG Chatbot",
+    theme=gr.themes.Soft(),
+    css="""
+        .gradio-container {
+            max-width: 700px; 
+            margin: auto; 
+            font-family: Arial, sans-serif;
+            display: flex;
+            flex-direction: column;
+            height: 100vh;
+        }
+        #header {
+            text-align: center; 
+            font-size: 1.5rem; 
+            font-weight: bold; 
+            color: #008080; 
+            padding: 0.125rem;
+            flex: 0 0 auto;
+        }
+        #input-container {
+            display: flex; 
+            align-items: center;
+            background-color: #f7f7f8;
+            padding: 0.25rem; 
+            border-radius: 8px;
+            margin-top: 0.25rem;
+            flex: 0 0 auto;
+        }
+        #chatbot {
+            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: 72vh; 
+        }
+        .gr-button-primary {
+            background-color: #008080;
+            border-color: #008080;
+        }
+        .gr-button-primary:hover {
+            background-color: #006666;
+        }
+    """
+) as demo:
 
     # Header
     gr.Markdown(
@@ -260,7 +102,6 @@ def clear_chat() -> Optional[List[Tuple[str, str]]]:
 
     # Chatbot Component
     chatbot = gr.Chatbot(
-        height=450,
         show_label=False,
         container=True,
         elem_id="chatbot"
@@ -277,6 +118,7 @@ def clear_chat() -> Optional[List[Tuple[str, str]]]:
         )
         send_button = gr.Button("Send", variant="primary", scale=1)
         clear_button = gr.Button("Clear Chat", variant="secondary")
+        
 
     # Button Functionality
     send_button.click(
@@ -302,3 +144,4 @@ def clear_chat() -> Optional[List[Tuple[str, str]]]:
         share=False,
         debug=True,
         show_error=True,)
+

capec-dataset/1000.csv

@@ -6,7 +6,7 @@ services:
   server:
     build:
       context: .
-      dockerfile: Dockerfile.server
+      dockerfile: src/docker-files/Dockerfile.server
     ports:
       - "8000:8000"
     networks:
@@ -20,7 +20,7 @@ services:
   client:
     build:
       context: .
-      dockerfile: Dockerfile.client
+      dockerfile: src/docker-files/Dockerfile.client
     networks:
       - capec-network
     ports:

capec-dataset/3000.csv

@@ -1,18 +1,14 @@
-# sentence-transformers==3.2.1
+fastapi==0.115.5
+websockets==12.0
 qdrant-client==1.12.0
 python-multipart==0.0.16
 scikit-learn==1.5.2
-langchain-core==0.3.13
+langchain-core==0.3.18
 gradio==4.44.1
 loguru==0.7.2
-groq==0.11.0
 python-dotenv==1.0.1
-fastapi
-websockets
-# google-auth-oauthlib==1.2.1
-# google-auth==2.35.0
-# google-api-python-client==2.151.0
 llama-index==0.11.21
 llama-index-vector-stores-qdrant==0.3.3
-llama-index-llms-groq==0.2.0
-llama-index-embeddings-huggingface==0.3.1
+langchain-groq==0.2.1
+llama-index-embeddings-huggingface==0.3.1
+langchain==0.3.7