vibevoice-tts

#!/usr/bin/env python3
"""
mock/text_preprocessor.py — Standalone mock of the vibevoice-tts formatter.

Reads JSON from stdin with shape { "text": str, "options": { ... } } and
writes a formatted JSON result to stdout with shape:

{
  "segments": [{ "text": str, "directives": [...] }],
  "ssml": "<speak>...</speak>",
  "originalText": str,
  "appliedOptions": { ... },
  "warnings": [...]
}

This is a MOCK for testing — it performs basic inline markup parsing
to demonstrate the format without depending on Node.js or @lmstudio/sdk.

Usage:
    echo '{"text": "Hello **world**", "options": {}}' | python mock/text_preprocessor.py
"""

import json
import re
import sys
from typing import Any


# ---------------------------------------------------------------------------
# Regex patterns (mirroring the TypeScript formatter)
# ---------------------------------------------------------------------------

RE_BOLD = re.compile(r"\*\*(.+?)\*\*")
RE_ITALIC = re.compile(r"(?<!\*)\*(?!\*)(.+?)(?<!\*)\*(?!\*)")
RE_SPEED = re.compile(r"\[speed:\s*([0-9]+(?:\.[0-9]+)?)\s*x\]", re.IGNORECASE)
RE_PITCH = re.compile(r"\[pitch:\s*(low|normal|high)\s*\]", re.IGNORECASE)
RE_PAUSE = re.compile(r"\[pause:\s*([0-9]+)\s*ms\]", re.IGNORECASE)


# ---------------------------------------------------------------------------
# Formatting option defaults and validation
# ---------------------------------------------------------------------------

SPEED_MIN = 0.5
SPEED_MAX = 2.0


def validate_options(raw: dict[str, Any]) -> tuple[dict[str, Any], list[str]]:
    """Validate and clamp formatting options. Returns (applied, warnings)."""
    warnings: list[str] = []
    result: dict[str, Any] = {}

    if raw.get("bold") is True:
        result["bold"] = True
    if raw.get("italic") is True:
        result["italic"] = True

    speed = raw.get("speed")
    if speed is not None:
        try:
            v = float(speed)
            if v < SPEED_MIN or v > SPEED_MAX:
                clamped = max(SPEED_MIN, min(SPEED_MAX, v))
                warnings.append(
                    f"Speed {v} outside [{SPEED_MIN}, {SPEED_MAX}], clamped to {clamped}"
                )
                result["speed"] = clamped
            else:
                result["speed"] = v
        except (ValueError, TypeError):
            warnings.append(f"Invalid speed value, using default 1.0")
            result["speed"] = 1.0

    pitch = raw.get("pitch")
    if pitch is not None:
        if pitch not in ("low", "normal", "high"):
            warnings.append(f"Invalid pitch '{pitch}', using normal")
            result["pitch"] = "normal"
        else:
            result["pitch"] = pitch

    result["pauses"] = raw.get("pauses", True)

    return result, warnings


# ---------------------------------------------------------------------------
# Inline markup parsing
# ---------------------------------------------------------------------------


def find_markup(text: str) -> list[dict[str, Any]]:
    """Find all inline markup in text. Returns list of match dicts."""
    results: list[dict[str, Any]] = []
    seen: set[str] = set()

    def add_match(match: re.Match, mtype: str, value: Any = None) -> None:
        raw = match.group(0)
        key = f"{match.start()}:{raw}"
        if key in seen:
            return
        seen.add(key)
        results.append(
            {
                "raw": raw,
                "type": mtype,
                "value": value,
                "start": match.start(),
                "end": match.end(),
            }
        )

    for m in RE_BOLD.finditer(text):
        add_match(m, "bold")
    for m in RE_ITALIC.finditer(text):
        add_match(m, "italic")
    for m in RE_SPEED.finditer(text):
        val = float(m.group(1))
        add_match(m, "speed", val)
    for m in RE_PITCH.finditer(text):
        add_match(m, "pitch", m.group(1))
    for m in RE_PAUSE.finditer(text):
        val = int(m.group(1))
        add_match(m, "pause", val)

    results.sort(key=lambda x: x["start"])

    return results


def _inner_delim_len(markup: dict[str, Any]) -> int:
    """Return the length of the leading/trailing delimiter for a markup."""
    if markup["type"] == "bold":
        return 2
    if markup["type"] == "italic":
        return 1
    # For bracketed directives, return 0 (no inner content splitting)
    return 0
    return results


# ---------------------------------------------------------------------------
# SSML generation
# ---------------------------------------------------------------------------


def escape_xml(text: str) -> str:
    """Escape XML special characters."""
    return (
        text.replace("&", "&amp;")
        .replace("<", "&lt;")
        .replace(">", "&gt;")
        .replace('"', "&quot;")
        .replace("'", "&apos;")
    )


def generate_ssml(
    segments: list[dict[str, Any]],
    options: dict[str, Any],
) -> str:
    """Generate SSML output from parsed segments."""
    global_attrs: list[str] = []
    if options.get("speed") is not None:
        global_attrs.append(f'rate="{options["speed"]}"')
    if options.get("pitch") is not None:
        global_attrs.append(f'pitch="{options["pitch"]}"')

    parts: list[str] = []
    for seg in segments:
        text = seg.get("text", "")
        directives = seg.get("directives", [])

        if not directives:
            parts.append(escape_xml(text))
            continue

        content = escape_xml(text)
        for d in reversed(directives):
            dtype = d["type"]
            if dtype == "bold":
                content = f'<emphasis level="strong">{content}</emphasis>'
            elif dtype == "italic":
                content = f'<emphasis level="moderate">{content}</emphasis>'
            elif dtype == "speed":
                content = f'<prosody rate="{d["value"]}">{content}</prosody>'
            elif dtype == "pitch":
                content = f'<prosody pitch="{d["value"]}">{content}</prosody>'
            elif dtype == "pause":
                ms = max(50, min(10000, int(d.get("value", 500))))
                content = f'<break time="{ms}ms"/>'

        parts.append(content)

    body = "".join(parts)
    if global_attrs:
        body = f"<prosody {' '.join(global_attrs)}>{body}</prosody>"

    return f"<speak>{body}</speak>"


