Files
sound-analyze/services/collector/ws_app.py

50 lines
1.4 KiB
Python

#!/usr/bin/env python3
"""FastAPI WebSocket endpoint with rate limiting support."""
from __future__ import annotations
import logging
from fastapi import FastAPI, WebSocket, WebSocketDisconnect, Query
from ws_manager import ConnectionManager
logger = logging.getLogger(__name__)
app = FastAPI(title="Audio Analyzer WebSocket")
manager = ConnectionManager()
@app.get("/health")
async def health() -> dict[str, str]:
"""Health check endpoint."""
return {"status": "ok"}
@app.websocket("/ws/live")
async def ws_live(
websocket: WebSocket,
hz: float = Query(default=10.0, ge=0.1, le=100.0, description="Update rate in Hz"),
) -> None:
"""
WebSocket endpoint for real-time audio data streaming.
Query parameters:
hz: Update rate in Hz (0.1 - 100.0, default: 10.0)
Examples:
- ws://localhost:8001/ws/live?hz=10 → 10 messages/sec
- ws://localhost:8001/ws/live?hz=1 → 1 message/sec
- ws://localhost:8001/ws/live?hz=30 → 30 messages/sec
"""
await manager.connect(websocket, rate_hz=hz)
try:
while True:
# Keep connection alive; client can send pings
await websocket.receive_text()
except WebSocketDisconnect:
pass
except Exception as e:
logger.debug("WS connection error: %s", e)
finally:
await manager.disconnect(websocket)