205 lines
7.1 KiB
Python
205 lines
7.1 KiB
Python
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import logging
|
|
import os
|
|
import time
|
|
from abc import ABC, abstractmethod
|
|
|
|
from realtime_voice_service.core.audio_pacer import AudioPacer
|
|
|
|
|
|
LOGGER = logging.getLogger("uvicorn.error")
|
|
|
|
|
|
def _float_env(name: str, default: float) -> float:
|
|
raw = os.getenv(name)
|
|
if raw is None:
|
|
return default
|
|
try:
|
|
return float(str(raw).strip())
|
|
except ValueError:
|
|
return default
|
|
|
|
|
|
class BaseMediaTransport(ABC):
|
|
def __init__(
|
|
self,
|
|
*,
|
|
transport_id: str,
|
|
sample_rate_hz: int = 16000,
|
|
frame_duration_ms: int = 20,
|
|
) -> None:
|
|
self._transport_id = transport_id
|
|
self._sample_rate_hz = sample_rate_hz
|
|
self._frame_duration_ms = frame_duration_ms
|
|
self._audio_pacer = AudioPacer(frame_bytes=self.frame_bytes)
|
|
self._send_lock = asyncio.Lock()
|
|
self._next_frame_monotonic: float | None = None
|
|
self._send_generation = 0
|
|
self._audio_log_interval_seconds = max(
|
|
_float_env("REALTIME_VOICE_AUDIO_LOG_INTERVAL_SECONDS", 1.0),
|
|
0.1,
|
|
)
|
|
self._last_tx_summary_monotonic = time.perf_counter()
|
|
self._tx_chunk_count = 0
|
|
self._tx_input_bytes = 0
|
|
self._tx_frame_count = 0
|
|
self._tx_frame_bytes = 0
|
|
self._tx_flush_count = 0
|
|
self._tx_clear_count = 0
|
|
|
|
@property
|
|
def transport_id(self) -> str:
|
|
return self._transport_id
|
|
|
|
@property
|
|
def sample_rate_hz(self) -> int:
|
|
return self._sample_rate_hz
|
|
|
|
@property
|
|
def frame_duration_ms(self) -> int:
|
|
return self._frame_duration_ms
|
|
|
|
@property
|
|
def frame_bytes(self) -> int:
|
|
return int((self.sample_rate_hz * self.frame_duration_ms / 1000.0) * 2)
|
|
|
|
@property
|
|
@abstractmethod
|
|
def protocol(self) -> str:
|
|
raise NotImplementedError
|
|
|
|
@abstractmethod
|
|
async def receive_audio(self) -> bytes | None:
|
|
raise NotImplementedError
|
|
|
|
async def send_audio(self, audio_chunk: bytes) -> None:
|
|
if not audio_chunk:
|
|
return
|
|
async with self._send_lock:
|
|
generation = self._send_generation
|
|
frames = self._audio_pacer.push(audio_chunk)
|
|
self._tx_chunk_count += 1
|
|
self._tx_input_bytes += len(audio_chunk)
|
|
self._tx_frame_count += len(frames)
|
|
self._tx_frame_bytes += sum(len(frame) for frame in frames)
|
|
self._maybe_log_tx_summary(generation=generation)
|
|
for frame in frames:
|
|
if generation != self._send_generation:
|
|
return
|
|
await self._pace_and_send_frame(frame, generation=generation)
|
|
|
|
async def flush_audio(self, *, pad_final_frame: bool = True) -> None:
|
|
async with self._send_lock:
|
|
generation = self._send_generation
|
|
buffered_before = self._audio_pacer.buffered_bytes
|
|
frames = self._audio_pacer.flush(pad_final_frame=pad_final_frame)
|
|
self._tx_flush_count += 1
|
|
self._tx_frame_count += len(frames)
|
|
self._tx_frame_bytes += sum(len(frame) for frame in frames)
|
|
LOGGER.info(
|
|
"transport %s flush_audio protocol=%s generation=%s frames=%s frame_bytes=%s "
|
|
"buffered_before=%s pad_final_frame=%s",
|
|
self.transport_id,
|
|
self.protocol,
|
|
generation,
|
|
len(frames),
|
|
sum(len(frame) for frame in frames),
|
|
buffered_before,
|
|
pad_final_frame,
|
|
)
|
|
self._maybe_log_tx_summary(generation=generation, force=True)
|
|
for frame in frames:
|
|
if generation != self._send_generation:
|
|
return
|
|
await self._pace_and_send_frame(frame, generation=generation)
|
|
|
|
def clear_buffer(self) -> None:
|
|
previous_generation = self._send_generation
|
|
buffered_before = self._audio_pacer.buffered_bytes
|
|
self._send_generation += 1
|
|
self._tx_clear_count += 1
|
|
self._audio_pacer.clear()
|
|
self._next_frame_monotonic = None
|
|
LOGGER.info(
|
|
"transport %s clear_buffer protocol=%s generation=%s->%s buffered_bytes=%s "
|
|
"tx_chunks=%s tx_frames=%s tx_bytes=%s clears=%s",
|
|
self.transport_id,
|
|
self.protocol,
|
|
previous_generation,
|
|
self._send_generation,
|
|
buffered_before,
|
|
self._tx_chunk_count,
|
|
self._tx_frame_count,
|
|
self._tx_frame_bytes,
|
|
self._tx_clear_count,
|
|
)
|
|
|
|
def discard_audio_buffer(self) -> None:
|
|
self.clear_buffer()
|
|
|
|
@abstractmethod
|
|
async def close(self) -> None:
|
|
raise NotImplementedError
|
|
|
|
@abstractmethod
|
|
async def _send_frame(self, frame: bytes) -> None:
|
|
raise NotImplementedError
|
|
|
|
async def _pace_and_send_frame(self, frame: bytes, *, generation: int) -> None:
|
|
if generation != self._send_generation:
|
|
return
|
|
frame_duration_seconds = self.frame_duration_ms / 1000.0
|
|
now = time.perf_counter()
|
|
if (
|
|
self._next_frame_monotonic is None
|
|
or now > (self._next_frame_monotonic + (frame_duration_seconds * 4.0))
|
|
):
|
|
self._next_frame_monotonic = now
|
|
sleep_for = self._next_frame_monotonic - now
|
|
if sleep_for > 0:
|
|
await asyncio.sleep(sleep_for)
|
|
if generation != self._send_generation:
|
|
return
|
|
send_started_monotonic = time.perf_counter()
|
|
await self._send_frame(frame)
|
|
send_duration_ms = int((time.perf_counter() - send_started_monotonic) * 1000.0)
|
|
if send_duration_ms > (self.frame_duration_ms * 2):
|
|
LOGGER.warning(
|
|
"transport %s slow_frame_send protocol=%s generation=%s duration_ms=%s frame_bytes=%s",
|
|
self.transport_id,
|
|
self.protocol,
|
|
generation,
|
|
send_duration_ms,
|
|
len(frame),
|
|
)
|
|
if generation != self._send_generation:
|
|
return
|
|
baseline = max(time.perf_counter(), self._next_frame_monotonic)
|
|
self._next_frame_monotonic = baseline + frame_duration_seconds
|
|
|
|
def _maybe_log_tx_summary(self, *, generation: int, force: bool = False) -> None:
|
|
now = time.perf_counter()
|
|
if not force and (now - self._last_tx_summary_monotonic) < self._audio_log_interval_seconds:
|
|
return
|
|
self._last_tx_summary_monotonic = now
|
|
LOGGER.info(
|
|
"transport %s audio_tx_summary protocol=%s generation=%s sample_rate=%s frame_ms=%s "
|
|
"frame_bytes=%s chunks=%s input_bytes=%s frames=%s frame_payload_bytes=%s "
|
|
"pacer_buffered=%s flushes=%s clears=%s",
|
|
self.transport_id,
|
|
self.protocol,
|
|
generation,
|
|
self.sample_rate_hz,
|
|
self.frame_duration_ms,
|
|
self.frame_bytes,
|
|
self._tx_chunk_count,
|
|
self._tx_input_bytes,
|
|
self._tx_frame_count,
|
|
self._tx_frame_bytes,
|
|
self._audio_pacer.buffered_bytes,
|
|
self._tx_flush_count,
|
|
self._tx_clear_count,
|
|
)
|