watch_sceen_ai/wsai/backends/whisper.py

"""Real STT via faster-whisper, run as a persistent out-of-venv worker.

faster-whisper needs its own interpreter (whisper312) and loading the model
costs several seconds, so we keep one worker process alive and stream
transcription requests to it (see whisper_worker.py for the protocol) — the
same warm-worker shape as MeloTTS on the TTS side.

`transcribe(wav)` is the core engine: hand it a wav path, get the recognised
text back. It is what closes the voice round trip (TTS wav -> STT text) and what
the Discord voice path will call once per detected utterance.

`utterances()` turns this into a SpeechToText source: it pulls finished-utterance
wav paths from an injected `audio_source` and yields a transcribed Utterance for
each. Until the Discord voice receiver is wired, `audio_source` is None and the
loop simply idles (like the eyes-free perception loop), while `transcribe()`
stays usable directly.

Env:
    WSAI_WHISPER_PYTHON    interpreter with faster-whisper installed
                           (default: /home/claude/jarvis-stt/whisper312/bin/python)
    WSAI_WHISPER_MODEL     model size/name (default: small)
    WSAI_WHISPER_DEVICE    cpu | cuda | auto  (default auto: GPU if present,
                           else CPU; the worker falls back to CPU if CUDA fails)
    WSAI_WHISPER_LANGUAGE  forced language, e.g. ko (default ko; "" = autodetect)
"""

from __future__ import annotations

import asyncio
import collections
import json
import logging
import os
import time
from pathlib import Path
from typing import AsyncIterator

from ..interfaces import Utterance

log = logging.getLogger("wsai.stt.whisper")

_DEFAULT_PYTHON = "/home/claude/jarvis-stt/whisper312/bin/python"


def _cuda_lib_dirs(python_exe: str) -> list[str]:
    """nvidia/*/lib dirs of the worker venv (cublas, cudnn, ...), for
    LD_LIBRARY_PATH so ctranslate2 can dlopen the CUDA runtime. Empty if the
    interpreter has no such packages (CPU-only install)."""
    import glob

    # Use the literal path, NOT .resolve(): the venv's bin/python is a symlink
    # into the uv-managed interpreter, and resolving it would jump out of the
    # venv and miss its site-packages/nvidia libs.
    venv = Path(python_exe).parent.parent  # .../bin/python -> venv root
    dirs = glob.glob(str(venv / "lib" / "python*" / "site-packages" / "nvidia" / "*" / "lib"))
    return sorted(set(dirs))


