esp32-server/app/tts/__init__.py

"""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}")