fix: add websocket in handlerToken

main.py

@@ -259,21 +259,24 @@ async def websocket_endpoint(websocket: WebSocket):
 
     await websocket.accept()
 
+    print(f'receive')
+
     try:
+        prompt = await websocket.receive_text()
+        res = QA(prompt)
+
         while True:
-            tokenMessageLLM.after_create(lambda now, old: print(f"{old} updated to {now}."))
 
-            print(tokenMessageLLM)
+            print(prompt)
+
+            tokenMessageLLM.after_update(lambda now, old: print(f"{old} updated to {now}."))
+
+            print(f'tokenMessageLLM')
+
 
-            data = await websocket.receive_text()
-            res = QA(data)
             print(res)
 
     except WebSocketDisconnect:
         print('disconnect')
     except RuntimeError as error:
         print(error)
-
-
-
-