splitbit_llm/voice/tts_output.py

"""TTS-Ready Output — formats model output for text-to-speech engines.

Strips markdown, code blocks, and special characters before TTS.
Number normalization ("100" → "one hundred").
Abbreviation expansion ("API" → "A-P-I").
Sentence queue with priority — first sentence goes to TTS immediately.
"""

from __future__ import annotations

import logging
import re
from typing import Any

logger = logging.getLogger(__name__)


# Common abbreviations to expand for TTS
ABBREVIATIONS = {
    "api": "A P I",
    "url": "U R L",
    "cpu": "C P U",
    "gpu": "G P U",
    "ram": "R A M",
    "ssd": "S S D",
    "html": "H T M L",
    "css": "C S S",
    "sql": "S Q L",
    "json": "J S O N",
    "xml": "X M L",
    "ai": "A I",
    "ml": "M L",
    "llm": "L L M",
    "gpt": "G P T",
    "ui": "U I",
    "ux": "U X",
    "os": "O S",
    "ip": "I P",
    "dns": "D N S",
    "vpn": "V P N",
    "pdf": "P D F",
    "csv": "C S V",
    "png": "P N G",
    "jpg": "J P G",
    "gif": "G I F",
    "npm": "N P M",
    "pip": "P I P",
    "git": "G I T",
    "ssh": "S S H",
    "tcp": "T C P",
    "udp": "U D P",
    "http": "H T T P",
    "https": "H T T P S",
}

# Number words for normalization
ONES = ["zero", "one", "two", "three", "four", "five", "six", "seven", "eight", "nine"]
TEENS = ["ten", "eleven", "twelve", "thirteen", "fourteen", "fifteen", "sixteen",
         "seventeen", "eighteen", "nineteen"]
TENS = ["", "", "twenty", "thirty", "forty", "fifty", "sixty", "seventy", "eighty", "ninety"]


def number_to_words(n: int) -> str:
    """Convert a number to its spoken word form."""
    if n == 0:
        return "zero"
    if n < 0:
        return "negative " + number_to_words(-n)

    parts: list[str] = []

    if n >= 1000000:
        parts.append(f"{number_to_words(n // 1000000)} million")
        n %= 1000000
    if n >= 1000:
        parts.append(f"{number_to_words(n // 1000)} thousand")
        n %= 1000
    if n >= 100:
        parts.append(f"{ONES[n // 100]} hundred")
        n %= 100
    if n >= 20:
        tens_part = n // 10
        ones_part = n % 10
        parts.append(TENS[tens_part])
        if ones_part > 0:
            parts[-1] += f" {ONES[ones_part]}"
    elif n >= 10:
        parts.append(TEENS[n - 10])
    elif n > 0:
        parts.append(ONES[n])

    return " ".join(parts)


class TTSOutputFormatter:
    """Formats model output for text-to-speech engines.

    - Strips markdown, code blocks, special characters
    - Normalizes numbers to words
    - Expands abbreviations
    - Cleans up whitespace
    """

    def __init__(self) -> None:
        self._stats = {
            "sentences_formatted": 0,
            "characters_stripped": 0,
        }

    def format(self, text: str) -> str:
        """Format text for TTS.

        Args:
            text: raw model output
        Returns:
            TTS-ready text
        """
        original_len = len(text)

        # Strip code blocks
        text = re.sub(r"```[\s\S]*?```", " (code block) ", text)
        text = re.sub(r"`[^`]+`", " (code) ", text)

        # Strip markdown formatting
        text = re.sub(r"\*\*([^*]+)\*\*", r"\1", text)  # bold
        text = re.sub(r"\*([^*]+)\*", r"\1", text)  # italic
        text = re.sub(r"__([^_]+)__", r"\1", text)  # bold
        text = re.sub(r"~~([^~]+)~~", r"\1", text)  # strikethrough
        text = re.sub(r"^#{1,6}\s+", "", text, flags=re.MULTILINE)  # headers
        text = re.sub(r"^>\s+", "", text, flags=re.MULTILINE)  # blockquote
        text = re.sub(r"^[-*]\s+", "", text, flags=re.MULTILINE)  # list items
        text = re.sub(r"\[([^\]]+)\]\([^)]+\)", r"\1", text)  # links

        # Strip special characters that TTS can't handle well
        text = text.replace("|", " ")
        text = text.replace("~", " ")
        text = text.replace("`", "")
        text = text.replace("*", "")
        text = text.replace("#", "")
        text = text.replace("_", " ")

        # Expand abbreviations
        words = text.split()
        expanded: list[str] = []
        for word in words:
            clean = re.sub(r"[^a-zA-Z]", "", word).lower()
            if clean in ABBREVIATIONS:
                # Preserve surrounding punctuation
                prefix = re.match(r"[^a-zA-Z]*", word).group(0)
                suffix = re.match(r".*?[^a-zA-Z]*$", word).group(0)
                # Actually just replace the word
                expanded.append(ABBREVIATIONS[clean])
            else:
                expanded.append(word)
        text = " ".join(expanded)

        # Normalize numbers
        text = re.sub(r"\b(\d+)\b", lambda m: number_to_words(int(m.group(1))), text)

        # Clean up whitespace
        text = re.sub(r"\s+", " ", text).strip()

        self._stats["sentences_formatted"] += 1
        self._stats["characters_stripped"] += max(0, original_len - len(text))

        return text

    def get_stats(self) -> dict[str, Any]:
        return {**self._stats}