LLeMbas/src/lembas/services/audio.py

"""Speech to text and text to speech, against OpenAI-shaped audio endpoints.

The same reasoning as the chat client: plain httpx rather than an SDK, because
the target is not api.openai.com so much as whisper.cpp's server, Speaches,
faster-whisper-server, Kokoro and anything else exposing ``/v1/audio/*``. They
agree on the request and disagree politely about the response, so this is
tolerant about what comes back.

Two endpoints, not one. A local install almost always runs transcription and
speech as separate processes -- they are different models on different
schedules -- and forcing them onto one base URL would mean the common case
could not be configured at all.
"""

from __future__ import annotations

import logging
import time
from collections.abc import AsyncIterator
from typing import Any

import httpx

from lembas.config import settings as env_settings
from lembas.services.crypto import decrypt
from lembas.services.llm.openai_client import (
    Endpoint,
    LLMError,
    describe_http_error,
    wrap_transport_error,
)

log = logging.getLogger(__name__)

# api.openai.com has no endpoint that lists voices, so when one is not offered
# these are what a caller can reasonably assume. Anything else -- Kokoro's sixty
# or so -- is discovered.
OPENAI_VOICES = ("alloy", "echo", "fable", "onyx", "nova", "shimmer")

# Formats every player in a browser can decode. opus is deliberately absent:
# some endpoints emit it in an ogg container that Safari will not play.
FORMATS = ("mp3", "wav", "flac", "aac")

# Discovery is cached because the voice list is read every time anyone opens
# their settings, and waking a model server to answer that is rude.
_VOICE_TTL = 300.0
_voice_cache: dict[str, tuple[float, list[str]]] = {}


def endpoint_for(config: dict[str, Any], side: str) -> Endpoint:
    """Build an Endpoint from the stored audio settings.

    `side` is "stt" or "tts". Endpoint is a frozen snapshot with the key
    already decrypted, so nothing downstream has to know the secret was ever
    encrypted -- or hold a database session while it streams.
    """
    base_url = (config.get(f"{side}_base_url") or "").strip()
    if not base_url:
        raise LLMError("No audio endpoint has been configured.")
    return Endpoint(
        base_url=base_url.rstrip("/"),
        api_key=decrypt(config.get(f"{side}_api_key_encrypted") or ""),
        extra_headers={},
        name=base_url,
    )


async def transcribe(
    endpoint: Endpoint,
    *,
    data: bytes,
    filename: str,
    content_type: str,
    model: str = "whisper-1",
    language: str = "",
) -> str:
    """Turn recorded audio into text.

    `model` is sent even to servers that ignore it: whisper.cpp serves one model
    and does not care, while a router in front of several will not dispatch
    without it. `language` is omitted when empty, which is what asks the server
    to detect it -- sending an empty string instead makes some of them fail.
    """
    form: dict[str, Any] = {"model": model, "response_format": "json"}
    if language:
        form["language"] = language

    try:
        async with httpx.AsyncClient(timeout=env_settings.request_timeout) as client:
            response = await client.post(
                endpoint.url("audio/transcriptions"),
                headers=_headers_without_content_type(endpoint),
                data=form,
                files={"file": (filename, data, content_type or "application/octet-stream")},
            )
            response.raise_for_status()
    except httpx.HTTPStatusError as exc:
        raise LLMError(describe_http_error(exc), status_code=exc.response.status_code) from exc
    except httpx.RequestError as exc:
        raise wrap_transport_error(exc, endpoint) from exc

    try:
        payload = response.json()
    except ValueError:
        # response_format=text is what some servers give regardless of the ask.
        return response.text.strip()

    if isinstance(payload, dict):
        text = payload.get("text")
        if isinstance(text, str):
            return text.strip()
        error = payload.get("error")
        if error:
            raise LLMError(str(error))
    raise LLMError("The transcription endpoint returned no text.")


async def speak(
    endpoint: Endpoint,
    text: str,
    *,
    model: str = "tts-1",
    voice: str = "",
    fmt: str = "mp3",
    speed: float = 1.0,
) -> tuple[str, AsyncIterator[bytes]]:
    """Synthesise speech, returning its content type and a byte stream.

    Streamed rather than buffered: a long reply is a lot of audio, and playback
    can start on the first chunk instead of after the last.
    """
    if not text.strip():
        raise LLMError("There is nothing to read out.")

    body: dict[str, Any] = {
        "model": model,
        "input": text,
        "response_format": fmt if fmt in FORMATS else "mp3",
    }
    if voice:
        body["voice"] = voice
    if speed and speed != 1.0:
        body["speed"] = speed

    client = httpx.AsyncClient(timeout=env_settings.request_timeout)
    try:
        request = client.build_request(
            "POST", endpoint.url("audio/speech"), headers=endpoint.headers(), json=body
        )
        response = await client.send(request, stream=True)
        if response.status_code >= 400:
            # Nothing has been read yet on a streaming response, and the error
            # detail is in the body.
            await response.aread()
            response.raise_for_status()
    except httpx.HTTPStatusError as exc:
        await client.aclose()
        raise LLMError(describe_http_error(exc), status_code=exc.response.status_code) from exc
    except httpx.RequestError as exc:
        await client.aclose()
        raise wrap_transport_error(exc, endpoint) from exc
    except Exception:
        await client.aclose()
        raise

    media_type = response.headers.get("content-type", f"audio/{body['response_format']}")

    async def stream() -> AsyncIterator[bytes]:
        # The client is closed here rather than by the caller: it has to outlive
        # this function, and a response abandoned without aclose leaks a socket.
        try:
            async for chunk in response.aiter_bytes():
                yield chunk
        finally:
            await response.aclose()
            await client.aclose()

    return media_type, stream()


