fix: add redis cache

main.py

@@ -6,7 +6,7 @@ import shutil
 import subprocess
 import redis
 import torch
-import threading
+import concurrent.futures
 
 from fastapi import FastAPI, HTTPException, UploadFile, WebSocket, WebSocketDisconnect
 from fastapi.staticfiles import StaticFiles
@@ -69,13 +69,6 @@ QA = RetrievalQA.from_chain_type(
     },
 )
 
-async def pubsub_listener(websocket, client_id):
-    pubsub = redisClient.pubsub()
-    pubsub.subscribe(f'{client_id}')
-
-    while True:
-        message = pubsub.listen()
-        await websocket.send_text(f'{message}')
 
 app = FastAPI(title="homepage-app")
 api_app = FastAPI(title="api app")
@@ -230,11 +223,17 @@ async def websocket_endpoint(websocket: WebSocket,  client_id: int):
     try:
         while True:
             prompt = await websocket.receive_text()
+            pubsub = redisClient.pubsub()
+            pubsub.subscribe(f'{client_id}')
+
+            with concurrent.futures.ThreadPoolExecutor() as executor:
+               executor.submit(QA(inputs=prompt, return_only_outputs=True, callbacks=[MyCustomSyncHandler()], tags=f'{client_id}', include_run_info=True))
 
-            pubsub_thread = threading.Thread(target=pubsub_listener(websocket, client_id))
-            pubsub_thread.start()
+            for item in pubsub.listen():
+                print(item)
+                if item['type'] == 'message':
+                    await websocket.send_text(f'{item["data"]}')
 
-            QA(inputs=prompt, return_only_outputs=True, callbacks=[MyCustomSyncHandler()], tags=f'{client_id}', include_run_info=True)
 
     except WebSocketDisconnect:
         print('disconnect')