update websocket again

mrrrme/backend/websocket.py

@@ -1,6 +1,5 @@
 """
-MrrrMe Backend - WebSocket Handler (FIXED - NO CRASHES)
-Complete error handling to prevent connection drops
+MrrrMe Backend - WebSocket Handler (FIXED - GLOBAL VARIABLE BUG)
 """
 
 from fastapi import WebSocket, WebSocketDisconnect
@@ -11,10 +10,10 @@ import cv2
 import io
 from PIL import Image
 import requests
-from .models.loader import (
-    face_processor, text_analyzer, whisper_worker, 
-    voice_worker, llm_generator, fusion_engine, models_ready
-)
+
+# ✅ FIX: Import the MODULE, not the variable
+from . import models as models_module
+
 from .session.manager import validate_token, save_message, load_user_history
 from .session.summary import generate_session_summary
 from .auth.database import get_db_connection
@@ -25,13 +24,7 @@ from .processing.speech import process_speech_end
 AVATAR_API = get_avatar_api_url()
 
 async def websocket_endpoint(websocket: WebSocket):
-    """
-    Main WebSocket endpoint handler
-    
-    ✅ FIXED: Comprehensive error handling
-    ✅ FIXED: Logs every step
-    ✅ FIXED: Never crashes unexpectedly
-    """
+    """Main WebSocket endpoint handler"""
     
     print("\n" + "="*80, flush=True)
     print("[WebSocket] 🔌 NEW CONNECTION", flush=True)
@@ -54,42 +47,39 @@ async def websocket_endpoint(websocket: WebSocket):
     user_preferences = {"voice": "female", "language": "en"}
     
     try:
-        # ===== STEP 1: AUTHENTICATION =====
+        # ===== AUTHENTICATION =====
         print("[WebSocket] ⏳ Waiting for auth message...", flush=True)
         
         try:
             auth_msg = await asyncio.wait_for(websocket.receive_json(), timeout=10.0)
             print(f"[WebSocket] 📨 Auth message received: {auth_msg.get('type')}", flush=True)
         except asyncio.TimeoutError:
-            print("[WebSocket] ❌ Auth timeout (10s)", flush=True)
-            await websocket.close(code=1008, reason="Authentication timeout")
-            return
-        except Exception as e:
-            print(f"[WebSocket] ❌ Auth receive error: {e}", flush=True)
+            print("[WebSocket] ❌ Auth timeout", flush=True)
+            await websocket.close(code=1008, reason="Auth timeout")
             return
         
         if auth_msg.get("type") != "auth":
-            print(f"[WebSocket] ❌ Wrong message type: {auth_msg.get('type')}", flush=True)
-            await websocket.send_json({"type": "error", "message": "Authentication required"})
-            await websocket.close(code=1008, reason="Auth required")
+            print(f"[WebSocket] ❌ Wrong type: {auth_msg.get('type')}", flush=True)
+            await websocket.send_json({"type": "error", "message": "Auth required"})
+            await websocket.close(code=1008)
             return
         
         token = auth_msg.get("token")
         print(f"[WebSocket] 🔑 Validating token: {token[:10] if token else 'None'}...", flush=True)
         
         if not token:
-            print(f"[WebSocket] ❌ No token provided", flush=True)
-            await websocket.send_json({"type": "error", "message": "No token provided"})
-            await websocket.close(code=1008, reason="No token")
+            print("[WebSocket] ❌ No token", flush=True)
+            await websocket.send_json({"type": "error", "message": "No token"})
+            await websocket.close(code=1008)
             return
         
-        # Validate token
+        # Validate
         session_data = validate_token(token)
         
         if not session_data:
-            print(f"[WebSocket] ❌ Invalid token", flush=True)
-            await websocket.send_json({"type": "error", "message": "Invalid session - please login again"})
-            await websocket.close(code=1008, reason="Invalid token")
+            print("[WebSocket] ❌ Invalid token", flush=True)
+            await websocket.send_json({"type": "error", "message": "Invalid session"})
+            await websocket.close(code=1008)
             return
         
         session_id = session_data['session_id']
