atlas-iac/services/hermes/scripts/hermes_stt_client.py

#!/usr/bin/env python3
"""Send one Hermes local-command STT request to the private Whisper service."""

from __future__ import annotations

import argparse
import json
import mimetypes
import os
import secrets
import sys
import time
from pathlib import Path
from urllib.error import URLError
from urllib.request import Request, urlopen

# The private Whisper service briefly restarts on image rolls; a transcription
# request that lands in that window must retry rather than crash the turn.
_MAX_ATTEMPTS = 5
_RETRY_BACKOFF_SECONDS = (0.5, 1.5, 3.0, 5.0)


LANGUAGE_SUFFIX = ".language"


def _normalize_language(value: object) -> str:
    """Accept only a bare ISO-639 code from the private Whisper response."""
    if not isinstance(value, str):
        return ""
    code = value.strip().lower()
    if not 2 <= len(code) <= 3 or not code.isascii() or not code.isalpha():
        return ""
    return code


def _write_result(output_dir: Path, stem: str, transcript: str, language: str) -> Path:
    """Write the .txt Hermes reads, plus the language sidecar when we have one.

    Hermes' local-command contract is "leave a .txt in --output-dir"; it globs
    ``*.txt`` and reads the first match. The sidecar deliberately uses another
    suffix so the transcript stays the only ``.txt`` in the directory.
    """
    output_dir.mkdir(parents=True, exist_ok=True)
    transcript_path = output_dir / f"{stem}.txt"
    transcript_path.write_text(transcript, encoding="utf-8")
    if language:
        (output_dir / f"{stem}{LANGUAGE_SUFFIX}").write_text(language, encoding="utf-8")
    return transcript_path


def _multipart(audio: Path, language: str, model: str) -> tuple[bytes, str]:
    boundary = f"atlas-hermes-{secrets.token_hex(12)}"
    mime = mimetypes.guess_type(audio.name)[0] or "application/octet-stream"
    chunks: list[bytes] = []

    def field(name: str, value: str) -> None:
        chunks.extend(
            [
                f"--{boundary}\r\n".encode(),
                f'Content-Disposition: form-data; name="{name}"\r\n\r\n'.encode(),
                value.encode(),
                b"\r\n",
            ]
        )

    field("language", language)
    field("model", model)
    chunks.extend(
        [
            f"--{boundary}\r\n".encode(),
            f'Content-Disposition: form-data; name="file"; filename="{audio.name}"\r\n'.encode(),
            f"Content-Type: {mime}\r\n\r\n".encode(),
            audio.read_bytes(),
            b"\r\n",
            f"--{boundary}--\r\n".encode(),
        ]
    )
    return b"".join(chunks), boundary


def _transcribe_with_retry(request: Request) -> dict:
    """POST to the Whisper service, retrying a brief outage (a service roll).

    A connection error means the private STT service is momentarily down
    (e.g. rolling to a new image); retry with backoff instead of surfacing a
    Python traceback into the conversation. On exhaustion, raise a plain
    RuntimeError so ``main`` can print one short line, not a stack trace.
    """
    last_error: Exception | None = None
    for attempt in range(_MAX_ATTEMPTS):
        try:
            with urlopen(request, timeout=120) as response:
                return json.loads(response.read().decode("utf-8"))
        except (URLError, ConnectionError, TimeoutError, OSError) as error:
            last_error = error
            if attempt < _MAX_ATTEMPTS - 1:
                delay = _RETRY_BACKOFF_SECONDS[min(attempt, len(_RETRY_BACKOFF_SECONDS) - 1)]
                time.sleep(delay)
    raise RuntimeError(
        "the speech service is temporarily unavailable; please try again"
    ) from last_error


def main() -> None:
    """Transcribe one file and emit the .txt contract Hermes expects."""
    parser = argparse.ArgumentParser()
    parser.add_argument("input_path", type=Path)
    parser.add_argument("--output-dir", required=True, type=Path)
    parser.add_argument("--language", default="auto")
    parser.add_argument("--model", default="small")
    args = parser.parse_args()

    body, boundary = _multipart(args.input_path, args.language, args.model)
    request = Request(
        os.getenv(
            "HERMES_STT_URL",
            "http://hermes-stt.hermes.svc.cluster.local:9000/v1/audio/transcriptions",
        ),
        data=body,
        headers={
            "Content-Type": f"multipart/form-data; boundary={boundary}",
            "Accept": "application/json",
        },
        method="POST",
    )
    result = _transcribe_with_retry(request)
    transcript = str(result.get("text") or "").strip()
    _write_result(
        args.output_dir,
        args.input_path.stem,
        transcript,
        _normalize_language(result.get("language")),
    )


if __name__ == "__main__":
    try:
        main()
    except RuntimeError as error:
        # A concise message, never a traceback, reaches the conversation.
        print(f"speech transcription unavailable: {error}", file=sys.stderr)
        sys.exit(1)