"""Thai TTS — interface, chunker, and barge-in queue (REQ-006/REQ-007). * :class:`TTSEngine` — interface. ``mock`` emits deterministic audio frames so the streaming path is testable on a CPU-only host. * :func:`chunk_text` — Thai-aware sentence/phrase splitting so TTS starts before the full answer exists (REQ-006, "first audio before answer done"). * :class:`TTSQueue` — async producer/consumer with :meth:`cancel` for barge-in (REQ-007): a user interrupt drops the rest of the audio without a race. """ from __future__ import annotations import asyncio import re from abc import ABC, abstractmethod # --------------------------------------------------------------------------- # Chunking (REQ-006) # --------------------------------------------------------------------------- # Thai sentence delimiters: ASCII punctuation plus the Thai text-marker ฯ and # the Arabic question mark Thai text commonly uses. We deliberately do NOT # split on Thai vowels/tone marks — doing so breaks words in half. _SENT_END = re.compile(r"[.!?。!?\u0e3f]") _MAX_CHUNK = 60 # hard ceiling so TTS starts promptly def chunk_text(text: str) -> list[str]: """Split Thai text into speakable chunks. Rules: split on sentence punctuation, **keeping the delimiter attached** to its fragment (JaiTTS needs the final "?" / "." for intonation); cap very long fragments on commas / natural pauses and then hard-slice; drop empty pieces. Deterministic and side-effect-free (unit-testable). """ text = (text or "").strip() if not text: return [] pieces: list[str] = [] pos = 0 spans: list[tuple[int, int]] = [] for m in _SENT_END.finditer(text): spans.append((pos, m.end())) pos = m.end() if pos < len(text): spans.append((pos, len(text))) for start, end in spans: frag = text[start:end].strip() if not frag: continue # Second pass: cap very long fragments on commas / "แล้ว" pauses. if len(frag) > _MAX_CHUNK: sub = re.split(r"[,,;;]+", frag) pieces.extend(s.strip() for s in sub if s.strip()) else: pieces.append(frag) # Hard cap: any piece still over _MAX_CHUNK is sliced, so TTS always # starts promptly no matter the input shape. capped: list[str] = [] for p in pieces: while len(p) > _MAX_CHUNK: capped.append(p[:_MAX_CHUNK]) p = p[_MAX_CHUNK:] if p: capped.append(p) pieces = capped # Merge only true micro-fragments (<= 4 chars) so we don't fire a TTS # call per token; sentence-level fragments stay separate. merged: list[str] = [] for p in pieces: if merged and len(merged[-1]) <= 4 and len(merged[-1]) + 1 + len(p) <= 24: merged[-1] = f"{merged[-1]} {p}" else: merged.append(p) return [m for m in merged if m] # --------------------------------------------------------------------------- # Engine (REQ-006) # --------------------------------------------------------------------------- class TTSEngine(ABC): async def synthesize(self, text: str): """Yield PCM audio chunks for ``text``. Streaming: yield as produced.""" raise NotImplementedError class MockTTS(TTSEngine): """Emits 16-byte PCM frames — enough to prove streaming + barge-in.""" async def synthesize(self, text: str): # Emit a couple of frames per chunk to simulate a stream. for _ in range(2): yield b"\x00\x01" * 8 await asyncio.sleep(0) class JaiTTSEngine(TTSEngine): def __init__(self, endpoint: str, voice: str = "default", sample_rate: int = 24000): self.endpoint = endpoint self.voice = voice self.sample_rate = sample_rate async def synthesize(self, text: str): import httpx # lazy async with httpx.AsyncClient() as client: async with client.stream( "POST", self.endpoint, json={"text": text, "voice": self.voice}, headers={"Accept": "application/octet-stream"}, ) as r: r.raise_for_status() async for chunk in r.aiter_bytes(4096): yield chunk # --------------------------------------------------------------------------- # Barge-in queue (REQ-007) # --------------------------------------------------------------------------- class TTSQueue: """Bounded queue of audio chunks with a clean cancel for barge-in. The consumer (audio-out) reads :meth:`get`; :meth:`cancel` drops all pending chunks and signals the consumer to stop — used when the user barges in while the AI is speaking. """ def __init__(self, maxsize: int = 64) -> None: self._q: asyncio.Queue = asyncio.Queue(maxsize=maxsize) self._cancelled = False self._done = asyncio.Event() async def put(self, chunk: bytes) -> None: if self._cancelled: return await self._q.put(chunk) async def get(self) -> bytes: item = await self._q.get() if item is _SENTINEL: self._done.set() return item def cancel(self) -> int: """Drop everything pending. Returns how many chunks were dropped.""" self._cancelled = True dropped = 0 while not self._q.empty(): try: self._q.get_nowait() dropped += 1 except asyncio.QueueEmpty: break self._q.put_nowait(_SENTINEL) self._done.set() return dropped @property def cancelled(self) -> bool: return self._cancelled def pending(self) -> int: return self._q.qsize() class _Sentinel: pass _SENTINEL = _Sentinel() def make_tts(cfg) -> TTSEngine: if cfg.engine == "mock": return MockTTS() if cfg.engine == "jaitts": return JaiTTSEngine(cfg.jaitts_endpoint, cfg.voice, cfg.sample_rate) raise ValueError(f"unknown tts engine: {cfg.engine}")