Spaces:
Running
Running
| from fastapi import APIRouter, WebSocket, WebSocketDisconnect | |
| from typing import List | |
| import asyncio | |
| import random | |
| import logging | |
| logger = logging.getLogger(__name__) | |
| router = APIRouter() | |
| class ConnectionManager: | |
| def __init__(self): | |
| self.active_connections: List[WebSocket] = [] | |
| async def connect(self, websocket: WebSocket): | |
| await websocket.accept() | |
| self.active_connections.append(websocket) | |
| logger.info(f"WebSocket connected. Total clients: {len(self.active_connections)}") | |
| def disconnect(self, websocket: WebSocket): | |
| if websocket in self.active_connections: | |
| self.active_connections.remove(websocket) | |
| logger.info(f"WebSocket disconnected. Total clients: {len(self.active_connections)}") | |
| async def send_personal_message(self, message: str, websocket: WebSocket): | |
| try: | |
| await websocket.send_text(message) | |
| except Exception as e: | |
| logger.error(f"Error sending message to ws: {e}") | |
| async def broadcast(self, message: str): | |
| for connection in self.active_connections[:]: | |
| try: | |
| await connection.send_text(message) | |
| except Exception: | |
| self.disconnect(connection) | |
| async def broadcast_json(self, data: dict): | |
| for connection in self.active_connections[:]: | |
| try: | |
| await connection.send_json(data) | |
| except Exception: | |
| self.disconnect(connection) | |
| manager = ConnectionManager() | |
| async def websocket_endpoint(websocket: WebSocket): | |
| await manager.connect(websocket) | |
| try: | |
| while True: | |
| data = await websocket.receive_text() | |
| # Echo back for ping/pong or process command | |
| await manager.send_personal_message(f'{{"type": "ACK", "message": "Received"}}', websocket) | |
| except WebSocketDisconnect: | |
| manager.disconnect(websocket) | |
| except Exception as e: | |
| logger.error(f"WebSocket error: {e}") | |
| manager.disconnect(websocket) | |
| async def push_realtime_updates(): | |
| """Background task to push live updates to all connected WS clients""" | |
| while True: | |
| try: | |
| await asyncio.sleep(5) # Update interval | |
| if len(manager.active_connections) > 0: | |
| # Simulated live dashboard metrics (in a real app, this would query a DB/Redis) | |
| live_data = { | |
| "type": "REALTIME_METRICS", | |
| "payload": { | |
| "active_users": random.randint(120, 150), | |
| "cpu_load": round(random.uniform(25.0, 65.0), 1), | |
| "memory_usage": round(random.uniform(40.0, 80.0), 1), | |
| "queries_per_sec": random.randint(10, 50), | |
| "anomaly_score": round(random.uniform(0.01, 0.05), 3) | |
| } | |
| } | |
| await manager.broadcast_json(live_data) | |
| except Exception as e: | |
| logger.error(f"Error in push_realtime_updates: {e}") | |
| await asyncio.sleep(5) | |