Update app.py

app.py

@@ -1,6 +1,8 @@
 
 from langfuse import Langfuse
 from langfuse.decorators import observe, langfuse_context
+from fastapi import WebSocketDisconnect
+import asyncio
 
 from config.config import settings
 from services.llama_generator import LlamaGenerator
@@ -68,15 +70,17 @@ async def detect_wakeword(audio_chunk: bytes) -> bool:
     # You might want to use libraries like Porcupine or build your own wake word detector
     return True
 
+
 @app.websocket("/ws")
 async def websocket_endpoint(websocket: WebSocket):
+    # Create the queue outside the try block
+    audio_queue = asyncio.Queue()
+    stream_task = None
+
     await websocket.accept()
     try:
-        # Use a queue to manage audio chunks
-        audio_queue = asyncio.Queue()
-        
         # Create a task to process the audio stream
-        stream_task = asyncio.create_task(process_audio_stream(audio_queue))
+        stream_task = asyncio.create_task(process_audio_stream(audio_queue, websocket))
         
         # Main receive loop
         while True:
@@ -104,20 +108,26 @@ async def websocket_endpoint(websocket: WebSocket):
         print(f"WebSocket endpoint error: {e}")
     
     finally:
-        # Cancel the stream processing task
-        stream_task.cancel()
+        # Cancel the stream processing task if it exists
+        if stream_task:
+            stream_task.cancel()
+            try:
+                await stream_task  # Wait for the task to be fully cancelled
+            except asyncio.CancelledError:
+                pass
+        
         try:
             await websocket.close(code=1000)
         except Exception as close_error:
             print(f"Error closing WebSocket: {close_error}")
 
-async def process_audio_stream(audio_queue: asyncio.Queue) -> AsyncGenerator[str, None]:
+async def process_audio_stream(audio_queue: asyncio.Queue, websocket: WebSocket) -> AsyncGenerator[str, None]:
     buffer = []
     is_speaking = False
     silence_frames = 0
     
-    while True:
-        try:
+    try:
+        while True:
             # Get audio data from queue with timeout
             try:
                 audio_data = await asyncio.wait_for(audio_queue.get(), timeout=5.0)
@@ -172,30 +182,30 @@ async def process_audio_stream(audio_queue: asyncio.Queue) -> AsyncGenerator[str
                                 bot_response_de = from_en_translation(response, desired_language)
                                 
                                 # Stream the response
-                                yield json.dumps({
+                                response_data = json.dumps({
                                     "user_text": user_speech_text,
                                     "response_de": bot_response_de,
                                     "response_en": response
                                 })
+                                await websocket.send_text(response_data)
                                 
                                 # Generate and stream audio response
                                 bot_voice = tts(bot_response_de, desired_language)
                                 bot_voice_bytes = tts_to_bytesio(bot_voice)
-                                yield json.dumps({
+                                audio_data = json.dumps({
                                     "audio": bot_voice_bytes.decode('latin1')
                                 })
+                                await websocket.send_text(audio_data)
                     
                     except Exception as processing_error:
                         print(f"Error processing speech utterance: {processing_error}")
+    
+    except asyncio.CancelledError:
+        # Handle task cancellation
+        print("Audio stream processing task cancelled")
+    except Exception as e:
+        print(f"Unexpected error in audio stream processing: {e}")
         
-        except asyncio.CancelledError:
-            # Handle task cancellation
-            break
-        except Exception as e:
-            print(f"Unexpected error in audio stream processing: {e}")
-            # Prevent tight error loop
-            await asyncio.sleep(1)
-
 @app.get("/", response_class=HTMLResponse)
 async def get_index():
     with open("static/index.html") as f: