Bolt-Pet/bolt_pet/audio/mic.py

"""Mic capture + simple energy-based VAD utterance recording.

Ported from desk_client/bolt_desk.py's record_utterance() — same tuning
knobs, same behavior. Kept independent of any UI/threading model so it can
be unit tested by feeding it a fake "stream" object.
"""

from __future__ import annotations

import io
import wave
from typing import Optional, Protocol

import numpy as np

from .. import config


class AudioStream(Protocol):
    """Minimal shape of the object record_utterance() needs — matches
    sounddevice.InputStream's .read(frames) -> (data, overflowed)."""

    def read(self, frames: int): ...


def pcm_to_wav_bytes(pcm: np.ndarray, sample_rate: int = config.SAMPLE_RATE) -> bytes:
    buf = io.BytesIO()
    with wave.open(buf, "wb") as wf:
        wf.setnchannels(1)
        wf.setsampwidth(2)
        wf.setframerate(sample_rate)
        wf.writeframes(pcm.tobytes())
    return buf.getvalue()


def rms(frame: np.ndarray) -> float:
    return float(np.sqrt(np.mean(frame.astype(np.float64) ** 2)))


def flush(stream, max_seconds: float = 10.0, sample_rate: int = config.SAMPLE_RATE) -> int:
    """Throw away whatever is already sitting in the mic's buffer. Returns the
    number of frames dropped.

    PortAudio keeps capturing into a ring buffer while nothing is reading it, so
    audio recorded during a long blocking stretch is still queued when the next
    read happens. That matters exactly once: at the end of a reply the pet is
    about to listen for an answer, and the last fraction of a second of its own
    TTS is in that buffer. It's above the VAD threshold, so `record_utterance`
    treats it as the start of your answer, Deepgram transcribes it, and Bolt is
    handed his own sentence as if you had said it. With barge-in on, the
    detector was draining the stream during playback and the window is small;
    with `BARGE_IN=false` nothing drains it at all.

    Only safe where the buffer is known to hold *nothing you said* — never
    before a wake-triggered recording, where the rest of "thunderbolt, what
    time is it" is legitimately queued and dropping it clips the request.

    *max_seconds* bounds a single call so this can't chase a stream that's
    filling as fast as it's read. Best-effort: a fake stream in tests has no
    `read_available` and this is a no-op, which is the correct behaviour for
    one."""
    try:
        available = int(getattr(stream, "read_available", 0) or 0)
    except (TypeError, ValueError):
        return 0
    if available <= 0:
        return 0
    frames = min(available, int(max_seconds * sample_rate))
    try:
        stream.read(frames)
    except Exception:
        return 0  # a mid-flush device error is the reader's problem, not ours
    return frames


def record_utterance(
    stream: AudioStream,
    should_continue=lambda: True,
    rms_threshold: int = None,
    silence_end_sec: float = None,
    max_utterance_s: float = None,
    min_utterance_s: float = None,
    grace_s: float = None,
    frame_len: int = config.FRAME_LEN,
    sample_rate: int = config.SAMPLE_RATE,
) -> Optional[np.ndarray]:
    """Capture one utterance from *stream*: wait for speech to start, stop
    after trailing silence. Returns None if nothing usable was heard.

    *should_continue* is polled each frame so a caller can cancel recording
    (e.g. the pet window was closed) without needing threading primitives
    baked into this function.

    *grace_s* is how long to wait for speech to *begin* before giving up.
    The controller stretches it for follow-up questions, where you're being
    asked something and need a moment to think rather than having just said
    the wake word on purpose.
    """
    rms_threshold = config.RMS_THRESHOLD if rms_threshold is None else rms_threshold
    silence_end_sec = config.SILENCE_END_SEC if silence_end_sec is None else silence_end_sec
    max_utterance_s = config.MAX_UTTERANCE_S if max_utterance_s is None else max_utterance_s
    min_utterance_s = config.MIN_UTTERANCE_S if min_utterance_s is None else min_utterance_s
    grace_s = config.GRACE_SECONDS if grace_s is None else grace_s

    frames: list[np.ndarray] = []
    started = False
    silence_frames = 0
    silence_limit = int(silence_end_sec * sample_rate / frame_len)
    max_frames = int(max_utterance_s * sample_rate / frame_len)
    grace_frames = int(grace_s * sample_rate / frame_len)  # how long to wait for speech to begin
    waited = 0

    while should_continue():
        chunk, _ = stream.read(frame_len)
        frame = np.asarray(chunk)[:, 0].copy()
        frame_rms = rms(frame)
        if not started:
            waited += 1
            if frame_rms >= rms_threshold:
                started = True
                frames.append(frame)
            elif waited > grace_frames:
                return None  # woke it up but said nothing
            continue
        frames.append(frame)
        if frame_rms < rms_threshold:
            silence_frames += 1
            if silence_frames >= silence_limit:
                break
        else:
            silence_frames = 0
        if len(frames) >= max_frames:
            break

    if not frames:
        return None
    pcm = np.concatenate(frames)
    if len(pcm) < min_utterance_s * sample_rate:
        return None
    return pcm


def open_input_stream():
    """Real sounddevice input stream, imported lazily so pure-logic tests
    (record_utterance with a fake stream) don't need PortAudio installed."""
    import sounddevice as sd

    return sd.InputStream(
        samplerate=config.SAMPLE_RATE,
        channels=1,
        dtype="int16",
        blocksize=config.FRAME_LEN,
        device=config.MIC_DEVICE,
    )