app.py

"""
app.py
------
FastAPI server for the open-source streaming voice agent.

All project files live in ONE directory — no subdirectories.

Endpoints
---------
GET  /                         → index.html
GET  /app.js                   → app.js
GET  /audio-capture-worklet.js → audio-capture-worklet.js
GET  /health                   → JSON status
WS   /ws/audio                 → bidirectional audio stream

WebSocket protocol
------------------
Client → Server:
  • Binary frames : PCM int16, mono, 16 kHz, 640-byte chunks (20 ms)
  • Text  frames  : JSON control  {"type": "start"} | {"type": "stop"}

Server → Client:
  • Binary frames : PCM int16, mono, 16 kHz (TTS output, variable size)
  • Text  frames  : JSON status
        {"type": "status",       "message": "..."}
        {"type": "transcript",   "text": "..."}
        {"type": "agent_start"}
        {"type": "agent_done",   "text": "...", "latency_ms": 000}
        {"type": "error",        "message": "..."}
        {"type": "interrupted"}
"""

from __future__ import annotations

import asyncio
import json
import logging
import os
import time
from pathlib import Path

import numpy as np
import torch
import uvicorn
from fastapi import FastAPI, WebSocket, WebSocketDisconnect
from fastapi.responses import FileResponse, JSONResponse

from models import STT_SR, ModelManager
from pipeline import StreamingPipeline, build_initial_conversation

# --------------------------------------------------------------------------- #
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s  %(levelname)-8s  %(name)s: %(message)s",
    datefmt="%H:%M:%S",
)
logger = logging.getLogger(__name__)

BASE_DIR = Path(__file__).parent   # All files live here

# --------------------------------------------------------------------------- #
app = FastAPI(title="Open-Source Voice Agent", version="1.0.0")

_models: ModelManager | None = None
_pipeline: StreamingPipeline | None = None


@app.on_event("startup")
async def on_startup():
    global _models, _pipeline
    logger.info("=== Loading models — this may take a few minutes on first run ===")
    loop = asyncio.get_event_loop()
    _models = await loop.run_in_executor(None, ModelManager)
    await loop.run_in_executor(None, _models.warm_up)
    _pipeline = StreamingPipeline(_models)
    logger.info("=== Server ready ===")


# --------------------------------------------------------------------------- #
#  Static file routes (flat — all files in same directory as app.py)          #
# --------------------------------------------------------------------------- #
@app.get("/")
async def root():
    return FileResponse(str(BASE_DIR / "index.html"))

@app.get("/app.js")
async def serve_appjs():
    return FileResponse(str(BASE_DIR / "app.js"), media_type="application/javascript")

@app.get("/audio-capture-worklet.js")
async def serve_worklet():
    return FileResponse(str(BASE_DIR / "audio-capture-worklet.js"), media_type="application/javascript")

@app.get("/health")
async def health():
    return JSONResponse({
        "status": "ready" if _models else "loading",
        "device": str(_models.device) if _models else "N/A",
        "stt": "whisper-base.en",
        "llm": "SmolLM2-1.7B-Instruct",
        "tts": "mms-tts-eng",
    })


# --------------------------------------------------------------------------- #
#  VAD-based speech segmenter                                                  #
# --------------------------------------------------------------------------- #
class SpeechSegmenter:
    """
    Wraps Silero VADIterator to accumulate PCM and fire when speech ends.

    Incoming audio is 640-byte chunks (320 int16 samples = 20 ms at 16 kHz).
    Silero needs exactly 512 float32 samples per call, so we buffer the
    difference and carry any leftover into the next iteration.
    """

    VAD_CHUNK = 512
    MIN_SPEECH_S = 0.30

    def __init__(self, models: ModelManager):
        self._vad = models.vad_iter
        self._vad.reset_states()
        self._pre_roll: list[np.ndarray] = []
        self._speech_buf: list[np.ndarray] = []
        self._vad_leftover = np.array([], dtype=np.float32)
        self._speaking = False

    def feed(self, pcm_bytes: bytes) -> np.ndarray | None:
        audio = (
            np.frombuffer(pcm_bytes, dtype=np.int16).astype(np.float32) / 32_768.0
        )
        self._pre_roll.append(audio)
        if len(self._pre_roll) > 5:
            self._pre_roll.pop(0)

        if self._speaking:
            self._speech_buf.append(audio)

        combined = np.concatenate([self._vad_leftover, audio])
        i = 0
        result = None

        while i + self.VAD_CHUNK <= len(combined):
            chunk = combined[i : i + self.VAD_CHUNK]
            i += self.VAD_CHUNK
            try:
                event = self._vad(
                    torch.from_numpy(chunk).float(), return_seconds=False
                )
            except Exception:
                event = None

            if event:
                if "start" in event and not self._speaking:
                    self._speaking = True
                    self._speech_buf = list(self._pre_roll)
                elif "end" in event and self._speaking:
                    self._speaking = False
                    if self._speech_buf:
                        speech = np.concatenate(self._speech_buf)
                        if len(speech) / STT_SR >= self.MIN_SPEECH_S:
                            result = speech
                    self._speech_buf = []

        self._vad_leftover = combined[i:]
        return result

    def reset(self):
        self._pre_roll = []
        self._speech_buf = []
        self._vad_leftover = np.array([], dtype=np.float32)
        self._speaking = False
        self._vad.reset_states()


