diff --git a/app/pipeline/ingestion.py b/app/pipeline/ingestion.py new file mode 100644 index 0000000..e85733e --- /dev/null +++ b/app/pipeline/ingestion.py @@ -0,0 +1,127 @@ +import numpy as np +from datetime import datetime +from typing import Optional + +# Silero VAD — docker-only; set to None when not available +try: + import torch + _vad_model = None + _TORCH_AVAILABLE = True +except ImportError: + _TORCH_AVAILABLE = False + _vad_model = None + + +def _load_vad(): + global _vad_model + if not _TORCH_AVAILABLE: + raise ImportError("torch not installed") + if _vad_model is None: + import torch + model, _ = torch.hub.load( + repo_or_dir="snakers4/silero-vad", + model="silero_vad", + force_reload=False, + trust_repo=True, + ) + _vad_model = model + return _vad_model + + +def vad_prob(chunk: bytes, sample_rate: int = 16000) -> float: + """Return speech probability (0-1) for a PCM int16 chunk.""" + import torch + model = _load_vad() + audio = np.frombuffer(chunk, dtype=np.int16).astype(np.float32) / 32768.0 + tensor = torch.FloatTensor(audio) + with torch.no_grad(): + prob = model(tensor, sample_rate).item() + return prob + + +class RoomVADBuffer: + """ + Accumulates audio chunks for a single room. + Emits (audio_bytes, start_time) when a speech segment ends. + """ + + SPEECH_THRESHOLD = 0.5 + + def __init__( + self, + silence_threshold_ms: int = 600, + min_speech_ms: int = 500, + chunk_ms: int = 250, + sample_rate: int = 16000, + ): + self.silence_threshold_ms = silence_threshold_ms + self.min_speech_ms = min_speech_ms + self.chunk_ms = chunk_ms + self.sample_rate = sample_rate + self._buffer: list = [] + self._speaking = False + self._silence_ms = 0 + self._speech_ms = 0 # track pure speech duration + self._start_time: Optional[datetime] = None + + def process_chunk( + self, chunk: bytes, timestamp: datetime + ) -> Optional[tuple]: + """ + Returns (audio_bytes, start_time) when a complete speech segment is ready. + Returns None while accumulating or during silence. + """ + prob = vad_prob(chunk) + + if prob >= self.SPEECH_THRESHOLD: + if not self._speaking: + self._speaking = True + self._start_time = timestamp + self._buffer = [] + self._speech_ms = 0 + self._buffer.append(chunk) + self._silence_ms = 0 + self._speech_ms += self.chunk_ms + elif self._speaking: + self._buffer.append(chunk) # include trailing silence + self._silence_ms += self.chunk_ms + if self._silence_ms >= self.silence_threshold_ms: + audio = b"".join(self._buffer) + start = self._start_time + speech_ms = self._speech_ms + self._buffer = [] + self._speaking = False + self._silence_ms = 0 + self._speech_ms = 0 + self._start_time = None + if speech_ms >= self.min_speech_ms: + return audio, start + + return None + + +# Per-room registry: room_name -> {"buffer": RoomVADBuffer, "muted": bool} +_rooms: dict = {} + + +def get_or_create_room(room_id: str, silence_ms: int, min_speech_ms: int) -> dict: + if room_id not in _rooms: + _rooms[room_id] = { + "buffer": RoomVADBuffer(silence_threshold_ms=silence_ms, min_speech_ms=min_speech_ms), + "muted": False, + } + return _rooms[room_id] + + +def mute_room(room_id: str) -> None: + if room_id in _rooms: + _rooms[room_id]["muted"] = True + + +def unmute_room(room_id: str) -> None: + if room_id in _rooms: + _rooms[room_id]["muted"] = False + + +def is_muted(room_id: str) -> bool: + return _rooms.get(room_id, {}).get("muted", False) diff --git a/tests/test_ingestion.py b/tests/test_ingestion.py new file mode 100644 index 0000000..6d2f3a0 --- /dev/null +++ b/tests/test_ingestion.py @@ -0,0 +1,55 @@ +import numpy as np +import pytest +from datetime import datetime, timezone +from unittest.mock import patch, MagicMock + + +def _make_chunk(samples: int = 4000, amplitude: float = 0.0) -> bytes: + """Create a fake PCM int16 chunk.""" + audio = (np.ones(samples) * amplitude * 32767).astype(np.int16) + return audio.tobytes() + + +@pytest.fixture +def buf(mocker): + """RoomVADBuffer with mocked VAD model.""" + mocker.patch( + "app.pipeline.ingestion.vad_prob", + side_effect=lambda chunk: 0.9 if np.frombuffer(chunk, np.int16).max() > 100 else 0.1, + ) + from app.pipeline.ingestion import RoomVADBuffer + return RoomVADBuffer(silence_threshold_ms=600, min_speech_ms=500, chunk_ms=250) + + +def test_silence_produces_no_segment(buf): + ts = datetime(2026, 1, 1, tzinfo=timezone.utc) + for _ in range(10): + result = buf.process_chunk(_make_chunk(amplitude=0.0), ts) + assert result is None + + +def test_speech_then_silence_produces_segment(buf): + ts = datetime(2026, 1, 1, tzinfo=timezone.utc) + # 4 speech chunks = 1 second of speech + for _ in range(4): + result = buf.process_chunk(_make_chunk(amplitude=0.5), ts) + assert result is None # still accumulating + # 3 silence chunks = 750ms silence (> 600ms threshold) + segment = None + for _ in range(3): + segment = buf.process_chunk(_make_chunk(amplitude=0.0), ts) + assert segment is not None + assert isinstance(segment, tuple) + audio_bytes, start_time = segment + assert len(audio_bytes) > 0 + + +def test_short_speech_discarded(buf): + ts = datetime(2026, 1, 1, tzinfo=timezone.utc) + # 1 speech chunk = 250ms (below 500ms min) + buf.process_chunk(_make_chunk(amplitude=0.5), ts) + # Now silence to trigger emission + result = None + for _ in range(3): + result = buf.process_chunk(_make_chunk(amplitude=0.0), ts) + assert result is None # discarded as too short