feat(local-session): Redis state store and audio segmentation engine

bots/local_audio_processing.py

@@ -0,0 +1,170 @@
+"""Turning a local session's glued audio into utterances, reusing the bots' VAD.
+
+This is the same webrtcVAD + utterance segmentation meeting bots use, with two local-only
+adjustments: a loudness meter that does not overflow, and driving the VAD on the session's
+own timeline (offsets) rather than the wall clock.
+"""
+
+import logging
+from datetime import timedelta
+
+import numpy as np
+
+from bots.bot_controller.per_participant_non_streaming_audio_input_manager import (
+    PerParticipantNonStreamingAudioInputManager,
+)
+from bots.models import AudioChunk, RecordingManager, Utterance
+
+logger = logging.getLogger(__name__)
+
+# A local recording is one person's device, so utterances are cut more aggressively than a
+# meeting bot's: the desktop wants lines quickly, and short clips transcribe faster.
+LOCAL_SILENCE_DURATION_LIMIT_SECONDS = 1
+LOCAL_UTTERANCE_SIZE_LIMIT_SECONDS = 30
+
+# webrtcvad only accepts 10/20/30ms frames, and the manager's "too large for VAD" guard
+# compares a BYTE length against a SAMPLE count -- so anything over ~15ms of audio silently
+# skips the VAD and is reported as speech. 10ms frames stay under that guard at every
+# supported rate, which is also what the bot adapters happen to emit.
+VAD_FRAME_MS = 10
+BYTES_PER_SAMPLE = 2
+INT16_FULL_SCALE = 32768.0
+
+
+def normalized_rms(audio_bytes):
+    """Loudness of a PCM frame, 0..1.
+
+    The shared calculate_normalized_rms() squares an int16 array *in int16*, which wraps for
+    any |sample| > 181 -- i.e. for all real speech -- so loud audio measures as silence.
+    Widening to float64 first is what the streaming manager already does.
+    """
+    samples = np.frombuffer(audio_bytes, dtype=np.int16).astype(np.float64)
+    if samples.size == 0:
+        return 0.0
+    return float(np.sqrt(np.mean(np.square(samples))) / INT16_FULL_SCALE)
+
+
+class LocalAudioInputManager(PerParticipantNonStreamingAudioInputManager):
+    """The bots' VAD manager, but with a loudness meter that does not overflow.
+
+    Only ``silence_detected`` is overridden, and only to swap in the corrected RMS. The
+    shared function is deliberately left alone so the bot path stays byte-identical; fixing
+    it there is worth doing, but as its own change with its own testing.
+    """
+
+    def silence_detected(self, chunk_bytes):
+        rms_value = normalized_rms(chunk_bytes)
+        if rms_value == 0:
+            self.diagnostic_info["total_chunks_marked_as_silent_due_to_rms_being_zero"] += 1
+            return True
+        if rms_value < 0.01:
+            self.diagnostic_info["total_chunks_marked_as_silent_due_to_rms_being_small"] += 1
+            return True
+        if not self.is_speech(chunk_bytes):
+            self.diagnostic_info["total_chunks_marked_as_silent_due_to_vad"] += 1
+            return True
+        return False
+
+
+def duration_ms(audio, sample_rate):
+    return int(len(audio) / ((sample_rate / 1000) * BYTES_PER_SAMPLE))
+
+
+def create_utterance(recording, participant, message):
+    """Mirror of BotController.process_individual_audio_chunk for a local session.
+
+    Keyed by a deterministic source_uuid (a globally-unique column) so a task retry cannot
+    duplicate an utterance it already wrote before failing.
+    """
+    from bots.tasks.process_utterance_task import process_utterance
+
+    audio_data = message["audio_data"]
+    sample_rate = message["sample_rate"]
+    source_uuid = f"local:{recording.id}:{participant.uuid}:{message['timestamp_ms']}"
+
+    if Utterance.objects.filter(source_uuid=source_uuid).exists():
+        logger.info(f"Local session {recording.bot.object_id}: utterance {source_uuid} already exists, skipping")
+        return
+
+    audio_chunk = AudioChunk.objects.create(
+        recording=recording,
+        audio_format=AudioChunk.AudioFormat.PCM,
+        timestamp_ms=message["timestamp_ms"],
+        duration_ms=duration_ms(audio_data, sample_rate),
+        sample_rate=sample_rate,
+        source=AudioChunk.Sources.PER_PARTICIPANT_AUDIO,
+        participant=participant,
+        is_blob_stored_remotely=False,
+        audio_blob=audio_data,
+    )
+    utterance = Utterance.objects.create(
+        source=Utterance.Sources.PER_PARTICIPANT_AUDIO,
+        async_transcription=None,
+        recording=recording,
+        participant=participant,
+        audio_chunk=audio_chunk,
+        timestamp_ms=audio_chunk.timestamp_ms,
+        duration_ms=audio_chunk.duration_ms,
+        source_uuid=source_uuid,
+    )
+
+    RecordingManager.set_recording_transcription_in_progress(recording)
+    process_utterance.delay(utterance.id)
+    logger.info(f"Local session {recording.bot.object_id}: queued utterance {utterance.id} ({audio_chunk.duration_ms}ms)")
+
+
+def build_manager(recording, participant, sample_rate):
+    def save_audio_chunk_callback(message):
+        create_utterance(recording, participant, message)
+
+    def get_participant_callback(speaker_id):
+        # Fixed for a local session; the manager only needs this to be non-None.
+        return {"participant_uuid": participant.uuid, "participant_full_name": participant.full_name}
+
+    return LocalAudioInputManager(
+        save_audio_chunk_callback=save_audio_chunk_callback,
+        get_participant_callback=get_participant_callback,
+        sample_rate=sample_rate,
+        utterance_size_limit=LOCAL_UTTERANCE_SIZE_LIMIT_SECONDS * sample_rate * BYTES_PER_SAMPLE,
+        silence_duration_limit=LOCAL_SILENCE_DURATION_LIMIT_SECONDS,
+        should_print_diagnostic_info=False,
+    )
+
+
+def feed(manager, source, audio, epoch, start_offset_ms, sample_rate):
+    """Drive the VAD frame by frame on the session's own timeline.
+
+    process_chunk() is called directly rather than add_chunk()/process_chunks(), because
+    process_chunks() probes for silence against datetime.utcnow() -- wall clock -- which
+    would force-flush this timeline instead of following it. Returns the byte count consumed
+    (whole frames); a trailing partial frame rolls into the next drain.
+    """
+    frame_bytes = int(sample_rate * VAD_FRAME_MS / 1000) * BYTES_PER_SAMPLE
+    for offset in range(0, len(audio), frame_bytes):
+        frame = audio[offset : offset + frame_bytes]
+        if len(frame) < frame_bytes:
+            break
+        frame_at = epoch + timedelta(milliseconds=start_offset_ms + (offset // frame_bytes) * VAD_FRAME_MS)
+        manager.process_chunk(source, frame_at, frame)
+    return len(audio) - (len(audio) % frame_bytes)
+
+
+def offset_of_buffered(manager, source, epoch):
+    """Where the still-buffered utterance began, as ms since the session started."""
+    started_at = manager.first_nonsilent_audio_time.get(source)
+    if started_at is None:
+        return None
+    return int((started_at - epoch).total_seconds() * 1000)
+
+
+def flush_remaining(manager, source, epoch, end_offset_ms):
+    """Emit whatever is still buffered, on our timeline rather than the wall clock.
+
+    flush_utterances() probes with datetime.utcnow(); against a session-relative timeline
+    that can compute negative silence and silently drop the final utterance, so the probe is
+    placed just past the end of the audio we actually have.
+    """
+    if not manager.utterances.get(source):
+        return
+    probe_at = epoch + timedelta(milliseconds=end_offset_ms, seconds=LOCAL_SILENCE_DURATION_LIMIT_SECONDS + 1)
+    manager.process_chunk(source, probe_at, None)

bots/local_session_store.py

@@ -0,0 +1,123 @@
+"""Redis-backed state for local recording sessions.
+
+A local session has no long-lived process (unlike a meeting bot), so the audio "memory"
+that a bot keeps in RAM is rebuilt here in Redis, per session and per source:
+
+* ``queue`` -- a FIFO list of uploaded, not-yet-processed segments.
+* ``tail``  -- the unfinished utterance left over after the last drain, glued to the front
+  of the next batch so the VAD only ever cuts on real silence.
+* ``lock``  -- ensures exactly one drain owns a source's tail at a time.
+"""
+
+import base64
+import json
+
+import redis
+from django.conf import settings
+
+# A local recording always has exactly these two audio sources (the desktop tags every
+# chunk with one). Kept here so the views, tasks and finalizer share one definition.
+MIC_SOURCE = "mic"
+SYSTEM_SOURCE = "system"
+LOCAL_SESSION_SOURCES = (MIC_SOURCE, SYSTEM_SOURCE)
+
+TAIL_TTL_SECONDS = 3600
+LOCK_TTL_SECONDS = 120
+QUEUE_TTL_SECONDS = 3600  # so an abandoned session's queued audio can't leak forever
+MAX_SEGMENTS_PER_DRAIN = 200
+
+
+def redis_client():
+    return redis.from_url(settings.REDIS_URL_WITH_PARAMS)
+
+
+def queue_key(bot_id, source):
+    return f"local_session_queue:{bot_id}:{source}"
+
+
+def tail_key(bot_id, source):
+    return f"local_session_tail:{bot_id}:{source}"
+
+
+def lock_key(bot_id, source):
+    return f"local_session_lock:{bot_id}:{source}"
+
+
+def enqueue_segment(bot_id, source, sequence, audio, sample_rate, offset_ms):
+    """Append an uploaded segment to its source's FIFO queue, ready for the next drain."""
+    payload = {
+        "sequence": sequence,
+        "audio": base64.b64encode(audio).decode(),
+        "sample_rate": sample_rate,
+        "offset_ms": offset_ms,
+    }
+    client = redis_client()
+    key = queue_key(bot_id, source)
+    client.rpush(key, json.dumps(payload))
+    client.expire(key, QUEUE_TTL_SECONDS)
+
+
+def clear_session_state(bot_id):
+    """Drop every Redis key for a session. Called when a session is deleted so its queue,
+    tail and lock don't outlive the DB rows (and so an in-flight drain finds nothing)."""
+    client = redis_client()
+    keys = []
+    for source in LOCAL_SESSION_SOURCES:
+        keys += [queue_key(bot_id, source), tail_key(bot_id, source), lock_key(bot_id, source)]
+    client.delete(*keys)
+
+
+def load_tail(client, bot_id, source):
+    """Unfinished-utterance audio plus the offset it began at and the last sequence seen."""
+    raw = client.get(tail_key(bot_id, source))
+    if not raw:
+        return {
+            "audio": b"",
+            "started_offset_ms": None,
+            "end_offset_ms": None,
+            "last_sequence": -1,
+            "sample_rate": None,
+        }
+    payload = json.loads(raw)
+    payload["audio"] = base64.b64decode(payload["audio"])
+    return payload
+
+
+def save_tail(client, bot_id, source, audio, started_offset_ms, end_offset_ms, last_sequence, sample_rate):
+    payload = {
+        "audio": base64.b64encode(bytes(audio)).decode(),
+        "started_offset_ms": started_offset_ms,
+        "end_offset_ms": end_offset_ms,
+        # Kept even when the audio is empty, so a replayed upload is still recognised.
+        "last_sequence": last_sequence,
+        # Carried so a final flush can rebuild the VAD without the caller supplying it.
+        "sample_rate": sample_rate,
+    }
+    client.set(tail_key(bot_id, source), json.dumps(payload), ex=TAIL_TTL_SECONDS)
+
+
+def pop_segments(client, bot_id, source, last_sequence, on_replay=None):
+    """FIFO drain, dropping replays.
+
+    Two kinds of replay have to die here: one already folded into the tail (sequence at or
+    below the tail's), and one duplicated *within this same batch* -- a retrying uploader can
+    land both copies before any drain runs, and checking only against the tail would let the
+    second copy through and splice the same audio in twice.
+    """
+    segments = []
+    seen_sequences = set()
+    for _ in range(MAX_SEGMENTS_PER_DRAIN):
+        raw = client.lpop(queue_key(bot_id, source))
+        if raw is None:
+            break
+        segment = json.loads(raw)
+        sequence = segment["sequence"]
+        if sequence <= last_sequence or sequence in seen_sequences:
+            if on_replay:
+                on_replay(sequence)
+            continue
+        seen_sequences.add(sequence)
+        segment["audio"] = base64.b64decode(segment["audio"])
+        segments.append(segment)
+    segments.sort(key=lambda s: s["sequence"])
+    return segments