fix error handle in main

main.py

@@ -1,4 +1,4 @@
-from fastapi import FastAPI, WebSocket
+from fastapi import FastAPI, WebSocket, WebSocketDisconnect
 from urllib.parse import urlparse, parse_qsl
 from transcribe.whisper_llm_serve import PyWhiperCppServe
 from uuid import uuid1
@@ -68,10 +68,12 @@ async def translate(websocket: WebSocket):
         client.set_lang(from_lang, to_lang)
         logger.info(f"Source lange: {from_lang}  -> Dst lange: {to_lang}")
         await websocket.accept()
-    while True:
-        frame_data = await get_audio_from_websocket(websocket)
-        client.add_frames(frame_data)
-
+    try:
+        while True:
+            frame_data = await get_audio_from_websocket(websocket)
+            client.add_frames(frame_data)
+    except WebSocketDisconnect:
+        return 
 
 if __name__ == '__main__':
     freeze_support()

transcribe/whisper_llm_serve.py

@@ -36,9 +36,10 @@ class PyWhiperCppServe(ServeClientBase):
         self.sample_rate = 16000
  
         self.send_ready_state()
-        
-        self.run_in_thread(self.speech_to_text)
-        self.run_in_thread(self.get_frame_from_queue)
+        self._translate_thread_stop = threading.Event()
+        self._frame_to_queue_thread_stop = threading.Event()
+        self.translate_thread = self.run_in_thread(self.speech_to_text)
+        self.frame_to_queue_thread = self.run_in_thread(self.get_frame_from_queue)
 
         self.text_sep = ""
 
@@ -46,6 +47,7 @@ class PyWhiperCppServe(ServeClientBase):
         t = threading.Thread(target=func)
         t.daemon = True
         t.start()
+        return t
 
     def send_ready_state(self):
         self.websocket.send(json.dumps({
@@ -72,7 +74,7 @@ class PyWhiperCppServe(ServeClientBase):
      
     
     def get_frame_from_queue(self,):
-        while True:
+        while not self._frame_to_queue_thread_stop.is_set():
             try:
                 frame_np = self._frame_queue.get(timeout=0.1)
                 with self.lock:
@@ -145,7 +147,7 @@ class PyWhiperCppServe(ServeClientBase):
 
     def speech_to_text(self):
         # c = 0
-        while True:
+        while not self._translate_thread_stop.is_set():
             if self.exit:
                 logger.info("Exiting speech to text thread")
                 break
@@ -221,11 +223,17 @@ class PyWhiperCppServe(ServeClientBase):
             )
     
     def send_to_client(self, data:TransResult):
-      
-        coro = self.websocket.send_text(
-            Message(result=data, request_id=self.client_uid).model_dump_json(by_alias=True)
-        )
-        asyncio.run(coro)
+        try:
+            coro = self.websocket.send_text(
+                Message(result=data, request_id=self.client_uid).model_dump_json(by_alias=True)
+            )
+            asyncio.run(coro)
+        except RuntimeError as e:
+            self.stop()
+            return 
+        except Exception as e:
+            logger.error(e)
+
     
 
     def get_audio_chunk_for_processing(self):
@@ -240,5 +248,6 @@ class PyWhiperCppServe(ServeClientBase):
         padded_audio = np.concatenate([silence, self.frames_np])
         return padded_audio.copy()
 
-    def cleanup(self):
-        return super().cleanup()#
+    def stop(self):
+        self._translate_thread_stop.set()
+        self._frame_to_queue_thread_stop.set()