Bolt-Console/bolt_console/dictation.py

"""Dictation: speak into the mic, get text in the composer.

Same route as Bolt-Pet: 16 kHz mono PCM streamed over a websocket to the Bolt
server's own `/desk/stt` relay, which forwards it to Deepgram and sends
Deepgram's JSON back untouched. No Deepgram account or key here — auth is the
desk API key — though the server does meter it against the same credits as a
chat turn.

One `Dictation` is one recording:

    d = Dictation(on_update, on_level, on_done)
    d.start()      # mic opens, audio streams up, on_update(final, interim) as words arrive
    d.stop()       # or silence / the length cap stops it; on_done(text, error) follows
    d.cancel()     # stop and throw the result away

Design points:

- **Live, with a safety net.** Audio goes up as it's captured so the words
  appear while you talk. Every frame is also buffered, and if the socket
  never opened or died, stop() opens a fresh one and sends the whole buffer —
  Bolt-Pet's one-shot fallback.
- **You decide when it's done**, with help: a click stops it, and so do
  `silence_seconds` of quiet after you've said something, or `max_seconds`.
- **No Qt, no numpy.** Callbacks run on background threads; the window turns
  them into signals. Audio is raw int16 bytes, and loudness is computed with
  the array module. That keeps this identical in the single-file Linux app.
- **Everything with hardware or network is injected** (`connect`,
  `open_stream`), so the whole flow is tested without a mic or a server.
"""

from __future__ import annotations

import array
import json
import math
import queue
import sys
import threading
import time
import urllib.parse
from typing import Callable, List, Optional

from . import config

SAMPLE_RATE = 16000
BLOCK_FRAMES = 1280           # 80 ms, same chunking as Bolt-Pet
SPEECH_RMS = 500              # int16 RMS above this counts as talking


def unavailable_reason() -> Optional[str]:
    """None if dictation can work here, otherwise a sentence saying why —
    shown as the mic button's tooltip instead of a crash."""
    if not config.is_configured():
        return "Set BOLT_SERVER_URL and DESK_API_KEY to use dictation"
    try:
        import websocket  # noqa: F401  (websocket-client)
    except Exception:
        return "Dictation needs websocket-client:  pip install websocket-client"
    try:
        import sounddevice  # noqa: F401
    except OSError:
        return ("Dictation needs the PortAudio library (e.g. sudo apt install libportaudio2)"
                if sys.platform.startswith("linux") else "PortAudio isn't available on this system")
    except Exception:
        return "Dictation needs sounddevice:  pip install sounddevice"
    return None


def mic_device():
    """MIC_DEVICE from .env: an index (as saved by Bolt-Pet's mic picker) or a
    name substring; blank means the system default."""
    raw = str(getattr(config, "MIC_DEVICE", "") or "").strip()
    if not raw:
        return None
    return int(raw) if raw.lstrip("-").isdigit() else raw


def connect(sample_rate: int = SAMPLE_RATE):
    """Open the server's /desk/stt websocket for one recording."""
    import websocket

    query = urllib.parse.urlencode({
        "session_id": config.SESSION_ID, "encoding": "linear16", "sample_rate": sample_rate,
    })
    url = config.SERVER_URL.replace("http", "ws", 1) + "/desk/stt?" + query
    return websocket.create_connection(url, header=[f"X-Desk-Api-Key: {config.API_KEY}"], timeout=10)


def open_stream(callback: Callable[[bytes], None], sample_rate: int = SAMPLE_RATE):
    """Start the default (or MIC_DEVICE) microphone, handing raw int16 mono
    bytes to *callback* from PortAudio's thread. Returns something with
    stop() and close()."""
    import sounddevice as sd

    stream = sd.RawInputStream(
        samplerate=sample_rate, channels=1, dtype="int16", blocksize=BLOCK_FRAMES,
        device=mic_device(), callback=lambda data, _frames, _time, _status: callback(bytes(data)),
    )
    stream.start()
    return stream


def rms(chunk: bytes) -> float:
    samples = array.array("h")
    samples.frombytes(chunk[: len(chunk) - len(chunk) % 2])
    if not samples:
        return 0.0
    return math.sqrt(sum(s * s for s in samples) / len(samples))


def parse_result(message) -> tuple:
    """(text, is_final) from a Deepgram results frame; ("", False) for
    anything else. Never raises — it runs on the reader thread."""
    try:
        if isinstance(message, bytes):
            message = message.decode("utf-8", "ignore")
        data = json.loads(message)
    except (TypeError, ValueError):
        return "", False
    if not isinstance(data, dict) or data.get("type") not in (None, "Results"):
        return "", False
    alternatives = (data.get("channel") or {}).get("alternatives") or []
    if not alternatives:
        return "", False
    text = str((alternatives[0] or {}).get("transcript") or "").strip()
    return text, bool(data.get("is_final") or data.get("speech_final"))


