atlas/home-voice-agent/wake-word/server.py

#!/usr/bin/env python3
"""
Wake-word detection HTTP/WebSocket server.

Provides endpoints for:
- Starting/stopping wake-word detection
- Receiving wake-word detection events
- Health check
"""

import logging
import json
import asyncio
from typing import Set
from fastapi import FastAPI, WebSocket, WebSocketDisconnect, HTTPException
from fastapi.responses import JSONResponse
from pydantic import BaseModel

from .detector import WakeWordDetector

logger = logging.getLogger(__name__)

app = FastAPI(title="Wake-Word Service", version="0.1.0")

# Global state
detector: WakeWordDetector = None
websocket_clients: Set[WebSocket] = set()
is_detecting = False


class DetectionCallback:
    """Callback to notify WebSocket clients of wake-word detection."""
    
    @staticmethod
    def on_detection():
        """Called when wake-word is detected."""
        asyncio.create_task(DetectionCallback._broadcast_detection())
    
    @staticmethod
    async def _broadcast_detection():
        """Broadcast detection to all WebSocket clients."""
        message = json.dumps({
            "type": "wake_word_detected",
            "wake_word": "hey atlas",
            "timestamp": asyncio.get_event_loop().time()
        })
        
        disconnected = set()
        for client in websocket_clients:
            try:
                await client.send_text(message)
            except Exception as e:
                logger.error(f"Error sending to client: {e}")
                disconnected.add(client)
        
        # Remove disconnected clients
        websocket_clients.difference_update(disconnected)


@app.get("/health")
async def health():
    """Health check endpoint."""
    return {
        "status": "healthy",
        "service": "wake-word",
        "is_detecting": is_detecting,
        "clients": len(websocket_clients)
    }


@app.post("/start")
async def start_detection():
    """Start wake-word detection."""
    global detector, is_detecting
    
    if is_detecting:
        return {"status": "already_running", "message": "Wake-word detection already running"}
    
    try:
        detector = WakeWordDetector(
            wake_word="hey atlas",
            threshold=0.5,
            on_detection=DetectionCallback.on_detection
        )
        detector.start()
        is_detecting = True
        
        return {
            "status": "started",
            "message": "Wake-word detection started"
        }
    except Exception as e:
        logger.error(f"Error starting detection: {e}")
        raise HTTPException(status_code=500, detail=str(e))


@app.post("/stop")
async def stop_detection():
    """Stop wake-word detection."""
    global detector, is_detecting
    
    if not is_detecting:
        return {"status": "not_running", "message": "Wake-word detection not running"}
    
    try:
        if detector:
            detector.stop()
        is_detecting = False
        
        return {
            "status": "stopped",
            "message": "Wake-word detection stopped"
        }
    except Exception as e:
        logger.error(f"Error stopping detection: {e}")
        raise HTTPException(status_code=500, detail=str(e))


@app.get("/status")
async def get_status():
    """Get current detection status."""
    return {
        "is_detecting": is_detecting,
        "clients": len(websocket_clients)
    }


@app.websocket("/events")
async def websocket_events(websocket: WebSocket):
    """WebSocket endpoint for wake-word detection events."""
    await websocket.accept()
    websocket_clients.add(websocket)
    logger.info(f"WebSocket client connected. Total clients: {len(websocket_clients)}")
    
    try:
        # Send initial status
        await websocket.send_json({
            "type": "connected",
            "status": "listening"
        })
        
        # Keep connection alive
        while True:
            try:
                # Wait for client messages (ping/pong)
                data = await asyncio.wait_for(websocket.receive_text(), timeout=30.0)
                # Echo back or handle commands
                await websocket.send_json({"type": "pong", "data": data})
            except asyncio.TimeoutError:
                # Send keepalive
                await websocket.send_json({"type": "keepalive"})
            except WebSocketDisconnect:
                break
    except Exception as e:
        logger.error(f"WebSocket error: {e}")
    finally:
        websocket_clients.discard(websocket)
        logger.info(f"WebSocket client disconnected. Total clients: {len(websocket_clients)}")


if __name__ == "__main__":
    import uvicorn
    logging.basicConfig(level=logging.INFO)
    uvicorn.run(app, host="0.0.0.0", port=8002)