@@ -98,114 +88,101 @@ async def websocket_endpoint(websocket: WebSocket):
         
         print(f"[WebSocket] ✅ Token validated for user: {username} (session: {session_id})", flush=True)
         
-        # Get user summary
+        # Get summary
         try:
-            user_summary = session_data.get('summary') or validate_token(token).get('summary')
-            if not user_summary:
-                conn = get_db_connection()
-                cursor = conn.cursor()
-                cursor.execute("SELECT summary_text FROM user_summaries WHERE user_id = ?", (user_id,))
-                summary_row = cursor.fetchone()
-                user_summary = summary_row[0] if summary_row else None
-                conn.close()
+            conn = get_db_connection()
+            cursor = conn.cursor()
+            cursor.execute("SELECT summary_text FROM user_summaries WHERE user_id = ?", (user_id,))
+            summary_row = cursor.fetchone()
+            user_summary = summary_row[0] if summary_row else None
+            conn.close()
         except Exception as e:
-            print(f"[WebSocket] ⚠️ Error loading summary: {e}", flush=True)
+            print(f"[WebSocket] ⚠️ Summary error: {e}", flush=True)
             user_summary = None
         
-        # Send authenticated confirmation
-        try:
-            await websocket.send_json({
-                "type": "authenticated",
-                "username": username,
-                "summary": user_summary
-            })
-            print(f"[WebSocket] ✅ Authenticated: {username} (user_id: {user_id})", flush=True)
-            if user_summary:
-                print(f"[WebSocket] 📖 Loaded summary: {user_summary[:60]}...", flush=True)
-        except Exception as e:
-            print(f"[WebSocket] ❌ Failed to send auth confirmation: {e}", flush=True)
-            return
+        # Send auth confirmation
+        await websocket.send_json({
+            "type": "authenticated",
+            "username": username,
+            "summary": user_summary
+        })
         
-        # Load conversation history
-        if llm_generator:
-            try:
-                llm_generator.clear_history()
-                user_history = load_user_history(user_id, limit=10)
-                
-                for role, content in user_history:
-                    llm_generator.conversation_history.append({
-                        "role": role,
-                        "content": content
-                    })
-                
-                if user_history:
-                    print(f"[WebSocket] 📚 Loaded {len(user_history)} messages from history", flush=True)
-            except Exception as e:
-                print(f"[WebSocket] ⚠️ Error loading history: {e}", flush=True)
+        print(f"[WebSocket] ✅ Authenticated: {username} (user_id: {user_id})", flush=True)
+        if user_summary:
+            print(f"[WebSocket] 📖 Loaded summary: {user_summary[:60]}...", flush=True)
         
-        # Wait for models if needed
-        if not models_ready:
-            print("[WebSocket] ⏳ Models still loading, asking client to wait...", flush=True)
+        # Load history
+        if models_module.llm_generator:
+            models_module.llm_generator.clear_history()
+            user_history = load_user_history(user_id, limit=10)
+            
+            for role, content in user_history:
+                models_module.llm_generator.conversation_history.append({
+                    "role": role,
+                    "content": content
+                })
+            
+            if user_history:
+                print(f"[WebSocket] 📚 Loaded {len(user_history)} messages", flush=True)
+        
+        # ✅ FIX: Check models_ready from the module, not imported variable
+        if not models_module.models_ready:
+            print(f"[WebSocket] ⏳ Models still loading (models_ready={models_module.models_ready}), asking client to wait...", flush=True)
             await websocket.send_json({
                 "type": "status",
-                "message": "AI models loading, please wait..."
+                "message": "Models loading..."
             })
             
             for i in range(900):
-                if models_ready:
+                if models_module.models_ready:
                     await websocket.send_json({
                         "type": "status",
-                        "message": "Models loaded! Ready."
+                        "message": "Models ready!"
                     })