class WhisperSTT:
    def __init__(
        self,
        *,
        python: str | None = None,
        model: str | None = None,
        device: str | None = None,
        language: str | None = None,
        audio_source: AsyncIterator[str] | None = None,
    ) -> None:
        self.python = python or os.environ.get("WSAI_WHISPER_PYTHON", _DEFAULT_PYTHON)
        self.model = model or os.environ.get("WSAI_WHISPER_MODEL", "small")
        self.device = device or os.environ.get("WSAI_WHISPER_DEVICE", "auto")
        # "" means autodetect; a real code like "ko" forces the language.
        env_lang = os.environ.get("WSAI_WHISPER_LANGUAGE", "ko")
        self.language = language if language is not None else (env_lang or None)
        self.audio_source = audio_source
        self._proc: asyncio.subprocess.Process | None = None
        self._lock = asyncio.Lock()
        self.load_ms: int | None = None
        self.resolved_device: str | None = None  # "cuda" | "cpu", known after start
        # Keep the worker's most recent stderr so a crash reports its real cause
        # instead of a bare JSONDecodeError. Bounded so it can't grow unbounded.
        self._stderr_tail: collections.deque[str] = collections.deque(maxlen=40)
        self._stderr_task: asyncio.Task | None = None

    async def _drain_stderr(self, stream: asyncio.StreamReader) -> None:
        # The worker redirects fd1 -> fd2, so model/library chatter lands on
        # stderr. If we PIPE but never read it, the pipe buffer fills and the
        # worker blocks. So we drain continuously, keeping only the last lines.
        try:
            while True:
                line = await stream.readline()
                if not line:
                    return
                self._stderr_tail.append(line.decode(errors="replace").rstrip())
        except asyncio.CancelledError:
            raise
        except Exception:  # draining must never crash the caller
            return

    def _stderr_hint(self) -> str:
        tail = "\n".join(self._stderr_tail)
        return f" worker stderr tail:\n{tail}" if tail else " (worker produced no stderr)"

    async def warmup(self) -> None:
        """Load the model now so the first real utterance is transcribed warm."""
        await self._ensure()

    async def _ensure(self) -> None:
        if self._proc is not None and self._proc.returncode is None:
            return
        env = {
            **os.environ,
            "WSAI_WHISPER_MODEL": self.model,
            "WSAI_WHISPER_DEVICE": self.device,
        }
        # ctranslate2 dlopens libcublas/libcudnn from the whisper venv's nvidia
        # pip packages; the dynamic loader only honours LD_LIBRARY_PATH captured
        # at exec, so inject those lib dirs into the child env here (harmless on
        # CPU). Without this the CUDA model loads but transcribe() dies with
        # "Library libcublas.so.12 is not found".
        lib_dirs = _cuda_lib_dirs(self.python)
        if lib_dirs:
            prev = env.get("LD_LIBRARY_PATH", "")
            env["LD_LIBRARY_PATH"] = ":".join(lib_dirs + ([prev] if prev else []))
        if self.language:
            env["WSAI_WHISPER_LANGUAGE"] = self.language
        repo_root = str(Path(__file__).resolve().parents[2])
        self._proc = await asyncio.create_subprocess_exec(
            self.python, "-m", "wsai.backends.whisper_worker",
            cwd=repo_root, env=env,
            stdin=asyncio.subprocess.PIPE,
            stdout=asyncio.subprocess.PIPE,
            stderr=asyncio.subprocess.PIPE,
        )
        self._stderr_tail.clear()
        assert self._proc.stderr is not None
        self._stderr_task = asyncio.create_task(self._drain_stderr(self._proc.stderr))
        ready = await self._proc.stdout.readline()
        if not ready:  # worker died before signalling ready
            await self._proc.wait()
            raise RuntimeError(
                f"whisper worker exited before ready (code {self._proc.returncode})."
                f"{self._stderr_hint()}"
            )
        try:
            info = json.loads(ready.decode())
        except json.JSONDecodeError as exc:
            raise RuntimeError(
                f"whisper worker sent invalid ready line {ready!r}: {exc}."
                f"{self._stderr_hint()}"
            ) from exc
        if not info.get("ready"):
            raise RuntimeError(
                f"whisper worker failed to start: {info}.{self._stderr_hint()}"
            )
        self.load_ms = info.get("ms")
        self.resolved_device = info.get("device")
        log.info(
            "whisper worker ready in %s ms on %s (model %s)",
            self.load_ms, info.get("device"), info.get("model"),
        )

    async def transcribe(self, wav_path: str, *, language: str | None = None) -> str:
        """Transcribe one wav file to text using the warm worker."""
        await self._ensure()
        req: dict[str, object] = {"wav": wav_path}
        lang = language if language is not None else self.language
        if lang:
            req["language"] = lang
        s = time.monotonic()
        async with self._lock:
            assert self._proc and self._proc.stdin and self._proc.stdout
            self._proc.stdin.write((json.dumps(req) + "\n").encode())
            await self._proc.stdin.drain()
            resp = await self._proc.stdout.readline()
        if not resp:
            raise RuntimeError(f"whisper worker closed unexpectedly.{self._stderr_hint()}")
        res = json.loads(resp.decode())
        if not res.get("ok"):
            raise RuntimeError(f"whisper transcribe failed: {res.get('error')}")
        log.debug(
            "transcribe %d ms (worker %s ms): %s",
            int((time.monotonic() - s) * 1000), res.get("ms"), res.get("text", "")[:60],
        )
        return res.get("text", "")

    async def utterances(self) -> AsyncIterator[Utterance]:
        """Yield an Utterance per finished-utterance wav from `audio_source`.

        With no audio source wired yet (Discord voice receiver pending) this
        idles and returns, leaving the voice loop dormant but valid."""
        if self.audio_source is None:
            return
        async for wav_path in self.audio_source:
            text = await self.transcribe(wav_path)
            if text:
                yield Utterance(text=text, ts=time.monotonic(), source="voice")

    async def aclose(self) -> None:
        if self._proc is not None and self._proc.returncode is None:
            try:
                self._proc.terminate()
                await asyncio.wait_for(self._proc.wait(), timeout=5)
            except (ProcessLookupError, asyncio.TimeoutError):
                pass
        if self._stderr_task is not None:
            self._stderr_task.cancel()
            try:
                await self._stderr_task
            except (asyncio.CancelledError, Exception):
                pass
            self._stderr_task = None
        self._proc = None