Spaces:
Runtime error
Runtime error
| """ | |
| WebSocket Connection Manager for AI Interviews | |
| Handles WebSocket connections, sessions, and message routing | |
| """ | |
| from fastapi import WebSocket, WebSocketDisconnect | |
| from typing import Dict, Optional, Callable | |
| from datetime import datetime | |
| import json | |
| import asyncio | |
| import logging | |
| logger = logging.getLogger(__name__) | |
| class ConnectionManager: | |
| """ | |
| Manages WebSocket connections and sessions for AI interviews | |
| """ | |
| def __init__(self): | |
| # Active WebSocket connections: session_id -> WebSocket | |
| self.active_connections: Dict[str, WebSocket] = {} | |
| # Session metadata: session_id -> session_data | |
| self.sessions: Dict[str, dict] = {} | |
| # Message handlers: message_type -> handler_function | |
| self.message_handlers: Dict[str, Callable] = {} | |
| logger.info("WebSocket ConnectionManager initialized") | |
| async def connect(self, session_id: str, websocket: WebSocket, metadata: Optional[dict] = None): | |
| """ | |
| Accept and register a new WebSocket connection | |
| """ | |
| await websocket.accept() | |
| self.active_connections[session_id] = websocket | |
| # Initialize session data | |
| self.sessions[session_id] = { | |
| "status": "connected", | |
| "connected_at": datetime.utcnow().isoformat(), | |
| "metadata": metadata or {}, | |
| "message_count": 0 | |
| } | |
| logger.info(f"✅ WebSocket connected: session={session_id}") | |
| # Send connection confirmation | |
| await self.send_message(session_id, { | |
| "type": "connection", | |
| "status": "connected", | |
| "session_id": session_id, | |
| "timestamp": datetime.utcnow().isoformat() | |
| }) | |
| def disconnect(self, session_id: str): | |
| """ | |
| Remove a WebSocket connection and mark session as disconnected | |
| """ | |
| if session_id in self.active_connections: | |
| del self.active_connections[session_id] | |
| logger.info(f"❌ WebSocket disconnected: session={session_id}") | |
| if session_id in self.sessions: | |
| self.sessions[session_id]["status"] = "disconnected" | |
| self.sessions[session_id]["disconnected_at"] = datetime.utcnow().isoformat() | |
| def is_connected(self, session_id: str) -> bool: | |
| """ | |
| Check if a session is currently connected | |
| """ | |
| return session_id in self.active_connections | |
| def get_session(self, session_id: str) -> Optional[dict]: | |
| """ | |
| Get session metadata | |
| """ | |
| return self.sessions.get(session_id) | |
| def update_session(self, session_id: str, data: dict): | |
| """ | |
| Update session metadata | |
| """ | |
| if session_id in self.sessions: | |
| self.sessions[session_id].update(data) | |
| async def send_message(self, session_id: str, message: dict): | |
| """ | |
| Send a message to a specific session | |
| """ | |
| if session_id in self.active_connections: | |
| try: | |
| await self.active_connections[session_id].send_json(message) | |
| # Update message count | |
| if session_id in self.sessions: | |
| self.sessions[session_id]["message_count"] += 1 | |
| logger.debug(f"📤 Message sent to session={session_id}, type={message.get('type')}") | |
| except Exception as e: | |
| logger.error(f"Error sending message to session={session_id}: {e}") | |
| self.disconnect(session_id) | |
| else: | |
| logger.warning(f"Cannot send message: session={session_id} not connected") | |
| async def send_text(self, session_id: str, text: str): | |
| """ | |
| Send a text message to a specific session | |
| """ | |
| if session_id in self.active_connections: | |
| try: | |
| await self.active_connections[session_id].send_text(text) | |
| except Exception as e: | |
| logger.error(f"Error sending text to session={session_id}: {e}") | |
| self.disconnect(session_id) | |
| async def send_bytes(self, session_id: str, data: bytes): | |
| """ | |
| Send binary data to a specific session | |
| """ | |
| if session_id in self.active_connections: | |
| try: | |
| await self.active_connections[session_id].send_bytes(data) | |
| except Exception as e: | |
| logger.error(f"Error sending bytes to session={session_id}: {e}") | |
| self.disconnect(session_id) | |
| async def broadcast(self, message: dict, exclude: Optional[list] = None): | |
| """ | |
| Broadcast a message to all connected sessions | |
| """ | |
| exclude = exclude or [] | |
| disconnected = [] | |
| for session_id, connection in self.active_connections.items(): | |
| if session_id not in exclude: | |
| try: | |
| await connection.send_json(message) | |
| except Exception as e: | |
| logger.error(f"Error broadcasting to session={session_id}: {e}") | |
| disconnected.append(session_id) | |
| # Clean up disconnected sessions | |
| for session_id in disconnected: | |
| self.disconnect(session_id) | |
| def register_handler(self, message_type: str, handler: Callable): | |
| """ | |
| Register a message handler for a specific message type | |
| """ | |
| self.message_handlers[message_type] = handler | |
| logger.info(f"Registered handler for message type: {message_type}") | |
| async def handle_message(self, session_id: str, message: dict): | |
| """ | |
| Route message to appropriate handler | |
| """ | |
| message_type = message.get("type") | |
| if message_type in self.message_handlers: | |
| try: | |
| await self.message_handlers[message_type](session_id, message) | |
| except Exception as e: | |
| logger.error(f"Error handling message type={message_type}: {e}") | |
| await self.send_message(session_id, { | |
| "type": "error", | |
| "message": f"Error processing {message_type}: {str(e)}" | |
| }) | |
| else: | |
| logger.warning(f"No handler for message type: {message_type}") | |
| await self.send_message(session_id, { | |
| "type": "error", | |
| "message": f"Unknown message type: {message_type}" | |
| }) | |
| def get_active_sessions(self) -> list: | |
| """ | |
| Get list of all active session IDs | |
| """ | |
| return list(self.active_connections.keys()) | |
| def get_session_count(self) -> int: | |
| """ | |
| Get count of active sessions | |
| """ | |
| return len(self.active_connections) | |
| def cleanup_session(self, session_id: str): | |
| """ | |
| Clean up session data | |
| """ | |
| if session_id in self.sessions: | |
| del self.sessions[session_id] | |
| logger.info(f"🧹 Session cleaned up: session={session_id}") | |
| # Global connection manager instance | |
| manager = ConnectionManager() | |
| # Utility functions | |
| async def send_error(session_id: str, error_message: str): | |
| """ | |
| Send an error message to a session | |
| """ | |
| await manager.send_message(session_id, { | |
| "type": "error", | |
| "message": error_message, | |
| "timestamp": datetime.utcnow().isoformat() | |
| }) | |
| async def send_status(session_id: str, status: str, details: Optional[dict] = None): | |
| """ | |
| Send a status update to a session | |
| """ | |
| message = { | |
| "type": "status", | |
| "status": status, | |
| "timestamp": datetime.utcnow().isoformat() | |
| } | |
| if details: | |
| message.update(details) | |
| await manager.send_message(session_id, message) | |