feat: add redis

main.py

@@ -55,8 +55,30 @@ QA = RetrievalQA.from_chain_type(
     },
 )
 
+def sendPromptChain(QA, user_prompt):
+
+    res = QA(user_prompt)
+
+    answer, docs = res["result"], res["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))
+        )
+
+    return prompt_response_dict;
+
+
 socket_manager = WebSocketManager()
 
+
+
 app = FastAPI(title="homepage-app")
 api_app = FastAPI(title="api app")
 
@@ -143,20 +165,7 @@ def predict(data: Predict):
     try:
         user_prompt = data.prompt
         if user_prompt:
-            res = QA(user_prompt)
-
-            answer, docs = res["result"], res["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))
-                )
+            prompt_response_dict = sendPromptChain(QA, user_prompt)
 
             return {"response": prompt_response_dict}
         else:
@@ -213,7 +222,7 @@ async def websocket_endpoint(websocket: WebSocket,  room_id: str, user_id: int):
     message = {
         "user_id": user_id,
         "room_id": room_id,
-        "message": f"User {user_id} connected to room - {room_id}"
+        "message": f"Student {user_id} connected to the classroom"
     }
 
     await socket_manager.broadcast_to_room(room_id, json.dumps(message))
@@ -228,23 +237,9 @@ async def websocket_endpoint(websocket: WebSocket,  room_id: str, user_id: int):
                 "message": data
             }
 
-            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 = sendPromptChain(QA, 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, prompt_response_dict)
 
     except WebSocketDisconnect:
         await socket_manager.remove_user_from_room(room_id, websocket)