Add comprehensive WebSocket error handling and message validation

- Add proper error handling for missing 'type' field in messages
- Validate message format and log detailed debugging information
- Handle connection and knowledge base requests properly
- Add graceful error responses for invalid messages
- This should fix the 'type' and 'text' KeyError messages in logs

enhanced_websocket_handler.py

@@ -43,6 +43,13 @@ async def handle_enhanced_websocket_connection(websocket: WebSocket):
         # Get initial connection data
         initial_data = await websocket.receive_json()
         
+        # Validate initial data
+        if not isinstance(initial_data, dict):
+            logger.warning(f"⚠️ Invalid initial data format: {type(initial_data)}")
+            initial_data = {}
+        
+        logger.info(f"📨 Initial connection data: {initial_data}")
+        
         # Extract user preferences
         if "preferences" in initial_data:
             session_data["user_preferences"].update(initial_data["preferences"])
@@ -88,36 +95,71 @@ async def handle_enhanced_websocket_connection(websocket: WebSocket):
             try:
                 data = await websocket.receive_json()
                 
-                if data["type"] == "text_message":
+                # Validate message format
+                if not isinstance(data, dict):
+                    logger.warning(f"⚠️ Invalid message format: {type(data)}")
+                    continue
+                
+                if "type" not in data:
+                    logger.warning(f"⚠️ Message missing 'type' field: {data}")
+                    continue
+                
+                message_type = data["type"]
+                logger.debug(f"📨 Received message type: {message_type}")
+                
+                if message_type == "text_message":
                     await handle_text_message(
                         websocket, data, session_data, 
                         use_hybrid, config, knowledge_base, graph
                     )
                 
-                elif data["type"] == "voice_message":
+                elif message_type == "voice_message":
                     await handle_voice_message(
                         websocket, data, session_data,
                         use_hybrid, config, knowledge_base, graph
                     )
                 
-                elif data["type"] == "preferences_update":
+                elif message_type == "preferences_update":
                     await handle_preferences_update(websocket, data, session_data)
                 
-                elif data["type"] == "get_voice_status":
+                elif message_type == "get_voice_status":
                     await websocket.send_json({
                         "type": "voice_status",
                         "data": voice_service.get_voice_status()
                     })
                 
-                elif data["type"] == "get_llm_status":
+                elif message_type == "get_llm_status":
                     await websocket.send_json({
                         "type": "llm_status", 
                         "data": hybrid_llm_service.get_provider_info()
                     })
+                
+                elif message_type == "connection":
+                    # Handle initial connection - already processed above
+                    logger.debug("📨 Connection message received (already processed)")
+                
+                elif message_type == "get_knowledge_bases":
+                    # Handle knowledge base request
+                    await websocket.send_json({
+                        "type": "knowledge_bases",
+                        "knowledge_bases": ["government_docs", "rajasthan_documents"]
+                    })
+                
+                else:
+                    logger.warning(f"⚠️ Unknown message type: {message_type}")
 
             except WebSocketDisconnect:
                 logger.info("🔌 WebSocket client disconnected.")
                 break
+            except Exception as e:
+                logger.error(f"❌ Error handling message: {e}")
+                try:
+                    await websocket.send_json({
+                        "type": "error",
+                        "message": f"An error occurred: {str(e)}"
+                    })
+                except:
+                    pass  # Connection might be closed
             except Exception as e:
                 logger.error(f"❌ Error handling message: {e}")
                 await websocket.send_json({