+                    print(f"[WebSocket] ✅ Models became ready after {i}s", flush=True)
                     break
                 await asyncio.sleep(1)
             
-            if not models_ready:
-                print("[WebSocket] ❌ Models failed to load within 15 minutes", flush=True)
-                await websocket.send_json({
-                    "type": "error",
-                    "message": "Models failed to load"
-                })
-                await websocket.close(code=1011, reason="Model loading failed")
+            if not models_module.models_ready:
+                print("[WebSocket] ❌ Models timeout", flush=True)
+                await websocket.send_json({"type": "error", "message": "Models timeout"})
+                await websocket.close(code=1011)
                 return
         
-        print(f"[WebSocket] ✅ Models ready, entering message loop for {username}", flush=True)
+        print(f"[WebSocket] ✅ Models confirmed ready, entering message loop for {username}", flush=True)
         
-        # ===== STEP 2: MESSAGE LOOP =====
+        # ===== MESSAGE LOOP =====
         message_count = 0
         
         while True:
             try:
-                print(f"[WebSocket] ⏳ Waiting for message from {username}...", flush=True)
+                print(f"[WebSocket] ⏳ Waiting for message from {username} (msg #{message_count+1})...", flush=True)
+                
                 data = await websocket.receive_json()
                 message_count += 1
                 msg_type = data.get("type")
                 
-                print(f"[WebSocket] 📨 Message #{message_count} from {username}: type={msg_type}", flush=True)
+                print(f"[WebSocket] 📨 Message #{message_count}: type={msg_type}", flush=True)
                 
-                # ============ PREFERENCES UPDATE ============
+                # ===== PREFERENCES =====
                 if msg_type == "preferences":
-                    try:
-                        if "voice" in data:
-                            user_preferences["voice"] = data.get("voice", "female")
-                        if "language" in data:
-                            user_preferences["language"] = data.get("language", "en")
-                        print(f"[Preferences] ✅ {username}: voice={user_preferences.get('voice')}, lang={user_preferences.get('language')}", flush=True)
-                    except Exception as e:
-                        print(f"[Preferences] ❌ Error: {e}", flush=True)
+                    if "voice" in data:
+                        user_preferences["voice"] = data.get("voice", "female")
+                    if "language" in data:
+                        user_preferences["language"] = data.get("language", "en")
+                    print(f"[Preferences] ✅ Updated: {user_preferences}", flush=True)
                     continue
                 
-                # ============ GREETING REQUEST ============
+                # ===== GREETING =====
                 elif msg_type == "request_greeting":
-                    print(f"[Greeting] 🤖 Generating greeting for {username}...", flush=True)
+                    print(f"[Greeting] 🤖 Generating for {username}...", flush=True)
                     
                     try:
-                        # Determine greeting text
-                        language = user_preferences.get("language", "en")
+                        lang = user_preferences.get("language", "en")
                         if user_summary:
-                            greeting_text = GREETINGS[language]["returning"].format(username=username)
+                            greeting_text = GREETINGS[lang]["returning"].format(username=username)
                         else:
-                            greeting_text = GREETINGS[language]["new"].format(username=username)
+                            greeting_text = GREETINGS[lang]["new"].format(username=username)
                         
                         print(f"[Greeting] 👋 Text: '{greeting_text}'", flush=True)
                         
@@ -217,7 +194,7 @@ async def websocket_endpoint(websocket: WebSocket):
                             voice_pref = user_preferences.get("voice", "female")
                             lang_pref = user_preferences.get("language", "en")
                             
