call-center/services/ai_voice_runtime_service/audiosocket.py

from __future__ import annotations

import asyncio
import io
import struct
import uuid
import wave
from dataclasses import dataclass

from services.shared.audioop_compat import audioop


AUDIO_SOCKET_PACKET_HANGUP = 0x00
AUDIO_SOCKET_PACKET_UUID = 0x01
AUDIO_SOCKET_PACKET_DTMF = 0x03
AUDIO_SOCKET_PACKET_PCM16 = 0x10


def normalize_media_uuid(value: str | bytes) -> str:
    if isinstance(value, bytes):
        try:
            return str(uuid.UUID(bytes=value)).lower()
        except (TypeError, ValueError, AttributeError):
            return str(uuid.UUID(value.decode("utf-8").strip())).lower()
    return str(uuid.UUID(str(value).strip())).lower()


def encode_packet(packet_type: int, payload: bytes = b"") -> bytes:
    return struct.pack("!BH", packet_type & 0xFF, len(payload)) + payload


def encode_audio_packet(pcm_bytes: bytes) -> bytes:
    return encode_packet(AUDIO_SOCKET_PACKET_PCM16, pcm_bytes)


async def read_packet(reader: asyncio.StreamReader, *, timeout: float) -> tuple[int, bytes]:
    header = await asyncio.wait_for(reader.readexactly(3), timeout=timeout)
    packet_type, payload_length = struct.unpack("!BH", header)
    payload = b""
    if payload_length:
        payload = await asyncio.wait_for(reader.readexactly(payload_length), timeout=timeout)
    return packet_type, payload


def pcm16le_to_wav_bytes(
    pcm_bytes: bytes,
    *,
    sample_rate_hz: int,
    channels: int = 1,
    sample_width_bytes: int = 2,
) -> bytes:
    handle = io.BytesIO()
    with wave.open(handle, "wb") as wav_file:
        wav_file.setnchannels(channels)
        wav_file.setsampwidth(sample_width_bytes)
        wav_file.setframerate(sample_rate_hz)
        wav_file.writeframes(pcm_bytes)
    return handle.getvalue()


def resample_pcm16le(
    pcm_bytes: bytes,
    *,
    input_rate_hz: int,
    output_rate_hz: int,
    channels: int = 1,
    sample_width_bytes: int = 2,
) -> bytes:
    if not pcm_bytes or input_rate_hz == output_rate_hz:
        return pcm_bytes
    mono_bytes = pcm_bytes
    if channels == 2:
        mono_bytes = audioop.tomono(pcm_bytes, sample_width_bytes, 0.5, 0.5)
    converted, _ = audioop.ratecv(
        mono_bytes,
        sample_width_bytes,
        1,
        input_rate_hz,
        output_rate_hz,
        None,
    )
    return converted


def chunk_audio(pcm_bytes: bytes, *, frame_bytes: int) -> list[bytes]:
    if frame_bytes <= 0:
        raise ValueError("frame_bytes must be positive")
    frames = [pcm_bytes[index : index + frame_bytes] for index in range(0, len(pcm_bytes), frame_bytes)]
    if frames and len(frames[-1]) < frame_bytes:
        frames[-1] = frames[-1] + (b"\x00" * (frame_bytes - len(frames[-1])))
    return frames


@dataclass(slots=True)
class EnergyVADResult:
    speech_started: bool = False
    utterance_pcm: bytes | None = None


class EnergyVAD:
    def __init__(
        self,
        *,
        frame_ms: int,
        min_speech_ms: int,
        trailing_silence_ms: int,
        max_turn_ms: int,
        rms_threshold: int = 250,
    ) -> None:
        self._frame_ms = max(frame_ms, 10)
        self._min_speech_ms = max(min_speech_ms, self._frame_ms)
        self._trailing_silence_ms = max(trailing_silence_ms, self._frame_ms)
        self._max_turn_ms = max(max_turn_ms, self._frame_ms)
        self._rms_threshold = max(rms_threshold, 1)
        self.reset()

    def reset(self) -> None:
        self._pending_pcm = bytearray()
        self._utterance_pcm = bytearray()
        self._speech_ms = 0
        self._silence_ms = 0
        self._active = False

    def feed(self, pcm_frame: bytes) -> EnergyVADResult:
        if not pcm_frame:
            return EnergyVADResult()

        is_speech = audioop.rms(pcm_frame, 2) >= self._rms_threshold
        result = EnergyVADResult()

        if not self._active:
            if is_speech:
                self._pending_pcm.extend(pcm_frame)
                self._speech_ms += self._frame_ms
                if self._speech_ms >= self._min_speech_ms:
                    self._active = True
                    self._utterance_pcm.extend(self._pending_pcm)
                    self._pending_pcm.clear()
                    self._silence_ms = 0
                    result.speech_started = True
            else:
                self._pending_pcm.clear()
                self._speech_ms = 0
            return result

        self._utterance_pcm.extend(pcm_frame)
        if is_speech:
            self._silence_ms = 0
        else:
            self._silence_ms += self._frame_ms

        if self._silence_ms >= self._trailing_silence_ms or len(self._utterance_pcm) >= self._max_pcm_bytes:
            result.utterance_pcm = bytes(self._utterance_pcm)
            self.reset()
        return result

    @property
    def _max_pcm_bytes(self) -> int:
        samples_per_ms = 8
        bytes_per_sample = 2
        return self._max_turn_ms * samples_per_ms * bytes_per_sample

    @property
    def is_active(self) -> bool:
        return self._active

    def snapshot_utterance_pcm(self) -> bytes:
        if not self._active or not self._utterance_pcm:
            return b""
        return bytes(self._utterance_pcm)