feat: add broadcast

main.py

@@ -1,13 +1,13 @@
-from typing import Any, Dict, Union
-
 import os
 import glob
 import shutil
 import subprocess
 import torch
+import json
 
 from fastapi import FastAPI, HTTPException, UploadFile, WebSocket, WebSocketDisconnect
 from fastapi.staticfiles import StaticFiles
+from websocket.socketManager import WebSocketManager
 
 from pydantic import BaseModel
 
@@ -55,6 +55,8 @@ QA = RetrievalQA.from_chain_type(
     },
 )
 
+socket_manager = WebSocketManager()
+
 app = FastAPI(title="homepage-app")
 api_app = FastAPI(title="api app")
 
@@ -162,8 +164,6 @@ def predict(data: Predict):
     except Exception as e:
         raise HTTPException(status_code=500, detail=f"Error occurred: {str(e)}")
 
-
-
 @api_app.post("/save_document/")
 async def create_upload_file(file: UploadFile):
     # Get the file size (in bytes)
@@ -204,31 +204,57 @@ async def create_upload_file(file: UploadFile):
 
     return {"filename": file.filename}
 
-@api_app.websocket("/ws/{client_id}")
-async def websocket_endpoint(websocket: WebSocket,  client_id: str):
+@api_app.websocket("/ws/{room_id}/{user_id}")
+async def websocket_endpoint(websocket: WebSocket,  room_id: str, user_id: int):
     global QA
 
-    await websocket.accept()
+    await socket_manager.add_user_to_room(room_id, websocket)
+
+    message = {
+        "user_id": user_id,
+        "room_id": room_id,
+        "message": f"User {user_id} connected to room - {room_id}"
+    }
+
+    await socket_manager.broadcast_to_room(room_id, json.dumps(message))
 
     try:
         while True:
-            user_prompt = await websocket.receive_text()
-            response = QA(inputs=user_prompt, return_only_outputs=True, tags=f'{client_id}', include_run_info=True)
-            answer, docs = response["result"], response["source_documents"]
+            data = await websocket.receive_text()
 