# ---------------------------------------------------------------------------
# Main pipeline
# ---------------------------------------------------------------------------

def extract_inline_directives(text: str) -> tuple[str, dict[str, Any]]:
    """Extract bracketed inline directives from text as global options.

    Removes [speed:Nx] and [pitch:level] from the text, returning their
    values as extracted options. [pause:Nms] is left in the text for
    segment-level handling.
    """
    clean = text
    extracted: dict[str, Any] = {}

    m = RE_SPEED.search(clean)
    if m:
        extracted["speed"] = float(m.group(1))
        clean = clean.replace(m.group(0), "")

    m = RE_PITCH.search(clean)
    if m:
        extracted["pitch"] = m.group(1)
        clean = clean.replace(m.group(0), "")

    # Collapse multiple spaces left by bracket removal
    clean = re.sub(r"  +", " ", clean).strip()

    return clean, extracted


def process_text(text: str, options: dict[str, Any]) -> dict[str, Any]:
    """Full pipeline: extract → validate → segment → SSML."""
    warnings: list[str] = []

    # Stage 1: Extract bracketed inline directives from text
    # [speed:1.5x] and [pitch:high] are removed from text and merged into options.
    # [pause:500ms] stays — it becomes a segment-level directive.
    clean_text, extracted_opts = extract_inline_directives(text)
    merged_opts = {**extracted_opts, **options}

    applied_opts, opt_warnings = validate_options(merged_opts)
    warnings.extend(opt_warnings)

    markups = find_markup(clean_text)

    if not markups and not any(
        [
            options.get("bold"),
            options.get("italic"),
            options.get("speed") is not None,
            options.get("pitch") is not None,
        ]
    ):
        segments = [{"text": clean_text, "directives": []}]
    else:
        # Build boundaries from markup positions
        # For text-wrapping markups (bold, italic), add inner content boundaries
        boundaries: list[int] = [0]
        for m in markups:
            if m["start"] not in boundaries:
                boundaries.append(m["start"])
            # Inner content boundaries for text-wrapping markups
            if m["type"] in ("bold", "italic"):
                delim = _inner_delim_len(m)
                inner_start = m["start"] + delim
                inner_end = m["end"] - delim
                if inner_start < inner_end:
                    if inner_start not in boundaries:
                        boundaries.append(inner_start)
                    if inner_end not in boundaries:
                        boundaries.append(inner_end)
            if m["end"] not in boundaries:
                boundaries.append(m["end"])

        boundaries = list(dict.fromkeys(boundaries))  # dedup preserving order
        boundaries.sort()
        if boundaries[-1] < len(clean_text):
            boundaries.append(len(clean_text))
        segments = []
        for i in range(len(boundaries) - 1):
            seg_start = boundaries[i]
            seg_end = boundaries[i + 1]
            raw_text = clean_text[seg_start:seg_end]
            # Collect applicable directives (using inner content range)
            directives: list[dict[str, Any]] = []
            for m in markups:
                inner_start_m = m["start"] + _inner_delim_len(m)
                inner_end_m = m["end"] - _inner_delim_len(m)
                if seg_start < inner_end_m and seg_end > inner_start_m:
                    d = {"type": m["type"]}
                    if m.get("value") is not None:
                        d["value"] = m["value"]
                    directives.append(d)

            # Skip delimiter-only segments (e.g. "**" delimiters)
            # BUT only if no directives were collected — pause directives
            # live on delimiter-matching segments and must be preserved.
            if not directives and any(m["raw"] == raw_text for m in markups):
                continue
            # Collect applicable directives (using inner content range)
            directives: list[dict[str, Any]] = []
            for m in markups:
                inner_start_m = m["start"] + _inner_delim_len(m)
                inner_end_m = m["end"] - _inner_delim_len(m)
                if seg_start < inner_end_m and seg_end > inner_start_m:
                    d = {"type": m["type"]}
                    if m.get("value") is not None:
                        d["value"] = m["value"]
                    directives.append(d)

            if raw_text or directives:
                segments.append({"text": raw_text, "directives": directives})

    ssml = generate_ssml(segments, applied_opts)

    return {
        "segments": segments,
        "originalText": text,
        "appliedOptions": applied_opts,
        "warnings": warnings,
        "ssml": ssml,
    }


def main() -> None:
    """Read JSON from stdin, write processed result to stdout."""
    try:
        raw_input = sys.stdin.read()
        if not raw_input.strip():
            print(
                json.dumps(
                    {
                        "error": "No input received",
                        "segments": [],
                        "ssml": "",
                        "originalText": "",
                        "appliedOptions": {},
                        "warnings": [],
                    },
                ),
            )
            return

        data = json.loads(raw_input)
        text = data.get("text", "")
        options = data.get("options", {})

        result = process_text(text, options)
        print(json.dumps(result, indent=2, ensure_ascii=False))

    except json.JSONDecodeError as e:
        print(json.dumps({"error": f"Invalid JSON input: {e}"}))
        sys.exit(1)
    except Exception as e:
        print(json.dumps({"error": f"Processing error: {e}"}))
        sys.exit(1)


if __name__ == "__main__":
    main()