STT-Assistant/src/stt_daemon/output.py

"""Output sinks for transcripts."""

import json
import logging
import wave
from datetime import datetime
from pathlib import Path

import numpy as np

from .transcriber import TranscriptResult

logger = logging.getLogger(__name__)


def write_wav(path: Path, audio: np.ndarray, sample_rate: int) -> None:
    """Write float32 mono audio in [-1, 1] to a 16-bit PCM WAV file."""
    pcm = (np.clip(audio, -1.0, 1.0) * 32767.0).astype(np.int16)
    with wave.open(str(path), "wb") as wav:
        wav.setnchannels(1)
        wav.setsampwidth(2)
        wav.setframerate(sample_rate)
        wav.writeframes(pcm.tobytes())


class TranscriptSink:
    """Emits transcripts to stdout and/or JSONL files."""

    def __init__(
        self,
        stdout: bool = True,
        jsonl_dir: str | None = None,
        jsonl: bool = False,
        save_audio: bool = False,
        audio_dir: str = "audio",
        sample_rate: int = 16000,
    ):
        """Initialize output sinks.

        Args:
            stdout: If True, emit to stdout.
            jsonl_dir: Directory for daily JSONL files.
            jsonl: If True (and jsonl_dir is set), append JSONL lines to daily
                files in that directory. Defaults off so agent-only runs do not
                persist transcripts to disk.
            save_audio: If True, write each transcript's audio as a WAV file.
            audio_dir: Directory for WAV files (dated subfolders).
            sample_rate: Audio sample rate, written into the WAV header.
        """
        self.stdout = stdout
        self.jsonl = jsonl
        self.jsonl_dir = Path(jsonl_dir) if (jsonl and jsonl_dir) else None
        self.save_audio = save_audio
        self.audio_dir = Path(audio_dir)
        self.sample_rate = sample_rate

    def emit(
        self,
        result: TranscriptResult,
        duration_s: float,
        rms: float = 0.0,
        reason: str = "unknown",
        audio: np.ndarray | None = None,
        transcribe_s: float | None = None,
    ) -> str | None:
        """Emit a transcript result.

        Args:
            result: TranscriptResult from transcriber.
            duration_s: Duration of the utterance in seconds.
            rms: Root-mean-square level of the utterance audio.
            reason: Why the utterance ended (end_of_turn/pause/hard_cap/flush).
            audio: Utterance samples (float32, mono); saved when save_audio is enabled.
            transcribe_s: Wall-clock seconds spent transcribing this utterance.
                None omits the field from the record.

        Returns:
            Path of the saved audio file, or None if nothing was saved/emitted.
        """
        if not result.text:
            return None

        # One timestamp for the whole record: stdout, JSONL name/ts, and the
        # WAV path must agree even when emit() straddles midnight.
        now = datetime.now()

        audio_path = None
        if self.save_audio and audio is not None:
            try:
                audio_path = self._save_audio(audio, reason, now)
            except Exception as e:
                # Best-effort side output: a failed save must not drop the transcript.
                logger.warning("Failed to save audio for this utterance: %s", e)

        # Stdout line: [HH:MM:SS] (lang) text  [rms=0.041 end_of_turn]
        if self.stdout:
            timestamp = now.strftime("%H:%M:%S")
            line = f"[{timestamp}] ({result.language}) {result.text}  [rms={rms:.3f} {reason}]"
            print(line, flush=True)

        # JSONL file per day
        if self.jsonl_dir:
            self.jsonl_dir.mkdir(parents=True, exist_ok=True)
            today = now.strftime("%Y-%m-%d")
            jsonl_path = self.jsonl_dir / f"{today}.jsonl"

            record = {
                "ts": now.isoformat(),
                "lang": result.language,
                "lang_prob": result.language_probability,
                "duration_s": duration_s,
                "rms": round(rms, 5),
                "end_reason": reason,
                "text": result.text,
            }
            if audio_path is not None:
                record["audio"] = str(audio_path)
            if transcribe_s is not None:
                record["transcribe_s"] = round(transcribe_s, 3)

            with open(jsonl_path, "a", encoding="utf-8") as f:
                json.dump(record, f, ensure_ascii=False)
                f.write("\n")

        return str(audio_path) if audio_path is not None else None

    def _save_audio(self, audio: np.ndarray, reason: str, now: datetime) -> Path:
        day_dir = self.audio_dir / now.strftime("%Y-%m-%d")
        day_dir.mkdir(parents=True, exist_ok=True)
        name = f"{now.strftime('%H%M%S')}-{now.strftime('%f')[:3]}-{reason}.wav"
        path = day_dir / name
        write_wav(path, audio, self.sample_rate)
        return path

    def close(self) -> None:
        """Clean up (no-op for file-based output)."""
        pass