first commit

backend/app/core/websocket/__init__.py

@@ -0,0 +1,4 @@
+"""__init__.py for websocket package"""
+
+from app.core.websocket.manager import manager, ConnectionManager
+from app.core.websocket.broadcaster import broadcaster, DataBroadcaster

backend/app/core/websocket/broadcaster.py

@@ -0,0 +1,93 @@
+"""Data broadcaster for WebSocket connections"""
+
+import asyncio
+from datetime import datetime
+from typing import Dict, Any, Optional
+
+from app.core.websocket.manager import manager
+
+
+class DataBroadcaster:
+    """Periodically broadcasts data to connected WebSocket clients"""
+
+    def __init__(self):
+        self.running = False
+        self.tasks: Dict[str, asyncio.Task] = {}
+
+    async def get_dashboard_stats(self) -> Dict[str, Any]:
+        """Get dashboard statistics"""
+        return {
+            "total_datasources": 9,
+            "active_datasources": 8,
+            "tasks_today": 45,
+            "success_rate": 97.8,
+            "last_updated": datetime.utcnow().isoformat(),
+            "alerts": {"critical": 0, "warning": 2, "info": 5},
+        }
+
+    async def broadcast_stats(self, interval: int = 5):
+        """Broadcast dashboard stats periodically"""
+        while self.running:
+            try:
+                stats = await self.get_dashboard_stats()
+                await manager.broadcast(
+                    {
+                        "type": "data_frame",
+                        "channel": "dashboard",
+                        "timestamp": datetime.utcnow().isoformat(),
+                        "payload": {"stats": stats},
+                    },
+                    channel="dashboard",
+                )
+            except Exception:
+                pass
+            await asyncio.sleep(interval)
+
+    async def broadcast_alert(self, alert: Dict[str, Any]):
+        """Broadcast an alert to all connected clients"""
+        await manager.broadcast(
+            {
+                "type": "alert_notification",
+                "timestamp": datetime.utcnow().isoformat(),
+                "data": {"alert": alert},
+            }
+        )
+
+    async def broadcast_gpu_update(self, data: Dict[str, Any]):
+        """Broadcast GPU cluster update"""
+        await manager.broadcast(
+            {
+                "type": "data_frame",
+                "channel": "gpu_clusters",
+                "timestamp": datetime.utcnow().isoformat(),
+                "payload": data,
+            }
+        )
+
+    async def broadcast_custom(self, channel: str, data: Dict[str, Any]):
+        """Broadcast custom data to a specific channel"""
+        await manager.broadcast(
+            {
+                "type": "data_frame",
+                "channel": channel,
+                "timestamp": datetime.utcnow().isoformat(),
+                "payload": data,
+            },
+            channel=channel if channel in manager.active_connections else "all",
+        )
+
+    def start(self):
+        """Start all broadcasters"""
+        if not self.running:
+            self.running = True
+            self.tasks["dashboard"] = asyncio.create_task(self.broadcast_stats(5))
+
+    def stop(self):
+        """Stop all broadcasters"""
+        self.running = False
+        for task in self.tasks.values():
+            task.cancel()
+        self.tasks.clear()
+
+
+broadcaster = DataBroadcaster()

backend/app/core/websocket/manager.py

@@ -0,0 +1,70 @@
+"""WebSocket Connection Manager"""
+
+import json
+import asyncio
+from typing import Dict, Set, Optional
+from datetime import datetime
+from fastapi import WebSocket
+import redis.asyncio as redis
+
+from app.core.config import settings
+
+
+class ConnectionManager:
+    """Manages WebSocket connections"""
+
+    def __init__(self):
+        self.active_connections: Dict[str, Set[WebSocket]] = {}  # user_id -> connections
+        self.redis_client: Optional[redis.Redis] = None
+
+    async def connect(self, websocket: WebSocket, user_id: str):
+        await websocket.accept()
+        if user_id not in self.active_connections:
+            self.active_connections[user_id] = set()
+        self.active_connections[user_id].add(websocket)
+
+        if self.redis_client is None:
+            redis_url = settings.REDIS_URL
+            if redis_url.startswith("redis://"):
+                self.redis_client = redis.from_url(redis_url, decode_responses=True)
+            else:
+                self.redis_client = redis.Redis(
+                    host=settings.REDIS_SERVER,
+                    port=settings.REDIS_PORT,
+                    db=settings.REDIS_DB,
+                    decode_responses=True,
+                )
+
+    def disconnect(self, websocket: WebSocket, user_id: str):
+        if user_id in self.active_connections:
+            self.active_connections[user_id].discard(websocket)
+            if not self.active_connections[user_id]:
+                del self.active_connections[user_id]
+
+    async def send_personal_message(self, message: dict, user_id: str):
+        if user_id in self.active_connections:
+            for connection in self.active_connections[user_id]:
+                try:
+                    await connection.send_json(message)
+                except Exception:
+                    pass
+
+    async def broadcast(self, message: dict, channel: str = "all"):
+        if channel == "all":
+            for user_id in self.active_connections:
+                await self.send_personal_message(message, user_id)
+        else:
+            await self.send_personal_message(message, channel)
+
+    async def close_all(self):
+        for user_id in self.active_connections:
+            for connection in self.active_connections[user_id]:
+                await connection.close()
+        self.active_connections.clear()
+
+
+manager = ConnectionManager()
+
+
+async def get_websocket_manager() -> ConnectionManager:
+    return manager