-            prompt_response_dict = {
-                "Prompt": user_prompt,
-                "Answer": answer,
+            message = {
+                "user_id": user_id,
+                "room_id": room_id,
+                "message": data
             }
 
-            prompt_response_dict["Sources"] = []
-            for document in docs:
-                prompt_response_dict["Sources"].append(
-                    (os.path.basename(str(document.metadata["source"])), str(document.page_content))
-                )
-            await websocket.send_json(prompt_response_dict)
+            await socket_manager.broadcast_to_room(room_id, json.dumps(message))
+
+            # user_prompt = await websocket.receive_text()
+            # response = QA(inputs=user_prompt, return_only_outputs=True, tags=f'{client_id}', include_run_info=True)
+            # answer, docs = response["result"], response["source_documents"]
+
+            # prompt_response_dict = {
+            #     "Prompt": user_prompt,
+            #     "Answer": answer,
+            # }
+
+            # prompt_response_dict["Sources"] = []
+            # for document in docs:
+            #     prompt_response_dict["Sources"].append(
+            #         (os.path.basename(str(document.metadata["source"])), str(document.page_content))
+            #     )
+            # await websocket.send_json(prompt_response_dict)
 
     except WebSocketDisconnect:
-        print('disconnect')
+        await socket_manager.remove_user_from_room(room_id, websocket)
+
+        message = {
+            "user_id": user_id,
+            "room_id": room_id,
+            "message": f"User {user_id} disconnected from room - {room_id}"
+        }
+
+        await socket_manager.broadcast_to_room(room_id, json.dumps(message))
     except RuntimeError as error:
         print(error)

redisPubSubManger.py

@@ -0,0 +1,69 @@
+import asyncio
+import redis.asyncio as aioredis
+import json
+from fastapi import WebSocket
+
+
+class RedisPubSubManager:
+    """
+        Initializes the RedisPubSubManager.
+
+    Args:
+        host (str): Redis server host.
+        port (int): Redis server port.
+    """
+
+    def __init__(self, host='localhost', port=6379):
+        self.redis_host = host
+        self.redis_port = port
+        self.pubsub = None
+
+    async def _get_redis_connection(self) -> aioredis.Redis:
+        """
+        Establishes a connection to Redis.
+
+        Returns:
+            aioredis.Redis: Redis connection object.
+        """
+        return aioredis.Redis(host=self.redis_host,
+                              port=self.redis_port,
+                              auto_close_connection_pool=False)
+
+    async def connect(self) -> None:
+        """
+        Connects to the Redis server and initializes the pubsub client.
+        """
+        self.redis_connection = await self._get_redis_connection()
+        self.pubsub = self.redis_connection.pubsub()
+
+    async def _publish(self, room_id: str, message: str) -> None:
+        """
+        Publishes a message to a specific Redis channel.
+
+        Args:
+            room_id (str): Channel or room ID.
+            message (str): Message to be published.
+        """
+        await self.redis_connection.publish(room_id, message)
+
+    async def subscribe(self, room_id: str) -> aioredis.Redis:
+        """
+        Subscribes to a Redis channel.
+
+        Args:
+            room_id (str): Channel or room ID to subscribe to.
+
+        Returns:
+            aioredis.ChannelSubscribe: PubSub object for the subscribed channel.
+        """
+        await self.pubsub.subscribe(room_id)
+        return self.pubsub
+
+    async def unsubscribe(self, room_id: str) -> None:
+        """
+        Unsubscribes from a Redis channel.
+
+        Args:
+            room_id (str): Channel or room ID to unsubscribe from.
+        """
+        await self.pubsub.unsubscribe(room_id)

requirements.txt

@@ -29,7 +29,7 @@ uvicorn
 fastapi
 websockets
 pydantic
-redis
+aioredis
 
 # Streamlit related
 streamlit

run.sh

@@ -1,5 +1,4 @@
 # Redis Support uncomment this lines
-# redis-cli --version
-# nohup redis-server &
+nohup redis-server &
 
 uvicorn "main:app" --port 7860 --host 0.0.0.0

test.txt

@@ -1 +0,0 @@
-dkdaniz is an avatar of instagram, create by daniel marques

webSocketManger.py

@@ -0,0 +1,70 @@
+class WebSocketManager:
+    def __init__(self):
+        """
+        Initializes the WebSocketManager.
+
+        Attributes:
+            rooms (dict): A dictionary to store WebSocket connections in different rooms.
+            pubsub_client (RedisPubSubManager): An instance of the RedisPubSubManager class for pub-sub functionality.
+        """
+        self.rooms: dict = {}
+        self.pubsub_client = RedisPubSubManager()
+
+    async def add_user_to_room(self, room_id: str, websocket: WebSocket) -> None:
+        """
+        Adds a user's WebSocket connection to a room.
+
+        Args:
+            room_id (str): Room ID or channel name.
+            websocket (WebSocket): WebSocket connection object.
+        """
+        await websocket.accept()
+
+        if room_id in self.rooms:
+            self.rooms[room_id].append(websocket)
+        else:
+            self.rooms[room_id] = [websocket]
+
+            await self.pubsub_client.connect()
+            pubsub_subscriber = await self.pubsub_client.subscribe(room_id)
+            asyncio.create_task(self._pubsub_data_reader(pubsub_subscriber))
+
+    async def broadcast_to_room(self, room_id: str, message: str) -> None:
+        """
+        Broadcasts a message to all connected WebSockets in a room.
+
+        Args:
+            room_id (str): Room ID or channel name.
+            message (str): Message to be broadcasted.
+        """
+        await self.pubsub_client._publish(room_id, message)
+
+    async def remove_user_from_room(self, room_id: str, websocket: WebSocket) -> None:
+        """
+        Removes a user's WebSocket connection from a room.
+
+        Args:
+            room_id (str): Room ID or channel name.
+            websocket (WebSocket): WebSocket connection object.
+        """
+        self.rooms[room_id].remove(websocket)
+
+        if len(self.rooms[room_id]) == 0:
+            del self.rooms[room_id]
+            await self.pubsub_client.unsubscribe(room_id)
+
+    async def _pubsub_data_reader(self, pubsub_subscriber):
+        """
+        Reads and broadcasts messages received from Redis PubSub.
+
+        Args:
+            pubsub_subscriber (aioredis.ChannelSubscribe): PubSub object for the subscribed channel.
+        """
+        while True:
+            message = await pubsub_subscriber.get_message(ignore_subscribe_messages=True)
+            if message is not None:
+                room_id = message['channel'].decode('utf-8')
+                all_sockets = self.rooms[room_id]
+                for socket in all_sockets:
+                    data = message['data'].decode('utf-8')
+                    await socket.send_text(data)