| import numpy as np |
| import threading |
| from typing import Optional |
| from dataclasses import dataclass |
|
|
| from ..config import CHUNK_SAMPLES, BUFFER_SAMPLES, SHIFT_SAMPLES, settings |
|
|
|
|
| @dataclass |
| class BufferStats: |
| total_written: int = 0 |
| total_read: int = 0 |
| overruns: int = 0 |
| underruns: int = 0 |
| current_fill: float = 0.0 |
| window_size: int = 0 |
| step_size: int = 0 |
|
|
|
|
| class RingBuffer: |
| """ |
| Thread-safe rolling FIFO ring buffer for 16kHz mono audio. |
| |
| Capacity: BUFFER_SAMPLES (the model's full input window, 24,000 samples/ |
| 1.5s @ 16kHz). Each inference consumes the latest full window and advances |
| the read position by SHIFT_SAMPLES (configurable rolling step). |
| """ |
|
|
| def __init__( |
| self, |
| capacity: int = BUFFER_SAMPLES, |
| window: int = BUFFER_SAMPLES, |
| shift: int = SHIFT_SAMPLES, |
| ): |
| self._capacity = capacity |
| self._window = window |
| self._shift = shift |
| self._buffer = np.zeros(capacity, dtype=np.float32) |
| self._write_pos = 0 |
| self._read_pos = 0 |
| self._filled = 0 |
| self._lock = threading.RLock() |
| self._stats = BufferStats() |
|
|
| @property |
| def capacity(self) -> int: |
| return self._capacity |
|
|
| @property |
| def window_size(self) -> int: |
| return self._window |
|
|
| @property |
| def shift(self) -> int: |
| return self._shift |
|
|
| @property |
| def filled(self) -> int: |
| with self._lock: |
| return self._filled |
|
|
| @property |
| def fill_ratio(self) -> float: |
| with self._lock: |
| return self._filled / self._capacity |
|
|
| @property |
| def stats(self) -> BufferStats: |
| with self._lock: |
| self._stats.current_fill = self.fill_ratio |
| return BufferStats( |
| total_written=self._stats.total_written, |
| total_read=self._stats.total_read, |
| overruns=self._stats.overruns, |
| underruns=self._stats.underruns, |
| current_fill=self._stats.current_fill, |
| ) |
|
|
| def write(self, data: np.ndarray) -> int: |
| """ |
| Write audio data to buffer. Returns number of samples written. |
| Handles wrap-around and overrun detection. |
| """ |
| if data.ndim != 1: |
| raise ValueError(f"Expected 1D array, got {data.ndim}D") |
|
|
| data = data.astype(np.float32, copy=False) |
| n_samples = len(data) |
|
|
| with self._lock: |
| if n_samples > self._capacity: |
| data = data[-self._capacity:] |
| n_samples = self._capacity |
| self._stats.overruns += 1 |
|
|
| remaining = self._capacity - self._write_pos |
| if n_samples <= remaining: |
| self._buffer[self._write_pos:self._write_pos + n_samples] = data |
| self._write_pos = (self._write_pos + n_samples) % self._capacity |
| else: |
| self._buffer[self._write_pos:] = data[:remaining] |
| self._buffer[:n_samples - remaining] = data[remaining:] |
| self._write_pos = (self._write_pos + n_samples) % self._capacity |
|
|
| self._filled = min(self._filled + n_samples, self._capacity) |
| self._stats.total_written += n_samples |
| self._stats.current_fill = self.fill_ratio |
|
|
| return n_samples |
|
|
| def read_latest(self, n_samples: int) -> np.ndarray: |
| """ |
| Read the most recent n_samples from buffer (zero-copy view when possible). |
| Returns zero-padded array if insufficient data. |
| """ |
| with self._lock: |
| if n_samples > self._capacity: |
| n_samples = self._capacity |
|
|
| if self._filled < n_samples: |
| self._stats.underruns += 1 |
| result = np.zeros(n_samples, dtype=np.float32) |
| available = self._filled |
| if available > 0: |
| start = (self._write_pos - available) % self._capacity |
| if start + available <= self._capacity: |
| result[-available:] = self._buffer[start:start + available] |
| else: |
| first_part = self._capacity - start |
| result[-available:-available + first_part] = self._buffer[start:] |
| result[-available + first_part:] = self._buffer[:available - first_part] |
| return result |
|
|
| start = (self._write_pos - n_samples) % self._capacity |
| if start + n_samples <= self._capacity: |
| result = self._buffer[start:start + n_samples].copy() |
| else: |
| first_part = self._capacity - start |
| result = np.concatenate([ |
| self._buffer[start:], |
| self._buffer[:n_samples - first_part] |
| ]) |
|
|
| self._stats.total_read += n_samples |
| return result |
|
|
| def get_processing_window(self) -> Optional[np.ndarray]: |
| """ |
| Get the latest full processing window (model window) for inference. |
| |
| Returns the most recent `window_size` samples (a full 1.5s window under |
| defaults). Returns None if insufficient data accumulated. Advances the |
| read position by `shift` samples after a successful read so windows |
| roll forward by the configured step. |
| """ |
| with self._lock: |
| n = min(self._window, self._capacity) |
| if self._filled < n: |
| return None |
|
|
| start = (self._write_pos - n) % self._capacity |
| if start + n <= self._capacity: |
| window = self._buffer[start:start + n].copy() |
| else: |
| first_part = self._capacity - start |
| window = np.concatenate([ |
| self._buffer[start:], |
| self._buffer[:n - first_part] |
| ]) |
|
|
| step = min(self._shift, self._capacity) |
| self._read_pos = (self._read_pos + step) % self._capacity |
| self._filled = max(0, self._filled - step) |
| self._stats.total_read += n |
| self._stats.current_fill = self.fill_ratio |
|
|
| return window |
|
|
| def clear(self) -> None: |
| """Reset buffer to empty state.""" |
| with self._lock: |
| self._buffer.fill(0.0) |
| self._write_pos = 0 |
| self._read_pos = 0 |
| self._filled = 0 |
| self._stats = BufferStats() |
|
|
| def is_ready_for_inference(self) -> bool: |
| """Check if buffer has enough data for one inference window.""" |
| with self._lock: |
| return self._filled >= min(self._window, self._capacity) |
|
|
|
|
| class MultiConnectionBufferManager: |
| """Manages separate ring buffers per WebSocket connection.""" |
|
|
| def __init__(self, max_connections: int = settings.max_connections): |
| self._max_connections = max_connections |
| self._buffers: dict[str, RingBuffer] = {} |
| self._lock = threading.RLock() |
|
|
| def get_buffer(self, connection_id: str) -> RingBuffer: |
| with self._lock: |
| if connection_id not in self._buffers: |
| if len(self._buffers) >= self._max_connections: |
| raise RuntimeError(f"Max connections ({self._max_connections}) reached") |
| self._buffers[connection_id] = RingBuffer() |
| return self._buffers[connection_id] |
|
|
| def remove_buffer(self, connection_id: str) -> None: |
| with self._lock: |
| self._buffers.pop(connection_id, None) |
|
|
| def get_all_stats(self) -> dict[str, BufferStats]: |
| with self._lock: |
| return {cid: buf.stats for cid, buf in self._buffers.items()} |
|
|
| def cleanup_stale(self, active_ids: set[str]) -> int: |
| with self._lock: |
| stale = set(self._buffers.keys()) - active_ids |
| for cid in stale: |
| del self._buffers[cid] |
| return len(stale) |
