realtime_voice_service/providers/stt.py

from __future__ import annotations

import asyncio
import audioop
import io
import json
import os
import wave

from realtime_voice_service.providers.base import BaseSTT


def _api_base() -> str:
    return (os.getenv("ELEVENLABS_API_BASE", "https://api.elevenlabs.io").strip() or "https://api.elevenlabs.io").rstrip("/")


def _timeout_seconds() -> float:
    raw = os.getenv("ELEVENLABS_TIMEOUT_SECONDS")
    if raw is None:
        return 30.0
    try:
        return max(float(raw.strip()), 1.0)
    except ValueError:
        return 30.0


def _pcm16le_to_wav_bytes(
    pcm_bytes: bytes,
    *,
    sample_rate_hz: int,
    channels: int = 1,
    sample_width_bytes: int = 2,
) -> bytes:
    handle = io.BytesIO()
    with wave.open(handle, "wb") as wav_file:
        wav_file.setnchannels(channels)
        wav_file.setsampwidth(sample_width_bytes)
        wav_file.setframerate(sample_rate_hz)
        wav_file.writeframes(pcm_bytes)
    return handle.getvalue()


def _resample_pcm16le(
    pcm_bytes: bytes,
    *,
    input_rate_hz: int,
    output_rate_hz: int,
) -> bytes:
    if not pcm_bytes or input_rate_hz == output_rate_hz:
        return pcm_bytes
    converted, _ = audioop.ratecv(
        pcm_bytes,
        2,
        1,
        input_rate_hz,
        output_rate_hz,
        None,
    )
    return converted


class ElevenLabsSTT(BaseSTT):
    def __init__(
        self,
        *,
        api_key: str | None = None,
        api_base: str | None = None,
        model_id: str | None = None,
        input_sample_rate_hz: int = 8000,
        target_sample_rate_hz: int = 16000,
        timeout_seconds: float | None = None,
        language_code: str | None = None,
    ) -> None:
        self._api_key = str(api_key if api_key is not None else os.getenv("ELEVENLABS_API_KEY", "")).strip()
        self._api_base = str(api_base or _api_base()).strip().rstrip("/")
        self._model_id = str(model_id or os.getenv("ELEVENLABS_STT_MODEL_ID", "scribe_v2")).strip() or "scribe_v2"
        self._input_sample_rate_hz = max(int(input_sample_rate_hz), 1)
        self._target_sample_rate_hz = max(int(target_sample_rate_hz), 1)
        self._timeout_seconds = max(float(timeout_seconds if timeout_seconds is not None else _timeout_seconds()), 1.0)
        self._language_code = str(language_code or os.getenv("ELEVENLABS_STT_LANGUAGE_CODE", "")).strip() or None

    async def transcribe(self, audio_bytes: bytes) -> str:
        if not audio_bytes:
            return ""
        if not self._api_key:
            raise RuntimeError("ELEVENLABS_API_KEY is required for ElevenLabs STT")

        pcm_bytes, sample_rate_hz = self._extract_pcm(audio_bytes)
        if sample_rate_hz != self._target_sample_rate_hz:
            pcm_bytes = _resample_pcm16le(
                pcm_bytes,
                input_rate_hz=sample_rate_hz,
                output_rate_hz=self._target_sample_rate_hz,
            )
            sample_rate_hz = self._target_sample_rate_hz

        wav_bytes = _pcm16le_to_wav_bytes(pcm_bytes, sample_rate_hz=sample_rate_hz)

        try:
            import aiohttp
        except Exception as exc:  # noqa: BLE001
            raise RuntimeError("The `aiohttp` package is required for ElevenLabs STT") from exc

        form = aiohttp.FormData()
        form.add_field("model_id", self._model_id)
        form.add_field("timestamps_granularity", "none")
        form.add_field("diarize", "false")
        if self._language_code:
            form.add_field("language_code", self._language_code)
        form.add_field(
            "file",
            wav_bytes,
            filename="turn.wav",
            content_type="audio/wav",
        )

        timeout = aiohttp.ClientTimeout(total=self._timeout_seconds)
        try:
            async with aiohttp.ClientSession(timeout=timeout) as session:
                async with session.post(
                    f"{self._api_base}/v1/speech-to-text",
                    headers={"xi-api-key": self._api_key},
                    data=form,
                ) as response:
                    payload_text = await response.text()
                    if response.status >= 400:
                        raise RuntimeError(
                            f"ElevenLabs STT returned HTTP {response.status}: {payload_text[:300]}"
                        )
        except (TimeoutError, asyncio.TimeoutError) as exc:
            raise RuntimeError("ElevenLabs STT request timed out") from exc
        except aiohttp.ClientError as exc:
            raise RuntimeError("ElevenLabs STT request failed") from exc

        try:
            payload = json.loads(payload_text)
        except json.JSONDecodeError as exc:
            raise RuntimeError("ElevenLabs STT returned invalid JSON") from exc
        return str(payload.get("text") or payload.get("transcript") or "").strip()

    def _extract_pcm(self, audio_bytes: bytes) -> tuple[bytes, int]:
        try:
            with wave.open(io.BytesIO(audio_bytes), "rb") as wav_file:
                pcm_bytes = wav_file.readframes(wav_file.getnframes())
                sample_width = wav_file.getsampwidth()
                channels = wav_file.getnchannels()
                sample_rate_hz = int(wav_file.getframerate() or self._input_sample_rate_hz)
                if sample_width != 2:
                    return audio_bytes, self._input_sample_rate_hz
                if channels == 2:
                    pcm_bytes = audioop.tomono(pcm_bytes, sample_width, 0.5, 0.5)
                return pcm_bytes, sample_rate_hz
        except (wave.Error, EOFError):
            return audio_bytes, self._input_sample_rate_hz