mkbk
529 lines
19 KiB
Python
529 lines
19 KiB
Python
from __future__ import annotations
|
|
|
|
import json
|
|
import queue
|
|
import shutil
|
|
import subprocess
|
|
import tempfile
|
|
import time
|
|
from collections.abc import Callable
|
|
from collections import deque
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
from .models import (
|
|
AudioFrame,
|
|
AudioSegment,
|
|
ErrorCode,
|
|
PlaybackResult,
|
|
ProviderError,
|
|
TransportHealth,
|
|
)
|
|
|
|
|
|
class AudioRingBuffer:
|
|
def __init__(self, max_duration_ms: int = 3000) -> None:
|
|
if max_duration_ms <= 0:
|
|
raise ValueError("max_duration_ms must be positive")
|
|
self.max_duration_ms = max_duration_ms
|
|
self._frames: deque[AudioFrame] = deque()
|
|
|
|
def append(self, frame: AudioFrame) -> None:
|
|
if self._frames and frame.timestamp_ms < self._frames[-1].timestamp_ms:
|
|
raise ValueError("audio frame timestamps must be non-decreasing")
|
|
self._frames.append(frame)
|
|
self._trim()
|
|
|
|
def extend(self, frames: list[AudioFrame]) -> None:
|
|
for frame in frames:
|
|
self.append(frame)
|
|
|
|
def frames(self) -> list[AudioFrame]:
|
|
return list(self._frames)
|
|
|
|
def clear(self) -> None:
|
|
self._frames.clear()
|
|
|
|
def _trim(self) -> None:
|
|
if not self._frames:
|
|
return
|
|
latest = self._frames[-1].timestamp_ms
|
|
cutoff = latest - self.max_duration_ms
|
|
while self._frames and self._frames[0].timestamp_ms < cutoff:
|
|
self._frames.popleft()
|
|
|
|
|
|
class MemoryAudioTransport:
|
|
def __init__(
|
|
self,
|
|
frames: list[AudioFrame] | None = None,
|
|
input_available: bool = True,
|
|
output_available: bool = True,
|
|
flush_clears_input: bool = True,
|
|
) -> None:
|
|
self._frames: deque[AudioFrame] = deque(frames or [])
|
|
self.played_segments: list[AudioSegment] = []
|
|
self.flush_count = 0
|
|
self.started = False
|
|
self._input_available = input_available
|
|
self._output_available = output_available
|
|
self._flush_clears_input = flush_clears_input
|
|
|
|
def start_input(
|
|
self, device_id: str | None = None, sample_rate: int = 16000, channels: int = 1
|
|
) -> None:
|
|
if not self._input_available:
|
|
raise ProviderError(
|
|
ErrorCode.AUDIO_INPUT_DEVICE_MISSING,
|
|
"memory input transport is disabled",
|
|
False,
|
|
"memory-transport",
|
|
"transport",
|
|
)
|
|
self.started = True
|
|
|
|
def read_frames(self, timeout_ms: int) -> list[AudioFrame]:
|
|
if not self.started:
|
|
return []
|
|
if not self._frames:
|
|
return []
|
|
return [self._frames.popleft()]
|
|
|
|
def inject(self, frame: AudioFrame) -> None:
|
|
self._frames.append(frame)
|
|
|
|
def play_pcm(self, segment: AudioSegment, interrupt: bool = False) -> PlaybackResult:
|
|
if not self._output_available:
|
|
error = ProviderError(
|
|
ErrorCode.AUDIO_OUTPUT_DEVICE_MISSING,
|
|
"memory output transport is disabled",
|
|
False,
|
|
"memory-transport",
|
|
"transport",
|
|
)
|
|
return PlaybackResult(False, 0, error)
|
|
if interrupt:
|
|
self.played_segments.clear()
|
|
self.played_segments.append(segment)
|
|
return PlaybackResult(True, segment.duration_ms)
|
|
|
|
def play_pcm_chunks(
|
|
self,
|
|
segment: AudioSegment,
|
|
*,
|
|
chunk_ms: int,
|
|
after_chunk: Callable[[AudioSegment, int], bool] | None = None,
|
|
should_stop: Callable[[], bool] | None = None,
|
|
) -> PlaybackResult:
|
|
elapsed_ms = 0
|
|
for chunk in _raw_pcm_chunks(segment, chunk_ms=chunk_ms):
|
|
if should_stop is not None and should_stop():
|
|
return PlaybackResult(True, elapsed_ms)
|
|
playback = self.play_pcm(chunk)
|
|
if playback.error:
|
|
return playback
|
|
elapsed_ms += chunk.duration_ms
|
|
if should_stop is not None and should_stop():
|
|
return PlaybackResult(True, elapsed_ms)
|
|
if after_chunk is not None and after_chunk(chunk, elapsed_ms):
|
|
return PlaybackResult(True, elapsed_ms)
|
|
return PlaybackResult(True, elapsed_ms)
|
|
|
|
def flush_input(self) -> int:
|
|
self.flush_count += 1
|
|
if not self._flush_clears_input:
|
|
return 0
|
|
count = len(self._frames)
|
|
self._frames.clear()
|
|
return count
|
|
|
|
def stop(self) -> None:
|
|
self.started = False
|
|
|
|
def health(self) -> TransportHealth:
|
|
return TransportHealth(
|
|
input_available=self._input_available,
|
|
output_available=self._output_available,
|
|
message="memory transport",
|
|
)
|
|
|
|
|
|
class FileReplayTransport(MemoryAudioTransport):
|
|
@classmethod
|
|
def from_jsonl(cls, path: str | Path) -> "FileReplayTransport":
|
|
frames: list[AudioFrame] = []
|
|
with Path(path).open("r", encoding="utf-8") as handle:
|
|
for line_no, line in enumerate(handle, start=1):
|
|
if not line.strip():
|
|
continue
|
|
item: dict[str, Any] = json.loads(line)
|
|
try:
|
|
frames.append(
|
|
AudioFrame(
|
|
pcm=bytes.fromhex(item.get("pcm_hex", "")),
|
|
sample_rate=int(item.get("sample_rate", 16000)),
|
|
channels=int(item.get("channels", 1)),
|
|
timestamp_ms=int(item["timestamp_ms"]),
|
|
frame_id=int(item.get("frame_id", line_no - 1)),
|
|
metadata=dict(item.get("metadata", {})),
|
|
)
|
|
)
|
|
except KeyError as exc:
|
|
raise ValueError(f"missing required fixture field on line {line_no}: {exc}") from exc
|
|
return cls(frames)
|
|
|
|
@staticmethod
|
|
def write_jsonl(path: str | Path, frames: list[AudioFrame]) -> None:
|
|
with Path(path).open("w", encoding="utf-8") as handle:
|
|
for frame in frames:
|
|
handle.write(
|
|
json.dumps(
|
|
{
|
|
"pcm_hex": frame.pcm.hex(),
|
|
"sample_rate": frame.sample_rate,
|
|
"channels": frame.channels,
|
|
"timestamp_ms": frame.timestamp_ms,
|
|
"frame_id": frame.frame_id,
|
|
"metadata": dict(frame.metadata),
|
|
},
|
|
ensure_ascii=False,
|
|
)
|
|
+ "\n"
|
|
)
|
|
|
|
|
|
class SoundDeviceAudioTransport:
|
|
def __init__(self, sounddevice_module: Any | None = None, output_device: str | None = None) -> None:
|
|
self._sd: Any | None = sounddevice_module
|
|
self._load_error: Exception | None = None
|
|
self._stream: Any | None = None
|
|
self._queue: queue.Queue[AudioFrame] = queue.Queue(maxsize=200)
|
|
self._frame_id = 0
|
|
self._start_time = time.monotonic()
|
|
self._sample_rate = 16000
|
|
self._channels = 1
|
|
self._output_device = output_device
|
|
if self._sd is not None:
|
|
return
|
|
try:
|
|
import sounddevice as sd # type: ignore[import-not-found]
|
|
|
|
self._sd = sd
|
|
except Exception as exc: # pragma: no cover - depends on optional system package
|
|
self._load_error = exc
|
|
|
|
def start_input(
|
|
self, device_id: str | None = None, sample_rate: int = 16000, channels: int = 1
|
|
) -> None:
|
|
if self._sd is None:
|
|
raise ProviderError(
|
|
ErrorCode.AUDIO_INPUT_DEVICE_MISSING,
|
|
"sounddevice is not installed or cannot be imported",
|
|
False,
|
|
"sounddevice-transport",
|
|
"transport",
|
|
)
|
|
self.stop()
|
|
self._sample_rate = sample_rate
|
|
self._channels = channels
|
|
self._frame_id = 0
|
|
self._start_time = time.monotonic()
|
|
try:
|
|
self._stream = self._sd.RawInputStream(
|
|
samplerate=sample_rate,
|
|
channels=channels,
|
|
dtype="int16",
|
|
device=_coerce_device_id(device_id),
|
|
callback=self._on_input,
|
|
)
|
|
self._stream.start()
|
|
except Exception as exc:
|
|
self._stream = None
|
|
raise ProviderError(
|
|
ErrorCode.AUDIO_INPUT_DEVICE_MISSING,
|
|
f"cannot open sounddevice input stream: {exc}",
|
|
False,
|
|
"sounddevice-transport",
|
|
"transport",
|
|
) from exc
|
|
|
|
def read_frames(self, timeout_ms: int) -> list[AudioFrame]:
|
|
if self._stream is None:
|
|
return []
|
|
timeout_s = max(0, timeout_ms) / 1000
|
|
try:
|
|
frames = [self._queue.get(timeout=timeout_s)]
|
|
except queue.Empty:
|
|
return []
|
|
while True:
|
|
try:
|
|
frames.append(self._queue.get_nowait())
|
|
except queue.Empty:
|
|
break
|
|
return frames
|
|
|
|
def play_pcm(self, segment: AudioSegment, interrupt: bool = False) -> PlaybackResult:
|
|
if self._sd is None:
|
|
return PlaybackResult(
|
|
False,
|
|
0,
|
|
ProviderError(
|
|
ErrorCode.AUDIO_OUTPUT_DEVICE_MISSING,
|
|
"sounddevice is not installed or cannot be imported",
|
|
False,
|
|
"sounddevice-transport",
|
|
"transport",
|
|
),
|
|
)
|
|
if not segment.pcm:
|
|
return PlaybackResult(
|
|
False,
|
|
0,
|
|
ProviderError(
|
|
ErrorCode.AUDIO_STREAM_UNDERRUN,
|
|
"cannot play empty audio segment",
|
|
True,
|
|
"sounddevice-transport",
|
|
"transport",
|
|
),
|
|
)
|
|
if segment.metadata.get("format") in {"aiff", "wav", "mp3", "m4a", "aac"}:
|
|
return _play_file_bytes_with_afplay(segment)
|
|
try:
|
|
with self._sd.RawOutputStream(
|
|
samplerate=segment.sample_rate,
|
|
channels=segment.channels,
|
|
dtype="int16",
|
|
device=_coerce_device_id(self._output_device),
|
|
) as stream:
|
|
stream.write(segment.pcm)
|
|
except Exception as exc:
|
|
return PlaybackResult(
|
|
False,
|
|
0,
|
|
ProviderError(
|
|
ErrorCode.AUDIO_OUTPUT_DEVICE_MISSING,
|
|
f"cannot play through sounddevice output stream: {exc}",
|
|
False,
|
|
"sounddevice-transport",
|
|
"transport",
|
|
),
|
|
)
|
|
return PlaybackResult(True, segment.duration_ms)
|
|
|
|
def play_pcm_chunks(
|
|
self,
|
|
segment: AudioSegment,
|
|
*,
|
|
chunk_ms: int,
|
|
after_chunk: Callable[[AudioSegment, int], bool] | None = None,
|
|
should_stop: Callable[[], bool] | None = None,
|
|
) -> PlaybackResult:
|
|
if self._sd is None or segment.metadata.get("format") in {"aiff", "wav", "mp3", "m4a", "aac"}:
|
|
playback = self.play_pcm(segment)
|
|
if playback.error:
|
|
return playback
|
|
if after_chunk is not None:
|
|
after_chunk(segment, segment.duration_ms)
|
|
return playback
|
|
if not segment.pcm:
|
|
return PlaybackResult(
|
|
False,
|
|
0,
|
|
ProviderError(
|
|
ErrorCode.AUDIO_STREAM_UNDERRUN,
|
|
"cannot play empty audio segment",
|
|
True,
|
|
"sounddevice-transport",
|
|
"transport",
|
|
),
|
|
)
|
|
elapsed_ms = 0
|
|
try:
|
|
with self._sd.RawOutputStream(
|
|
samplerate=segment.sample_rate,
|
|
channels=segment.channels,
|
|
dtype="int16",
|
|
device=_coerce_device_id(self._output_device),
|
|
) as stream:
|
|
for chunk in _raw_pcm_chunks(segment, chunk_ms=chunk_ms):
|
|
if should_stop is not None and should_stop():
|
|
return PlaybackResult(True, elapsed_ms)
|
|
stream.write(chunk.pcm)
|
|
elapsed_ms += chunk.duration_ms
|
|
if should_stop is not None and should_stop():
|
|
return PlaybackResult(True, elapsed_ms)
|
|
if after_chunk is not None and after_chunk(chunk, elapsed_ms):
|
|
return PlaybackResult(True, elapsed_ms)
|
|
except Exception as exc:
|
|
return PlaybackResult(
|
|
False,
|
|
elapsed_ms,
|
|
ProviderError(
|
|
ErrorCode.AUDIO_OUTPUT_DEVICE_MISSING,
|
|
f"cannot play through sounddevice output stream: {exc}",
|
|
False,
|
|
"sounddevice-transport",
|
|
"transport",
|
|
),
|
|
)
|
|
return PlaybackResult(True, elapsed_ms)
|
|
|
|
def stop(self) -> None:
|
|
if self._stream is None:
|
|
return None
|
|
try:
|
|
self._stream.stop()
|
|
except Exception:
|
|
pass
|
|
try:
|
|
self._stream.close()
|
|
except Exception:
|
|
pass
|
|
self._stream = None
|
|
while not self._queue.empty():
|
|
try:
|
|
self._queue.get_nowait()
|
|
except queue.Empty:
|
|
break
|
|
return None
|
|
|
|
def flush_input(self) -> int:
|
|
count = 0
|
|
while not self._queue.empty():
|
|
try:
|
|
self._queue.get_nowait()
|
|
count += 1
|
|
except queue.Empty:
|
|
break
|
|
return count
|
|
|
|
def health(self) -> TransportHealth:
|
|
if self._sd is None:
|
|
return TransportHealth(False, False, "sounddevice unavailable")
|
|
try:
|
|
devices = list(self._sd.query_devices())
|
|
except Exception as exc:
|
|
return TransportHealth(False, False, f"sounddevice query failed: {exc}")
|
|
input_available = any(int(device.get("max_input_channels", 0)) > 0 for device in devices)
|
|
output_available = any(int(device.get("max_output_channels", 0)) > 0 for device in devices)
|
|
return TransportHealth(input_available, output_available, f"sounddevice devices={len(devices)}")
|
|
|
|
def device_report(self) -> dict[str, Any]:
|
|
health = self.health()
|
|
devices: list[dict[str, Any]] = []
|
|
if self._sd is not None:
|
|
try:
|
|
for index, device in enumerate(self._sd.query_devices()):
|
|
devices.append(
|
|
{
|
|
"index": index,
|
|
"name": str(device.get("name", "")),
|
|
"max_input_channels": int(device.get("max_input_channels", 0)),
|
|
"max_output_channels": int(device.get("max_output_channels", 0)),
|
|
"default_samplerate": float(device.get("default_samplerate", 0.0)),
|
|
}
|
|
)
|
|
except Exception:
|
|
devices = []
|
|
return {
|
|
"ok": health.input_available and health.output_available,
|
|
"input_available": health.input_available,
|
|
"output_available": health.output_available,
|
|
"message": health.message,
|
|
"devices": devices,
|
|
}
|
|
|
|
def _on_input(self, indata: Any, frames: int, _time_info: Any, status: Any) -> None:
|
|
timestamp_ms = int((time.monotonic() - self._start_time) * 1000)
|
|
duration_ms = int(frames / self._sample_rate * 1000) if self._sample_rate else 0
|
|
metadata: dict[str, Any] = {"duration_ms": duration_ms}
|
|
if status:
|
|
metadata["sounddevice_status"] = str(status)
|
|
frame = AudioFrame(
|
|
pcm=bytes(indata),
|
|
sample_rate=self._sample_rate,
|
|
channels=self._channels,
|
|
timestamp_ms=timestamp_ms,
|
|
frame_id=self._frame_id,
|
|
metadata=metadata,
|
|
)
|
|
self._frame_id += 1
|
|
if self._queue.full():
|
|
try:
|
|
self._queue.get_nowait()
|
|
except queue.Empty:
|
|
pass
|
|
self._queue.put_nowait(frame)
|
|
|
|
|
|
def sounddevice_device_report() -> dict[str, Any]:
|
|
return SoundDeviceAudioTransport().device_report()
|
|
|
|
|
|
def _coerce_device_id(device_id: str | None) -> int | str | None:
|
|
if device_id is None or device_id == "":
|
|
return None
|
|
return int(device_id) if str(device_id).isdigit() else device_id
|
|
|
|
|
|
def _play_file_bytes_with_afplay(segment: AudioSegment) -> PlaybackResult:
|
|
suffix = "." + str(segment.metadata.get("format", "aiff")).lstrip(".")
|
|
if not shutil.which("afplay"):
|
|
return PlaybackResult(
|
|
False,
|
|
0,
|
|
ProviderError(
|
|
ErrorCode.AUDIO_OUTPUT_DEVICE_MISSING,
|
|
"afplay is not available for file-backed audio playback",
|
|
False,
|
|
"afplay",
|
|
"transport",
|
|
),
|
|
)
|
|
with tempfile.NamedTemporaryFile(suffix=suffix) as handle:
|
|
handle.write(segment.pcm)
|
|
handle.flush()
|
|
try:
|
|
subprocess.run(["afplay", handle.name], check=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
|
|
except (FileNotFoundError, subprocess.CalledProcessError) as exc:
|
|
return PlaybackResult(
|
|
False,
|
|
0,
|
|
ProviderError(
|
|
ErrorCode.AUDIO_OUTPUT_DEVICE_MISSING,
|
|
f"afplay failed: {exc}",
|
|
False,
|
|
"afplay",
|
|
"transport",
|
|
),
|
|
)
|
|
return PlaybackResult(True, segment.duration_ms)
|
|
|
|
|
|
def _raw_pcm_chunks(segment: AudioSegment, *, chunk_ms: int) -> list[AudioSegment]:
|
|
if chunk_ms <= 0 or segment.duration_ms <= chunk_ms:
|
|
return [segment]
|
|
bytes_per_ms = max(1, int(segment.sample_rate * segment.channels * 2 / 1000))
|
|
chunk_bytes = max(2 * segment.channels, bytes_per_ms * chunk_ms)
|
|
chunk_bytes -= chunk_bytes % (2 * segment.channels)
|
|
chunks: list[AudioSegment] = []
|
|
offset = 0
|
|
start_ms = segment.start_time_ms
|
|
while offset < len(segment.pcm):
|
|
data = segment.pcm[offset : offset + chunk_bytes]
|
|
duration_ms = max(1, int(len(data) / bytes_per_ms))
|
|
chunks.append(
|
|
AudioSegment(
|
|
data,
|
|
segment.sample_rate,
|
|
segment.channels,
|
|
start_ms,
|
|
start_ms + duration_ms,
|
|
dict(segment.metadata),
|
|
)
|
|
)
|
|
offset += len(data)
|
|
start_ms += duration_ms
|
|
return chunks
|