async def voices(endpoint: Endpoint, *, refresh: bool = False) -> list[str]:
    """Voices the speech endpoint offers, newest answer cached briefly.

    Falls back to the OpenAI six on a 404, which is not an error: the official
    API simply has no such endpoint, and its voices are a fixed list everyone
    already knows.
    """
    key = endpoint.base_url
    cached = _voice_cache.get(key)
    if cached and not refresh and time.monotonic() - cached[0] < _VOICE_TTL:
        return cached[1]

    try:
        async with httpx.AsyncClient(timeout=15.0) as client:
            response = await client.get(
                endpoint.url("audio/voices"), headers=endpoint.headers()
            )
            if response.status_code == 404:
                found = list(OPENAI_VOICES)
                _voice_cache[key] = (time.monotonic(), found)
                return found
            response.raise_for_status()
            payload = response.json()
    except httpx.HTTPStatusError as exc:
        raise LLMError(describe_http_error(exc), status_code=exc.response.status_code) from exc
    except httpx.RequestError as exc:
        raise wrap_transport_error(exc, endpoint) from exc
    except ValueError as exc:
        raise LLMError("The endpoint returned a response that was not JSON.") from exc

    found = _parse_voices(payload)
    if not found:
        found = list(OPENAI_VOICES)
    _voice_cache[key] = (time.monotonic(), found)
    return found


def _parse_voices(payload: Any) -> list[str]:
    """Pull voice names out of whatever shape the server chose.

    Kokoro answers ``{"voices": [{"id": "af_heart", ...}]}``; older builds and
    some others answer ``{"voices": ["af_heart", ...]}``; a couple return the
    bare list. All three are the same information.
    """
    entries = payload
    if isinstance(payload, dict):
        for field in ("voices", "data"):
            if isinstance(payload.get(field), list):
                entries = payload[field]
                break
    if not isinstance(entries, list):
        return []

    names: list[str] = []
    for entry in entries:
        if isinstance(entry, str) and entry:
            names.append(entry)
        elif isinstance(entry, dict):
            name = entry.get("id") or entry.get("name") or entry.get("voice")
            if isinstance(name, str) and name:
                names.append(name)
    # Sorted and de-duplicated: sixty voices in the server's arbitrary order is
    # not a list anyone can pick from.
    return sorted(dict.fromkeys(names))


def _headers_without_content_type(endpoint: Endpoint) -> dict[str, str]:
    """Endpoint headers minus Content-Type.

    httpx sets the multipart Content-Type itself, including the boundary.
    Leaving the JSON one in place overrides it and the server sees a body it
    cannot parse.
    """
    return {k: v for k, v in endpoint.headers().items() if k.lower() != "content-type"}


def forget_voices() -> None:
    """Drop the discovery cache. Used when an administrator changes the URL."""
    _voice_cache.clear()


def template_flags(db, user) -> dict[str, Any]:
    """What the chat templates need to know about audio.

    Lives here rather than in one page module because a message bubble is
    rendered from four places -- the chat page, the two message endpoints, and
    the SSE stream, which has no request at all -- and each of them needs the
    same three booleans. Getting one of them wrong is how a speaker button ends
    up on a page that cannot use it.
    """
    from lembas.security import permissions
    from lembas.services import settings_store

    config = settings_store.audio(db)
    allowed = permissions.resolve(db, user)
    listen = bool(config.get("tts_enabled")) and allowed.get("audio.listen", False)
    preferences = (user.settings_json or {}).get("audio") or {} if user else {}

    return {
        "audio": config,
        "user_audio": preferences,
        "can_dictate": bool(config.get("stt_enabled"))
        and allowed.get("audio.transcribe", False),
        "can_listen": listen,
        # Only meaningful when can_listen; the template guards on both.
        "audio_autoplay": listen
        and bool(preferences.get("autoplay", config.get("tts_autoplay"))),
    }


__all__ = [
    "FORMATS",
    "OPENAI_VOICES",
    "LLMError",
    "endpoint_for",
    "forget_voices",
    "speak",
    "transcribe",
    "voices",
]