feat: add broadcast

.dockerignore

@@ -2,8 +2,10 @@
 !*.py
 !*.txt
 !*.html
+!*.sh
 !requirements.txt
 !SOURCE_DOCUMENTS
 !static
-!*.sh
+!websocket
+
 

prompt_template_utils.py

@@ -1,92 +0,0 @@
-"""
-This file implements prompt template for llama based models.
-Modify the prompt template based on the model you select.
-This seems to have significant impact on the output of the LLM.
-"""
-
-from langchain.memory import ConversationBufferMemory
-from langchain.prompts import PromptTemplate
-
-# this is specific to Llama-2.
-
-# system_prompt = """You are a helpful assistant, you will use the context and documents provided in the training to answer users questions.
-# Read the context provided before answering questions and think step by step. If you can't answer a user's question based on the
-# context provided, inform the user. Don't use any other information to answer the user."""
-
-# system_prompt = """You are a helpful assistant, and you will use the context and documents provided in the training to answer users' questions. Please read the context provided carefully before responding to questions and follow a step-by-step thought process. If you cannot answer a user's question based on the provided context, please inform the user. Do not use any other information to answer the user. Provide a detailed response based on the content of locally trained documents."""
-
-system_prompt = """It's a useful assistant who will use the context and documents provided in the training to answer users' questions.
-Read the context provided before answering the questions and think step by step. If you can't answer, just say "I don't know" and don't try to put together an answer to respond to the user."""
-
-def get_prompt_template(system_prompt=system_prompt, promptTemplate_type=None, history=False):
-    if promptTemplate_type == "llama":
-        B_INST, E_INST = "[INST]", "[/INST]"
-        B_SYS, E_SYS = "<<SYS>>\n", "\n<</SYS>>\n\n"
-        SYSTEM_PROMPT = B_SYS + system_prompt + E_SYS
-        if history:
-            instruction = """
-            Context: {history} \n {context}
-            User: {question}"""
-
-            prompt_template = B_INST + SYSTEM_PROMPT + instruction + E_INST
-            prompt = PromptTemplate(input_variables=["history", "context", "question"], template=prompt_template)
-        else:
-            instruction = """
-            Context: {context}
-            User: {question}"""
-
-            prompt_template = B_INST + SYSTEM_PROMPT + instruction + E_INST
-            prompt = PromptTemplate(input_variables=["context", "question"], template=prompt_template)
-    elif promptTemplate_type == "mistral":
-        B_INST, E_INST = "<s>[INST] ", " [/INST]"
-        if history:
-            prompt_template = (
-                B_INST
-                + system_prompt
-                + """
-
-            Context: {history} \n {context}
-            User: {question}"""
-                + E_INST
-            )
-            prompt = PromptTemplate(input_variables=["history", "context", "question"], template=prompt_template)
-        else:
-            prompt_template = (
-                B_INST
-                + system_prompt
-                + """
-
-            Context: {context}
-            User: {question}"""
-                + E_INST
-            )
-            prompt = PromptTemplate(input_variables=["context", "question"], template=prompt_template)
-    else:
-        # change this based on the model you have selected.
-        if history:
-            prompt_template = (
-                system_prompt
-                + """
-
-            Context: {history} \n {context}
-            User: {question}
-            Answer:"""
-            )
-            prompt = PromptTemplate(input_variables=["history", "context", "question"], template=prompt_template)
-        else:
-            prompt_template = (
-                system_prompt
-                + """
-
-            Context: {context}
-            User: {question}
-            Answer:"""
-            )
-            prompt = PromptTemplate(input_variables=["context", "question"], template=prompt_template)
-
-    memory = ConversationBufferMemory(input_key="question", memory_key="history")
-
-    return (
-        prompt,
-        memory,
-    )

redisPubSubManger.py → websocket/socketManager.py

@@ -67,3 +67,76 @@ class RedisPubSubManager:
             room_id (str): Channel or room ID to unsubscribe from.
         """
         await self.pubsub.unsubscribe(room_id)
+
+
+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)