sentence_buffer.py

"""
sentence_buffer.py
------------------
Accumulates LLM token stream and emits complete sentences for TTS synthesis.

Key design: We must detect REAL sentence boundaries without false-positives on
abbreviations (Dr., Mr., Inc.), decimal numbers (3.14), ellipses (...) etc.
Each emitted sentence is immediately handed off to TTS while the LLM continues
generating — this is the source of the streaming overlap latency savings.
"""

import re
from typing import List

# Abbreviations that must NOT trigger sentence split even when followed by a capital letter
ABBREVIATIONS = {
    "mr", "mrs", "ms", "dr", "prof", "sr", "jr", "rev", "gen", "sgt",
    "cpl", "pvt", "capt", "maj", "col", "lt", "cmdr", "adm",
    "inc", "corp", "ltd", "llc", "co", "dept", "est",
    "jan", "feb", "mar", "apr", "jun", "jul", "aug", "sep", "oct", "nov", "dec",
    "vs", "etc", "approx", "max", "min", "avg", "no", "vol", "fig",
    "st", "ave", "blvd", "rd",
}

# Pattern: a word ending with a period that looks like an abbreviation
_ABBREV_RE = re.compile(
    r'\b(' + '|'.join(re.escape(a) for a in ABBREVIATIONS) + r')\.$',
    re.IGNORECASE,
)

# Decimal numbers like 3.14 or $19.99 must not split
_DECIMAL_RE = re.compile(r'\d+\.$')

# Sentence-ending punctuation
_SENTENCE_END_RE = re.compile(r'[.!?]+')


def _is_false_boundary(text_before_dot: str) -> bool:
    """Return True if the period at the end of text_before_dot is NOT a sentence end."""
    stripped = text_before_dot.rstrip()
    if _ABBREV_RE.search(stripped):
        return True
    if _DECIMAL_RE.search(stripped):
        return True
    return False


class SentenceBuffer:
    """
    Feed LLM tokens one at a time; call flush() at stream end.

    Usage
    -----
    buf = SentenceBuffer(min_length=12)
    for token in llm_stream():
        sentences = buf.add(token)
        for s in sentences:
            audio = tts.synthesize(s)   # start immediately
    final = buf.flush()
    if final:
        audio = tts.synthesize(final)
    """

    def __init__(self, min_length: int = 10):
        """
        min_length : int
            Minimum character count before a boundary is accepted.
            Avoids firing TTS on single-word fragments like "Hi."
        """
        self.min_length = min_length
        self._buf = ""

    # ------------------------------------------------------------------
    def add(self, token: str) -> List[str]:
        """
        Append a token and return a list of complete sentences ready for TTS.
        Typically returns [] or a one-element list.
        """
        self._buf += token
        return self._extract()

    def flush(self) -> str:
        """Return whatever is left in the buffer (end of stream)."""
        remaining = self._buf.strip()
        self._buf = ""
        return remaining

    def reset(self):
        self._buf = ""

    # ------------------------------------------------------------------
    def _extract(self) -> List[str]:
        sentences = []

        while True:
            match = _SENTENCE_END_RE.search(self._buf)
            if not match:
                break

            end_pos = match.end()
            candidate = self._buf[:end_pos].strip()

            if len(candidate) < self.min_length:
                break  # Too short — wait for more tokens

            # Check for false boundary
            before_punct = self._buf[:match.start()]
            if _is_false_boundary(before_punct):
                # Advance past this dot and keep scanning
                self._buf = self._buf[end_pos:]
                if sentences:
                    sentences[-1] += candidate  # Merge into previous
                else:
                    # Keep in buffer with content that was before
                    self._buf = candidate + " " + self._buf
                    break
                continue

            # Real sentence boundary — check there's enough after the punctuation
            after = self._buf[end_pos:].lstrip()

            # If the very next char is another sentence-ender, keep accumulating
            if after and after[0] in '.!?':
                break

            sentences.append(candidate)
            self._buf = self._buf[end_pos:].lstrip()

        return sentences