whisper-stt-cli/stt_service.py

#!/usr/bin/env python3
"""STT-Service – lokaler Spracherkennungs-Dienst (Pendant zu chatterbox-tts-cli).

Start:
    .venv/bin/uvicorn stt_service:app --host 0.0.0.0 --port 9998

Endpunkte:
    GET  /health                  – Status, Modell, Gerät, Warteschlange
    POST /v1/audio/transcriptions – OpenAI-kompatibel (multipart: file, language, prompt,
                                    response_format=json|text|verbose_json)
    POST /transcribe              – Rohdaten im Body (für Mikrocontroller, ohne multipart):
                                    WAV/MP3/… oder ?format=pcm&sample_rate=16000&channels=1
    WS   /ws/stream               – Live-Streaming: s16le-PCM rein, Äußerungsende per VAD,
                                    Transkript-Events raus
    TCP  :10300 (Wyoming)         – Home Assistant (STT_WYOMING_PORT, leer = aus)

Nicht blockierend: Die Inferenz läuft in genau einem Worker-Thread (GPU-schonend),
der Event-Loop bleibt frei. Ist die Warteschlange voll, wird sofort mit 503
abgelehnt statt Anfragen unbegrenzt zu stauen.

Konfiguration per Umgebung (siehe README): STT_MODEL, STT_DEVICE, STT_LANGUAGE,
STT_MAX_QUEUE, STT_VAD_*, STT_WYOMING_HOST/PORT, STT_PRELOAD …
"""

from __future__ import annotations

import asyncio
import json
import logging
import os
from concurrent.futures import ThreadPoolExecutor
from contextlib import asynccontextmanager
from dataclasses import dataclass
from functools import partial
from typing import Callable

import numpy as np
from fastapi import FastAPI, File, Form, HTTPException, Query, Request, UploadFile, WebSocket
from fastapi.responses import PlainTextResponse
from starlette.websockets import WebSocketDisconnect, WebSocketState

import stt_core as core

log = logging.getLogger("stt.service")

__version__ = core.__version__


class Overloaded(RuntimeError):
    """Warteschlange voll."""


@dataclass
class ServiceSettings:
    max_queue: int = 8
    max_upload_mb: int = 50
    preload: bool = True
    wyoming_host: str = "0.0.0.0"
    wyoming_port: int | None = 10300

    @classmethod
    def from_env(cls) -> "ServiceSettings":
        d = cls()
        port = os.environ.get("STT_WYOMING_PORT", str(d.wyoming_port))
        return cls(
            max_queue=core._env("STT_MAX_QUEUE", d.max_queue, int),
            max_upload_mb=core._env("STT_MAX_UPLOAD_MB", d.max_upload_mb, int),
            preload=core._env("STT_PRELOAD", d.preload, bool),
            wyoming_host=core._env("STT_WYOMING_HOST", d.wyoming_host),
            wyoming_port=int(port) if port.strip() else None,
        )


class TranscriptionWorker:
    """Serialisiert Inferenz in einem Thread; begrenzt die Anzahl wartender Aufträge."""

    def __init__(self, engine: core.SttEngine, max_queue: int = 8):
        self.engine = engine
        self.max_queue = max_queue
        self.pending = 0  # nur im Event-Loop verändert → kein Lock nötig
        self._executor = ThreadPoolExecutor(max_workers=1, thread_name_prefix="stt-worker")

    async def run(self, fn: Callable, *args, **kwargs):
        loop = asyncio.get_running_loop()
        return await loop.run_in_executor(self._executor, partial(fn, *args, **kwargs))

    async def transcribe(self, audio: np.ndarray, **kwargs) -> core.Transcript:
        if self.pending >= self.max_queue:
            raise Overloaded(f"Warteschlange voll ({self.pending} Aufträge)")
        self.pending += 1
        try:
            return await self.run(self.engine.transcribe, audio, **kwargs)
        finally:
            self.pending -= 1

    def shutdown(self) -> None:
        self._executor.shutdown(wait=False, cancel_futures=True)


