Spaces:
Running
Running
File size: 876 Bytes
2ecc4a7 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 | from fastapi import WebSocket
from typing import Dict
import json
from datetime import datetime
class WSManager:
def __init__(self):
# key: f"{user_id}:{job_id}"
self._connections: Dict[str, WebSocket] = {}
async def connect(self, job_id: str, websocket: WebSocket, user_id: str):
await websocket.accept()
key = f"{user_id}:{job_id}"
self._connections[key] = websocket
async def disconnect(self, job_id: str, user_id: str):
key = f"{user_id}:{job_id}"
if key in self._connections:
del self._connections[key]
async def emit(self, job_id: str, user_id: str, event: dict):
key = f"{user_id}:{job_id}"
ws = self._connections.get(key)
if ws:
event["timestamp"] = datetime.utcnow().isoformat()
await ws.send_json(event)
ws_manager = WSManager()
|