# --------------------------------------------------------------------------- #
#  WebSocket handler                                                           #
# --------------------------------------------------------------------------- #
@app.websocket("/ws/audio")
async def ws_audio(ws: WebSocket):
    await ws.accept()
    logger.info("Client connected: %s", ws.client)

    if _models is None or _pipeline is None:
        await ws.send_text(json.dumps({
            "type": "error", "message": "Models still loading, please wait."
        }))
        await ws.close()
        return

    loop = asyncio.get_event_loop()
    segmenter = SpeechSegmenter(_models)
    conversation = build_initial_conversation()

    async def _send(obj: dict):
        await ws.send_text(json.dumps(obj))

    try:
        while True:
            msg = await ws.receive()

            # ── control frame ─────────────────────────────────────────── #
            if "text" in msg and msg["text"]:
                try:
                    ctrl = json.loads(msg["text"])
                except json.JSONDecodeError:
                    continue
                if ctrl.get("type") == "start":
                    segmenter.reset()
                    _models.vad_reset()
                    conversation = build_initial_conversation()
                    logger.info("Session reset by client.")
                    await _send({"type": "status", "message": "Session started."})
                elif ctrl.get("type") == "stop":
                    break
                continue

            # ── audio frame ───────────────────────────────────────────── #
            if "bytes" not in msg or not msg["bytes"]:
                continue

            speech_audio = segmenter.feed(msg["bytes"])
            if speech_audio is None:
                continue

            # ── transcribe ────────────────────────────────────────────── #
            logger.info("Speech: %.2f s → transcribing…", len(speech_audio) / STT_SR)
            await _send({"type": "status", "message": "Transcribing…"})

            transcript: str = await loop.run_in_executor(
                None, _models.transcribe, speech_audio
            )

            if not transcript or len(transcript.strip()) < 2:
                await _send({"type": "status", "message": "Couldn't hear clearly. Try again."})
                continue

            logger.info("Transcript: '%s'", transcript)
            await _send({"type": "transcript", "text": transcript})

            # ── LLM + TTS streaming pipeline ──────────────────────────── #
            conversation.append({"role": "user", "content": transcript})
            await _send({"type": "agent_start"})

            t0 = time.perf_counter()
            try:
                response_text = await _pipeline.process(conversation, ws, loop)
            except Exception as exc:
                logger.error("Pipeline error: %s", exc, exc_info=True)
                await _send({"type": "error", "message": f"Pipeline error: {exc}"})
                continue

            latency_ms = int((time.perf_counter() - t0) * 1000)
            conversation.append({"role": "assistant", "content": response_text})
            await _send({"type": "agent_done", "text": response_text, "latency_ms": latency_ms})
            logger.info("Turn done in %d ms.", latency_ms)

    except WebSocketDisconnect:
        logger.info("Client disconnected: %s", ws.client)
    except Exception as exc:
        logger.error("WS error: %s", exc, exc_info=True)
        try:
            await _send({"type": "error", "message": str(exc)})
        except Exception:
            pass
    finally:
        logger.info("WS session closed: %s", ws.client)


# --------------------------------------------------------------------------- #
if __name__ == "__main__":
    port = int(os.environ.get("PORT", 7860))
    uvicorn.run(
        "app:app",
        host="0.0.0.0",
        port=port,
        log_level="info",
        ws_ping_interval=30,
        ws_ping_timeout=60,
    )