def create_app(
    engine: core.SttEngine | None = None,
    settings: ServiceSettings | None = None,
    vad_config: core.VadConfig | None = None,
    prob_fn_factory: Callable[[], Callable] | None = None,
) -> FastAPI:
    """App-Fabrik; Tests injizieren eine Fake-Engine bzw. eine VAD-Funktion."""
    engine = engine or core.SttEngine(core.EngineConfig.from_env())
    settings = settings or ServiceSettings.from_env()
    vad_config = vad_config or core.VadConfig.from_env()
    prob_fn_factory = prob_fn_factory or core.SileroProb
    worker = TranscriptionWorker(engine, settings.max_queue)

    @asynccontextmanager
    async def lifespan(app: FastAPI):
        preload_future = None
        if settings.preload:
            # Nicht awaiten: Dienst ist sofort erreichbar (/health = loading);
            # erste Anfragen reihen sich im Worker hinter dem Laden ein.
            preload_future = asyncio.ensure_future(worker.run(engine.warmup))
            preload_future.add_done_callback(_log_preload_result)
        wyoming_server = None
        if settings.wyoming_port:
            wyoming_server = await _start_wyoming(worker, settings)
        app.state.wyoming = wyoming_server
        try:
            yield
        finally:
            if wyoming_server is not None:
                await wyoming_server.stop()
            if preload_future is not None and not preload_future.done():
                preload_future.cancel()
            worker.shutdown()

    app = FastAPI(title="STT-Service", version=__version__, lifespan=lifespan)
    app.state.engine = engine
    app.state.worker = worker

    # --- HTTP -----------------------------------------------------------------

    @app.get("/health")
    async def health():
        return {
            "status": "ok" if engine.status == "ready" else engine.status,
            "model": engine.active_model or engine.config.model,
            "device": engine.active_device,
            "compute_type": engine.active_compute_type,
            "default_language": engine.config.language or "auto",
            "queue": worker.pending,
            "max_queue": worker.max_queue,
            "load_seconds": engine.load_seconds,
            "wyoming_port": settings.wyoming_port,
            "error": engine.error,
            "version": __version__,
        }

    async def _run_transcription(audio: np.ndarray, **kwargs) -> core.Transcript:
        try:
            return await worker.transcribe(audio, **kwargs)
        except Overloaded as exc:
            raise HTTPException(status_code=503, detail=str(exc), headers={"Retry-After": "2"})
        except RuntimeError as exc:
            log.exception("Transkription fehlgeschlagen")
            raise HTTPException(status_code=500, detail=str(exc))

    def _check_size(size: int) -> None:
        if size > settings.max_upload_mb * 1024 * 1024:
            raise HTTPException(status_code=413, detail=f"Audio größer als {settings.max_upload_mb} MB")

    async def _decode(data: bytes) -> np.ndarray:
        try:
            # Dekodieren (PyAV) kostet bei langen Dateien spürbar CPU → nicht im Event-Loop
            return await asyncio.to_thread(core.decode_audio_bytes, data)
        except core.AudioDecodeError as exc:
            raise HTTPException(status_code=400, detail=str(exc))

    @app.post("/v1/audio/transcriptions")
    async def openai_transcriptions(
        file: UploadFile = File(...),
        model: str | None = Form(default=None),  # akzeptiert, aber ignoriert (ein Modell pro Dienst)
        language: str | None = Form(default=None),
        prompt: str | None = Form(default=None),
        response_format: str = Form(default="json"),
        temperature: float | None = Form(default=None),  # ignoriert
    ):
        data = await file.read()
        _check_size(len(data))
        audio = await _decode(data)
        tr = await _run_transcription(audio, language=language, initial_prompt=prompt)
        if response_format == "text":
            return PlainTextResponse(tr.text)
        if response_format == "verbose_json":
            return {
                "task": "transcribe",
                "language": tr.language,
                "duration": tr.duration,
                "text": tr.text,
                "segments": [
                    {"id": i, "start": s.start, "end": s.end, "text": s.text,
                     "avg_logprob": s.avg_logprob, "no_speech_prob": s.no_speech_prob}
                    for i, s in enumerate(tr.segments)
                ],
            }
        if response_format != "json":
            raise HTTPException(status_code=400, detail=f"response_format '{response_format}' nicht unterstützt")
        return {"text": tr.text}

    @app.post("/transcribe")
    async def transcribe_raw(
        request: Request,
        format: str = Query(default="auto", pattern="^(auto|pcm)$"),
        sample_rate: int = Query(default=16000, ge=4000, le=192000),
        channels: int = Query(default=1, ge=1, le=16),
        channel: int | None = Query(default=None, ge=0),
        language: str | None = Query(default=None),
        prompt: str | None = Query(default=None),
    ):
        data = await request.body()
        _check_size(len(data))
        if not data:
            raise HTTPException(status_code=400, detail="Leerer Body")
        if format == "pcm":
            try:
                audio = core.resample(core.pcm16_to_float(data, channels, channel), sample_rate)
            except ValueError as exc:
                raise HTTPException(status_code=400, detail=str(exc))
        else:
            audio = await _decode(data)
        tr = await _run_transcription(audio, language=language, initial_prompt=prompt)
        return tr.to_dict()

    # --- WebSocket-Streaming --------------------------------------------------

    @app.websocket("/ws/stream")
    async def ws_stream(ws: WebSocket):
        await ws.accept()
        session = _StreamSession(ws, worker, vad_config, prob_fn_factory)
        await session.run()

    return app