-                            print(f"[Greeting] 🔊 Requesting TTS (voice={voice_pref}, lang={lang_pref})...", flush=True)
+                            print(f"[Greeting] 🔊 Requesting TTS...", flush=True)
                             
                             avatar_response = requests.post(
                                 f"{AVATAR_API}/speak",
@@ -235,10 +212,10 @@ async def websocket_endpoint(websocket: WebSocket):
                                 visemes = avatar_data.get("visemes")
                                 print(f"[Greeting] ✅ TTS generated", flush=True)
                             else:
-                                print(f"[Greeting] ⚠️ TTS failed: {avatar_response.status_code}", flush=True)
+                                print(f"[Greeting] ⚠️ TTS status: {avatar_response.status_code}", flush=True)
                         
                         except requests.exceptions.ConnectionError:
-                            print(f"[Greeting] ⚠️ Avatar service not available (will send text-only)", flush=True)
+                            print(f"[Greeting] ⚠️ Avatar offline - text-only", flush=True)
                         except Exception as tts_err:
                             print(f"[Greeting] ⚠️ TTS error: {tts_err}", flush=True)
                         
@@ -257,39 +234,29 @@ async def websocket_endpoint(websocket: WebSocket):
                         else:
                             response_data["text_only"] = True
                         
-                        # Send greeting
-                        print(f"[Greeting] 📤 Sending response to {username}...", flush=True)
+                        print(f"[Greeting] 📤 Sending to {username}...", flush=True)
                         await websocket.send_json(response_data)
-                        print(f"[Greeting] ✅ Sent to {username}", flush=True)
+                        print(f"[Greeting] ✅ Sent", flush=True)
                         
-                        # Save to database
                         save_message(session_id, "assistant", greeting_text, "Neutral")
                         
-                    except Exception as greeting_err:
-                        print(f"[Greeting] ❌ Error: {greeting_err}", flush=True)
+                    except Exception as err:
+                        print(f"[Greeting] ❌ Error: {err}", flush=True)
                         import traceback
                         traceback.print_exc()
-                        
-                        try:
-                            await websocket.send_json({
-                                "type": "error",
-                                "message": f"Greeting failed: {str(greeting_err)}"
-                            })
-                        except:
-                            pass
                 
-                # ============ VIDEO FRAME ============
+                # ===== VIDEO FRAME =====
                 elif msg_type == "video_frame":
                     try:
                         img_data = base64.b64decode(data["frame"].split(",")[1])
                         img = Image.open(io.BytesIO(img_data))
                         frame = cv2.cvtColor(np.array(img), cv2.COLOR_RGB2BGR)
                         
-                        processed_frame, result = face_processor.process_frame(frame)
-                        face_emotion = face_processor.get_last_emotion() or "Neutral"
-                        face_confidence = face_processor.get_last_confidence() or 0.0
-                        face_probs = face_processor.get_last_probs()
-                        face_quality = getattr(face_processor, 'get_last_quality', lambda: 0.5)()
+                        processed_frame, result = models_module.face_processor.process_frame(frame)
+                        face_emotion = models_module.face_processor.get_last_emotion() or "Neutral"
+                        face_confidence = models_module.face_processor.get_last_confidence() or 0.0
+                        face_probs = models_module.face_processor.get_last_probs()
+                        face_quality = getattr(models_module.face_processor, 'get_last_quality', lambda: 0.5)()
                         
                         await websocket.send_json({
                             "type": "face_emotion",
@@ -301,23 +268,20 @@ async def websocket_endpoint(websocket: WebSocket):
                     except Exception as e:
                         print(f"[Video] ❌ Error: {e}", flush=True)
                 
-                # ============ AUDIO CHUNK ============
+                # ===== AUDIO CHUNK =====
                 elif msg_type == "audio_chunk":
                     try:
                         audio_data = base64.b64decode(data["audio"])
-                        
-                        # Convert to float32
                         audio_array = np.frombuffer(audio_data, dtype=np.int16).astype(np.float32) / 32768.0
                         
-                        # Feed to whisper
-                        if whisper_worker:
-                            whisper_worker.add_audio(audio_array)
+                        if models_module.whisper_worker:
+                            models_module.whisper_worker.add_audio(audio_array)
                         
                         audio_buffer.append(audio_array)
                         
                         if len(audio_buffer) >= 5:
-                            if voice_worker:
-                                voice_probs, voice_emotion = voice_worker.get_probs()
+                            if models_module.voice_worker:
+                                voice_probs, voice_emotion = models_module.voice_worker.get_probs()
                                 await websocket.send_json({
                                     "type": "voice_emotion",
                                     "emotion": voice_emotion
@@ -327,10 +291,10 @@ async def websocket_endpoint(websocket: WebSocket):
                     except Exception as e:
                         print(f"[Audio] ❌ Error: {e}", flush=True)
                 
-                # ============ SPEECH END ============
+                # ===== SPEECH END =====
                 elif msg_type == "speech_end":
                     transcription = data.get("text", "").strip()
-                    print(f"\n[Speech] 🎤 User finished: '{transcription}'", flush=True)
+                    print(f"\n[Speech] 🎤 '{transcription}'", flush=True)
                     
                     try:
                         await process_speech_end(
@@ -338,60 +302,44 @@ async def websocket_endpoint(websocket: WebSocket):
                             username, user_summary, user_preferences
                         )
                     except Exception as e:
-                        print(f"[Speech] ❌ Processing error: {e}", flush=True)
+                        print(f"[Speech] ❌ Error: {e}", flush=True)
                         import traceback
                         traceback.print_exc()
                         
-                        try:
-                            await websocket.send_json({
-                                "type": "error",
-                                "message": f"Processing failed: {str(e)}"
-                            })
-                        except:
-                            pass
+                        await websocket.send_json({
+                            "type": "error",
+                            "message": f"Error: {str(e)}"
+                        })
                 
                 else:
-                    print(f"[WebSocket] ⚠️ Unknown message type: {msg_type}", flush=True)
+                    print(f"[WebSocket] ⚠️ Unknown type: {msg_type}", flush=True)
             
             except WebSocketDisconnect:
-                print(f"[WebSocket] 🔌 {username} disconnected gracefully", flush=True)
+                print(f"[WebSocket] 🔌 {username} disconnected", flush=True)
                 break
             
-            except Exception as loop_err:
-                print(f"[WebSocket] ⚠️ Error in message loop: {loop_err}", flush=True)
+            except Exception as err:
+                print(f"[WebSocket] ⚠️ Loop error: {err}", flush=True)
                 import traceback
                 traceback.print_exc()
-                
-                try:
-                    await websocket.send_json({
-                        "type": "error",
-                        "message": f"Server error: {str(loop_err)}"
-                    })
-                except:
-                    pass
-                
-                # Don't break - continue receiving messages
                 await asyncio.sleep(0.1)
     
     except WebSocketDisconnect:
         print(f"[WebSocket] 🔌 {username or 'User'} disconnected", flush=True)
     
     except Exception as outer_err:
-        print(f"[WebSocket] ❌ Fatal error: {outer_err}", flush=True)
+        print(f"[WebSocket] ❌ Fatal: {outer_err}", flush=True)
         import traceback
         traceback.print_exc()
     
     finally:
-        # Generate summary on disconnect
         if session_id and user_id and username:
             print(f"[WebSocket] 📝 Generating summary for {username}...", flush=True)
             try:
                 summary = await generate_session_summary(session_id, user_id)
                 if summary:
-                    print(f"[Summary] ✅ Saved for {username}", flush=True)
-                else:
-                    print(f"[Summary] ⏭️ Skipped (not enough messages)", flush=True)
-            except Exception as summary_err:
-                print(f"[Summary] ❌ Error: {summary_err}", flush=True)
+                    print(f"[Summary] ✅ Saved", flush=True)
+            except Exception as e:
+                print(f"[Summary] ❌ Error: {e}", flush=True)
         
-        print(f"[WebSocket] 👋 Connection fully closed for {username or 'Unknown'}", flush=True)
+        print(f"[WebSocket] 👋 Closed for {username or 'Unknown'}", flush=True)