107 lines
3.6 KiB
Python
107 lines
3.6 KiB
Python
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import math
|
|
import re
|
|
import struct
|
|
from abc import ABC, abstractmethod
|
|
from collections.abc import AsyncGenerator
|
|
|
|
|
|
class BaseSTT(ABC):
|
|
@abstractmethod
|
|
async def transcribe(self, audio_bytes: bytes) -> str:
|
|
raise NotImplementedError
|
|
|
|
|
|
class BaseLLM(ABC):
|
|
@abstractmethod
|
|
async def generate_stream(self, text: str, context: list) -> AsyncGenerator[str, None]:
|
|
raise NotImplementedError
|
|
|
|
|
|
class BaseTTS(ABC):
|
|
@abstractmethod
|
|
async def synthesize_stream(self, text: str) -> AsyncGenerator[bytes, None]:
|
|
raise NotImplementedError
|
|
|
|
|
|
class MockSTT(BaseSTT):
|
|
def __init__(
|
|
self,
|
|
*,
|
|
latency_ms: int = 40,
|
|
sample_rate_hz: int = 8000,
|
|
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)"
|
|
|
|
|
|
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[str, 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 chunk
|
|
|
|
|
|
class MockTTS(BaseTTS):
|
|
def __init__(
|
|
self,
|
|
*,
|
|
sample_rate_hz: int = 8000,
|
|
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: str) -> AsyncGenerator[bytes, None]:
|
|
word_count = max(len(text.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)
|