def _log_preload_result(future: asyncio.Future) -> None:
    if future.cancelled():
        return
    exc = future.exception()
    if exc is not None:
        log.error("Vorladen des Modells fehlgeschlagen: %s", exc)


class _StreamSession:
    """Eine WebSocket-Verbindung.

    Protokoll (Client → Server):
        Text  {"type":"start", "sample_rate":16000, "channels":1, "channel":null,
               "language":"de", "prompt":null, "vad":true,
               "silence_ms":700, "threshold":0.5, "max_utterance_s":30}   (optional)
        Binär s16le-PCM (beliebige Stückgröße)
        Text  {"type":"end"}   – laufende Äußerung abschließen (bzw. ohne VAD: alles
                                 bisher Gesendete transkribieren); danach "done"
        Text  {"type":"close"} – Verbindung nach ausstehenden Transkripten schließen

    Server → Client:
        {"type":"ready", ...}  {"type":"speech_start","t":1.23}
        {"type":"transcript","text":...,"language":...,"duration":...,"processing_s":...,
         "reason":"silence|max_length|flush","t":...}
        {"type":"done"}  {"type":"error","detail":...}
    """

    def __init__(self, ws: WebSocket, worker: TranscriptionWorker,
                 vad_config: core.VadConfig, prob_fn_factory: Callable):
        self.ws = ws
        self.worker = worker
        self.base_vad = vad_config
        self.prob_fn_factory = prob_fn_factory
        self.send_lock = asyncio.Lock()
        self.jobs: asyncio.Queue = asyncio.Queue()
        self._configure({})

    def _configure(self, opts: dict) -> None:
        self.sample_rate = int(opts.get("sample_rate", 16000))
        self.channels = int(opts.get("channels", 1))
        self.channel = opts.get("channel")
        self.language = opts.get("language")
        self.prompt = opts.get("prompt")
        self.use_vad = bool(opts.get("vad", True))
        vad = core.VadConfig(
            threshold=float(opts.get("threshold", self.base_vad.threshold)),
            min_speech_ms=int(opts.get("min_speech_ms", self.base_vad.min_speech_ms)),
            silence_ms=int(opts.get("silence_ms", self.base_vad.silence_ms)),
            pre_roll_ms=int(opts.get("pre_roll_ms", self.base_vad.pre_roll_ms)),
            max_utterance_s=float(opts.get("max_utterance_s", self.base_vad.max_utterance_s)),
        )
        self.resampler = core.StreamResampler(self.sample_rate)
        self.segmenter = core.UtteranceSegmenter(vad, self.prob_fn_factory()) if self.use_vad else None
        self.buffer: list[np.ndarray] = []
        self.buffered_samples = 0
        self.max_buffer_samples = int(vad.max_utterance_s * core.SAMPLE_RATE)

    async def send(self, payload: dict) -> None:
        if self.ws.application_state != WebSocketState.CONNECTED:
            return
        async with self.send_lock:
            await self.ws.send_json(payload)

    async def run(self) -> None:
        consumer = asyncio.create_task(self._consume())
        await self.send({"type": "ready", "sample_rate": self.sample_rate, "vad": self.use_vad})
        try:
            while True:
                message = await self.ws.receive()
                if message["type"] == "websocket.disconnect":
                    break
                if message.get("bytes") is not None:
                    await self._on_audio(message["bytes"])
                elif message.get("text") is not None:
                    if not await self._on_control(message["text"]):
                        break
        except WebSocketDisconnect:
            pass
        finally:
            await self.jobs.put(None)
            try:
                await asyncio.wait_for(consumer, timeout=120)
            except (asyncio.TimeoutError, asyncio.CancelledError):
                consumer.cancel()
            if self.ws.application_state == WebSocketState.CONNECTED:
                await self.ws.close()

    async def _on_audio(self, data: bytes) -> None:
        try:
            audio = core.pcm16_to_float(data, self.channels, self.channel)
        except ValueError as exc:
            await self.send({"type": "error", "detail": str(exc)})
            return
        audio = self.resampler.process(audio)
        if self.segmenter is None:
            self.buffer.append(audio)
            self.buffered_samples += audio.size
            if self.buffered_samples >= self.max_buffer_samples:
                await self._enqueue_buffer("max_length")
            return
        for event in self.segmenter.feed(audio):
            await self._handle_vad_event(event)

    async def _handle_vad_event(self, event: core.VadEvent) -> None:
        if event.kind == "speech_start":
            await self.send({"type": "speech_start", "t": round(event.time_s, 2)})
        else:
            await self.jobs.put((event.audio, event.reason, event.time_s))

    async def _enqueue_buffer(self, reason: str) -> None:
        if self.buffer:
            audio = np.concatenate(self.buffer)
            self.buffer, self.buffered_samples = [], 0
            await self.jobs.put((audio, reason, None))

    async def _on_control(self, raw: str) -> bool:
        try:
            msg = json.loads(raw)
        except ValueError:
            await self.send({"type": "error", "detail": "Ungültiger Steuer-Frame (kein JSON)"})
            return True
        mtype = msg.get("type")
        if mtype == "start":
            self._configure(msg)
            await self.send({"type": "ready", "sample_rate": self.sample_rate, "vad": self.use_vad})
        elif mtype in ("end", "close"):
            tail = self.resampler.process(np.zeros(0, np.float32), last=True)
            if self.segmenter is not None:
                events = self.segmenter.feed(tail) + self.segmenter.flush()
                for event in events:
                    await self._handle_vad_event(event)
            else:
                if tail.size:
                    self.buffer.append(tail)
                await self._enqueue_buffer("flush")
            self.resampler = core.StreamResampler(self.sample_rate)
            await self.jobs.put("done")
            return mtype == "end"
        else:
            await self.send({"type": "error", "detail": f"Unbekannter Frame-Typ: {mtype!r}"})
        return True

    async def _consume(self) -> None:
        """Arbeitet Äußerungen der Reihe nach ab; der Empfang läuft derweil weiter."""
        while True:
            job = await self.jobs.get()
            if job is None:
                return
            if job == "done":
                await self.send({"type": "done"})
                continue
            audio, reason, t = job
            try:
                tr = await self.worker.transcribe(audio, language=self.language,
                                                  initial_prompt=self.prompt)
            except Overloaded as exc:
                await self.send({"type": "error", "detail": str(exc)})
                continue
            except Exception as exc:  # Verbindung soll Einzelfehler überleben
                log.exception("Transkription im Stream fehlgeschlagen")
                await self.send({"type": "error", "detail": str(exc)})
                continue
            payload = {"type": "transcript", **tr.to_dict(with_segments=False), "reason": reason}
            if t is not None:
                payload["t"] = round(t, 2)
            await self.send(payload)


async def _start_wyoming(worker: TranscriptionWorker, settings: ServiceSettings):
    try:
        import stt_wyoming
    except ImportError as exc:
        log.warning("Wyoming deaktiviert (Paket 'wyoming' fehlt): %s", exc)
        return None
    try:
        server = await stt_wyoming.start_server(worker, settings.wyoming_host, settings.wyoming_port)
        log.info("Wyoming-Server auf tcp://%s:%s", settings.wyoming_host, settings.wyoming_port)
        return server
    except OSError as exc:
        log.error("Wyoming-Server nicht startbar (%s:%s): %s",
                  settings.wyoming_host, settings.wyoming_port, exc)
        return None


logging.basicConfig(
    level=os.environ.get("STT_LOG_LEVEL", "INFO").upper(),
    format="%(asctime)s %(levelname)s %(name)s: %(message)s",
)
for _noisy in ("httpx", "faster_whisper"):  # Modell-Downloads / "Processing audio …" pro Anfrage
    logging.getLogger(_noisy).setLevel(logging.WARNING)

app = create_app()

if __name__ == "__main__":
    import uvicorn

    uvicorn.run(app, host=os.environ.get("STT_HOST", "0.0.0.0"),
                port=int(os.environ.get("STT_PORT", "9998")))