feat: add websocket

main.py

@@ -213,7 +213,26 @@ async def create_upload_file(file: UploadFile):
 
 @api_app.websocket("/ws/{user_id}")
 async def websocket_endpoint_student(websocket: WebSocket, user_id: str):
-    global QA
+    DB = Chroma(
+        persist_directory=PERSIST_DIRECTORY,
+        embedding_function=EMBEDDINGS,
+        client_settings=CHROMA_SETTINGS,
+    )
+
+    RETRIEVER = DB.as_retriever()
+
+    newInstanceQA = RetrievalQA.from_chain_type(
+        llm=LLM,
+        chain_type="stuff",
+        retriever=RETRIEVER,
+        return_source_documents=SHOW_SOURCES,
+        chain_type_kwargs={
+            "prompt": prompt,
+            "memory": memory
+        },
+    )
+
+    QA = socket_manager.get_instance_qa(user_id, newInstanceQA)
 
     message = {
         "message": f"Student {user_id} connected"
@@ -243,7 +262,26 @@ async def websocket_endpoint_student(websocket: WebSocket, user_id: str):
 
 @api_app.websocket("/ws/{room_id}/{user_id}")
 async def websocket_endpoint_room(websocket: WebSocket, room_id: str, user_id: str):
-    global QA
+    DB = Chroma(
+        persist_directory=PERSIST_DIRECTORY,
+        embedding_function=EMBEDDINGS,
+        client_settings=CHROMA_SETTINGS,
+    )
+
+    RETRIEVER = DB.as_retriever()
+
+    newInstanceQA = RetrievalQA.from_chain_type(
+        llm=LLM,
+        chain_type="stuff",
+        retriever=RETRIEVER,
+        return_source_documents=SHOW_SOURCES,
+        chain_type_kwargs={
+            "prompt": prompt,
+            "memory": memory
+        },
+    )
+
+    QA = socket_manager.get_instance_qa(room_id, newInstanceQA)
 
     message = {
         "message": f"Student {user_id} connected to the classroom"

prompt_template_utils.py

@@ -15,8 +15,9 @@ from langchain.prompts import PromptTemplate
 
 # 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."""
+system_prompt = """It's a useful assistant that 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. Your answer cannot be more than 10 sentences long.
+If you can't answer, just say "I don't know" and don't try to work out an answer to respond to the user."""
 
 def get_prompt_template(system_prompt=system_prompt, promptTemplate_type=None, history=False):
     if promptTemplate_type == "llama":

websocket/socketManager.py

@@ -1,9 +1,13 @@
+from typing import Any
 import asyncio
 import redis.asyncio as aioredis
 import json
 from fastapi import WebSocket
 
 
+
+
+
 class RedisPubSubManager:
     """
         Initializes the RedisPubSubManager.
@@ -80,6 +84,7 @@ class WebSocketManager:
             pubsub_client (RedisPubSubManager): An instance of the RedisPubSubManager class for pub-sub functionality.
         """
         self.rooms: dict = {}
+        self.qa: dict = {}
         self.pubsub_client = RedisPubSubManager()
 
     async def add_user_to_room(self, room_id: str, websocket: WebSocket) -> None:
@@ -96,7 +101,6 @@ class WebSocketManager:
             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))
@@ -120,6 +124,7 @@ class WebSocketManager:
             websocket (WebSocket): WebSocket connection object.
         """
         self.rooms[room_id].remove(websocket)
+        self.qa.pop(room_id, None)
 
         if len(self.rooms[room_id]) == 0:
             del self.rooms[room_id]
@@ -140,3 +145,11 @@ class WebSocketManager:
                 for socket in all_sockets:
                     data = message['data'].decode('utf-8')
                     await socket.send_text(data)
+
+    async def get_instance_qa(self, room_id: str, QA: Any):
+        if room_id in self.qa:
+            return self.qa[room_id]
+
+        self.qa[room_id] = QA
+        return self.qa[room_id]
+