class Dictation:
    def __init__(
        self,
        on_update: Callable[[str, str], None],
        on_level: Callable[[float], None],
        on_done: Callable[[str, Optional[str]], None],
        *,
        connect_fn: Callable = None,
        open_stream_fn: Callable = None,
        silence_seconds: float = 3.0,
        max_seconds: float = 300.0,
        sample_rate: int = SAMPLE_RATE,
        on_stopping: Callable[[], None] = None,
    ):
        self._on_update, self._on_level, self._on_done = on_update, on_level, on_done
        # Recording over (click, silence or cap), transcript not back yet.
        self._on_stopping = on_stopping or (lambda: None)
        self._connect = connect_fn or connect
        self._open_stream = open_stream_fn or open_stream
        self._silence_seconds = silence_seconds
        self._max_seconds = max_seconds
        self._rate = sample_rate
        self._frames: "queue.Queue[bytes]" = queue.Queue()
        self._buffer: List[bytes] = []
        self._finals: List[str] = []
        self._interim = ""
        self._lock = threading.Lock()
        self._stop = threading.Event()
        self._cancelled = False
        self._socket = None
        self._socket_ok = False
        self._stream = None
        self._reader = None
        self._pump = None

    # ── control (any thread) ─────────────────────────────────────────────

    def start(self) -> None:
        """Open the mic (raises if that fails — nothing to record without it)
        and try the live socket (failure there is fine: the fallback has the
        audio)."""
        try:
            self._socket = self._connect(self._rate)
            self._socket_ok = True
            self._reader = threading.Thread(target=self._read_loop, name="stt-reader", daemon=True)
            self._reader.start()
        except Exception:
            self._socket, self._socket_ok = None, False
        try:
            self._stream = self._open_stream(self._frames.put, self._rate)
        except Exception:
            self._close_socket()
            raise
        self._pump = threading.Thread(target=self._pump_loop, name="stt-pump", daemon=True)
        self._pump.start()

    def stop(self) -> None:
        self._stop.set()

    def cancel(self) -> None:
        self._cancelled = True
        self._stop.set()

    @property
    def live(self) -> bool:
        """Whether words are arriving as you speak (vs. only after stop)."""
        return self._socket_ok

    # ── the pump: mic frames -> socket, buffer, level, auto-stop ─────────

    def _pump_loop(self) -> None:
        started = time.monotonic()
        heard_speech, last_speech = False, started
        while not self._stop.is_set():
            try:
                chunk = self._frames.get(timeout=0.1)
            except queue.Empty:
                chunk = None
            now = time.monotonic()
            if chunk:
                self._buffer.append(chunk)
                self._send(chunk)
                loudness = rms(chunk)
                self._safe(self._on_level, min(1.0, math.sqrt(loudness / 8000.0)))
                if loudness >= SPEECH_RMS:
                    heard_speech, last_speech = True, now
            if heard_speech and now - last_speech >= self._silence_seconds:
                break
            if now - started >= self._max_seconds:
                break
        self._finish()

    def _send(self, chunk: bytes) -> None:
        if not self._socket_ok:
            return
        try:
            self._socket.send_binary(chunk)
        except Exception:
            self._socket_ok = False  # the fallback will send the buffer instead

    def _finish(self) -> None:
        self._stop.set()
        self._safe(self._on_stopping)
        for step in ("stop", "close"):
            try:
                getattr(self._stream, step)()
            except Exception:
                pass
        self._safe(self._on_level, 0.0)
        if self._cancelled:
            self._close_socket()
            self._safe(self._on_done, "", None)
            return
        # Whatever was still queued when we stopped belongs to the recording.
        while True:
            try:
                chunk = self._frames.get_nowait()
            except queue.Empty:
                break
            self._buffer.append(chunk)
            self._send(chunk)

        # Live socket healthy to the end: its transcript is the answer, even
        # an empty one (silence) — re-sending would pay to transcribe it twice.
        error = None
        if self._socket_ok:
            self._flush_socket()
        else:
            error = self._fallback()
        text = self._text()
        if not text and error is None:
            error = "Didn't catch anything — try again a little closer to the mic."
        self._safe(self._on_done, text, error)

    def _flush_socket(self) -> None:
        """Deepgram flushes its last results after CloseStream; wait for them,
        bounded so a hung socket can't hold the text hostage."""
        try:
            self._socket.send(json.dumps({"type": "CloseStream"}))
        except Exception:
            self._socket_ok = False
        if self._reader is not None:
            self._reader.join(timeout=4)
        self._close_socket()

    def _fallback(self) -> Optional[str]:
        """The live socket never worked or died part-way: one fresh connection,
        the whole buffer, from an empty transcript (a half-live session's words
        would otherwise come back twice)."""
        self._close_socket()
        if self._reader is not None:
            self._reader.join(timeout=1)
        with self._lock:
            self._finals, self._interim = [], ""
        if not self._buffer:
            return None
        try:
            self._socket = self._connect(self._rate)
        except Exception as exc:
            return f"Couldn't reach the transcription server: {exc}"
        self._socket_ok = True
        self._reader = threading.Thread(target=self._read_loop, name="stt-reader", daemon=True)
        self._reader.start()
        audio = b"".join(self._buffer)
        for offset in range(0, len(audio), BLOCK_FRAMES * 2 * 8):
            self._send(audio[offset:offset + BLOCK_FRAMES * 2 * 8])
        self._flush_socket()
        return None

    def _close_socket(self) -> None:
        socket, self._socket = self._socket, None
        if socket is not None:
            try:
                socket.close()
            except Exception:
                pass

    # ── the reader: transcripts back ─────────────────────────────────────

    def _read_loop(self) -> None:
        socket = self._socket
        while socket is not None:
            try:
                message = socket.recv()
            except Exception:
                break
            if not message:
                break
            text, is_final = parse_result(message)
            with self._lock:
                if is_final:
                    if text:
                        self._finals.append(text)
                    self._interim = ""
                elif text:
                    self._interim = text
                finals, interim = " ".join(self._finals), self._interim
            if text:
                self._safe(self._on_update, finals, interim)

    def _text(self) -> str:
        with self._lock:
            # An interim that never became final is still better than nothing.
            return " ".join(self._finals + ([self._interim] if self._interim else [])).strip()

    @staticmethod
    def _safe(callback, *args) -> None:
        try:
            callback(*args)
        except Exception:
            pass