realtime_voice_service/providers/base.py

from __future__ import annotations

import asyncio
import math
import re
import struct
from abc import ABC, abstractmethod
from collections.abc import AsyncGenerator
from collections.abc import AsyncIterable
from collections.abc import Awaitable
from collections.abc import Callable
from dataclasses import dataclass


PartialTranscriptCallback = Callable[[str], Awaitable[None] | None]


class BaseSTTStream(ABC):
    @abstractmethod
    async def push_audio(self, audio_chunk: bytes) -> None:
        raise NotImplementedError

    @abstractmethod
    async def finish(self) -> str:
        raise NotImplementedError

    @abstractmethod
    async def cancel(self) -> None:
        raise NotImplementedError


class BaseSTT(ABC):
    @abstractmethod
    async def transcribe(self, audio_bytes: bytes) -> str:
        raise NotImplementedError

    async def start_stream(
        self,
        *,
        partial_callback: PartialTranscriptCallback | None = None,
    ) -> BaseSTTStream | None:
        del partial_callback
        return None


class BaseLLM(ABC):
    @abstractmethod
    async def generate_stream(self, text: str, context: list) -> AsyncGenerator[LLMStreamEvent, None]:
        raise NotImplementedError


class BaseTTS(ABC):
    @abstractmethod
    async def synthesize_stream(self, text_stream: AsyncIterable[str]) -> AsyncGenerator[bytes, None]:
        raise NotImplementedError


@dataclass(slots=True)
class LLMStreamEvent:
    type: str
    content: str | None = None
    name: str | None = None
    tool_call_id: str | None = None


class MockSTT(BaseSTT):
    def __init__(
        self,
        *,
        latency_ms: int = 40,
        sample_rate_hz: int = 16000,
        scripted_transcripts: list[str] | None = None,
    ) -> None:
        self._latency_ms = max(latency_ms, 0)
        self._sample_rate_hz = sample_rate_hz
        self._scripted_transcripts = list(scripted_transcripts or [])
        self._call_count = 0

    async def transcribe(self, audio_bytes: bytes) -> str:
        await asyncio.sleep(self._latency_ms / 1000.0)
        self._call_count += 1
        if self._scripted_transcripts:
            return self._scripted_transcripts.pop(0)
        duration_ms = int(((len(audio_bytes) // 2) / float(self._sample_rate_hz)) * 1000.0)
        return f"mock user utterance {self._call_count} ({duration_ms} ms)"

    async def start_stream(
        self,
        *,
        partial_callback: PartialTranscriptCallback | None = None,
    ) -> BaseSTTStream | None:
        return _MockSTTStream(parent=self, partial_callback=partial_callback)


class MockLLM(BaseLLM):
    def __init__(
        self,
        *,
        token_delay_ms: int = 35,
        scripted_responses: list[str] | None = None,
    ) -> None:
        self._token_delay_ms = max(token_delay_ms, 0)
        self._scripted_responses = list(scripted_responses or [])

    async def generate_stream(self, text: str, context: list) -> AsyncGenerator[LLMStreamEvent, None]:
        del context
        response_text = (
            self._scripted_responses.pop(0)
            if self._scripted_responses
            else f"Mock reply: {text}"
        )
        chunks = re.findall(r"\S+\s*", response_text) or [response_text]
        for chunk in chunks:
            await asyncio.sleep(self._token_delay_ms / 1000.0)
            yield LLMStreamEvent(type="text", content=chunk)


class MockTTS(BaseTTS):
    def __init__(
        self,
        *,
        sample_rate_hz: int = 16000,
        chunk_duration_ms: int = 40,
        chunk_delay_ms: int = 15,
        tone_hz: float = 440.0,
        amplitude: int = 9000,
        milliseconds_per_word: int = 140,
    ) -> None:
        self._sample_rate_hz = sample_rate_hz
        self._chunk_duration_ms = max(chunk_duration_ms, 20)
        self._chunk_delay_ms = max(chunk_delay_ms, 0)
        self._tone_hz = tone_hz
        self._amplitude = amplitude
        self._milliseconds_per_word = max(milliseconds_per_word, 40)

    async def synthesize_stream(self, text_stream: AsyncIterable[str]) -> AsyncGenerator[bytes, None]:
        async for text in text_stream:
            normalized = text.strip()
            if not normalized:
                continue
            word_count = max(len(normalized.split()), 1)
            total_ms = min(word_count * self._milliseconds_per_word, 2400)
            total_samples = max(int(self._sample_rate_hz * (total_ms / 1000.0)), 1)
            chunk_samples = max(int(self._sample_rate_hz * (self._chunk_duration_ms / 1000.0)), 1)

            for start in range(0, total_samples, chunk_samples):
                end = min(start + chunk_samples, total_samples)
                pcm = bytearray()
                for index in range(start, end):
                    angle = 2.0 * math.pi * self._tone_hz * (index / float(self._sample_rate_hz))
                    sample = int(self._amplitude * math.sin(angle))
                    pcm.extend(struct.pack("<h", sample))
                await asyncio.sleep(self._chunk_delay_ms / 1000.0)
                yield bytes(pcm)


class _MockSTTStream(BaseSTTStream):
    def __init__(
        self,
        *,
        parent: MockSTT,
        partial_callback: PartialTranscriptCallback | None = None,
    ) -> None:
        self._parent = parent
        self._partial_callback = partial_callback
        self._audio_buffer = bytearray()
        self._cancelled = False
        self._partial_emitted = False

    async def push_audio(self, audio_chunk: bytes) -> None:
        if self._cancelled or not audio_chunk:
            return
        self._audio_buffer.extend(audio_chunk)
        if self._partial_callback is None or self._partial_emitted:
            return
        duration_ms = int(((len(self._audio_buffer) // 2) / float(self._parent._sample_rate_hz)) * 1000.0)
        if duration_ms < 600:
            return
        self._partial_emitted = True
        partial = f"mock partial utterance {self._parent._call_count + 1}"
        maybe_awaitable = self._partial_callback(partial)
        if maybe_awaitable is not None:
            await maybe_awaitable

    async def finish(self) -> str:
        if self._cancelled:
            return ""
        return await self._parent.transcribe(bytes(self._audio_buffer))

    async def cancel(self) -> None:
        self._cancelled = True
        self._audio_buffer.clear()