blurt/blurt_client.py

"""UI-free client for Blurt's transcription providers."""

from __future__ import annotations

import base64
import io
import logging
import threading
from dataclasses import dataclass, field
from typing import Callable

import numpy as np
import requests
from scipy.io import wavfile

OPENAI_TRANSCRIPTIONS_URL = "https://api.openai.com/v1/audio/transcriptions"
OPENROUTER_TRANSCRIPTIONS_URL = "https://openrouter.ai/api/v1/audio/transcriptions"
OPENROUTER_APP_URL = "https://git.felixm.de/felixm/blurt"
OPENROUTER_APP_TITLE = "Blurt"
SAMPLE_RATE = 16_000

TIMEOUT_CONNECT_SECONDS = 10.0
TIMEOUT_BASE_SECONDS = 30.0
TIMEOUT_PER_AUDIO_MINUTE = 30.0
TIMEOUT_MAX_SECONDS = 300.0


@dataclass(frozen=True)
class TranscriptionConfig:
    model: str
    language: str
    provider: str
    api_key: str = field(repr=False)


class TranscriptionError(Exception):
    """The provider answered but returned no usable transcript."""


def transcription_timeout(seconds_of_audio: float) -> tuple[float, float]:
    read = TIMEOUT_BASE_SECONDS + TIMEOUT_PER_AUDIO_MINUTE * (seconds_of_audio / 60.0)
    return TIMEOUT_CONNECT_SECONDS, min(read, TIMEOUT_MAX_SECONDS)


class TranscriptionClient:
    """Encode audio and send requests using one HTTP session per worker thread.

    ``close`` must be called only after callers have stopped their worker pool.
    """

    def __init__(self, session_factory: Callable[[], requests.Session] = requests.Session):
        self._session_factory = session_factory
        self._local = threading.local()
        self._sessions: list[requests.Session] = []
        self._sessions_lock = threading.Lock()
        self._closed = False

    def _session(self) -> requests.Session:
        if self._closed:
            raise RuntimeError("TranscriptionClient is closed")
        session = getattr(self._local, "session", None)
        if session is not None:
            return session
        with self._sessions_lock:
            if self._closed:
                raise RuntimeError("TranscriptionClient is closed")
            session = self._session_factory()
            self._sessions.append(session)
        self._local.session = session
        return session

    def transcribe(self, audio: np.ndarray, config: TranscriptionConfig) -> str:
        if audio.dtype != np.int16 or audio.ndim != 1:
            raise ValueError("audio must be a one-dimensional int16 array at 16 kHz")
        if config.provider not in {"openai", "openrouter"}:
            raise ValueError(f"Unsupported transcription provider: {config.provider}")

        buffer = io.BytesIO()
        wavfile.write(buffer, SAMPLE_RATE, audio)
        wav_bytes = buffer.getvalue()
        seconds = len(audio) / SAMPLE_RATE
        timeout = transcription_timeout(seconds)
        logging.info(
            "Starting transcription: provider=%s model=%s language=%s audio_seconds=%.1f "
            "wav_bytes=%d connect_timeout=%.0fs read_timeout=%.0fs",
            config.provider,
            config.model,
            config.language,
            seconds,
            len(wav_bytes),
            timeout[0],
            timeout[1],
        )

        if config.provider == "openrouter":
            result = self._transcribe_openrouter(wav_bytes, config, timeout)
        else:
            result = self._transcribe_openai(wav_bytes, config, timeout)

        text = result.get("text", "").lstrip()
        logging.info("Transcription text length: %d", len(text))
        if "usage" in result:
            logging.info("Transcription usage: %s", result["usage"])
        if not text:
            error = result.get("error")
            message = (
                error.get("message", "Unknown error")
                if isinstance(error, dict)
                else "Unknown error"
            )
            logging.error("Transcription response did not include text")
            raise TranscriptionError(f"Transcription failed: {message}")
        return text

    def _transcribe_openai(self, wav_bytes, config, timeout) -> dict:
        data = {"model": config.model, "language": config.language}
        if config.model == "gpt-4o-transcribe-diarize":
            data.update(response_format="diarized_json", chunking_strategy="auto")
        response = self._session().post(
            OPENAI_TRANSCRIPTIONS_URL,
            headers={"Authorization": f"Bearer {config.api_key}"},
            files={"file": ("blurt.wav", wav_bytes, "audio/wav")},
            data=data,
            timeout=timeout,
        )
        self._log_response(response)
        response.raise_for_status()
        return response.json()

    def _transcribe_openrouter(self, wav_bytes, config, timeout) -> dict:
        response = self._session().post(
            OPENROUTER_TRANSCRIPTIONS_URL,
            headers={
                "Authorization": f"Bearer {config.api_key}",
                "HTTP-Referer": OPENROUTER_APP_URL,
                "X-OpenRouter-Title": OPENROUTER_APP_TITLE,
            },
            json={
                "model": config.model,
                "input_audio": {
                    "data": base64.b64encode(wav_bytes).decode("ascii"),
                    "format": "wav",
                },
                "language": config.language,
            },
            timeout=timeout,
        )
        self._log_response(response)
        response.raise_for_status()
        return response.json()

    @staticmethod
    def _log_response(response: requests.Response) -> None:
        logging.info(
            "Transcription response: status=%s content_type=%s bytes=%d",
            response.status_code,
            response.headers.get("content-type", ""),
            len(response.content),
        )
        if not response.ok:
            logging.error("Transcription error response: %s", response.text[:1000])

    def close(self) -> None:
        with self._sessions_lock:
            if self._closed:
                return
            self._closed = True
            sessions, self._sessions = self._sessions, []
        for session in sessions:
            session.close()