python-services/voice_pipeline.py

# voice_pipeline.py
# Reusable Voice Turn Execution Pipeline with callback hooks and thread-safe cancellation
#
# v4 — Issues 4 + 5 fixed (TTS late-start, foreign-script leakage)
#
#   Issue 4 (text appears instantly, TTS starts speaking late):
#     on_token() streams every token to the frontend the instant it
#     arrives (this is why text looks instant). But TTS only fired once
#     a FULL sentence boundary (۔ ، ! ?) was detected via SENTENCE_END_RE.
#     If the LLM's first sentence was long (15-30 words), TTS audio
#     couldn't start until the ENTIRE sentence had streamed in —
#     1-2 seconds of "text is there, voice isn't".
#
#     Fix: for the FIRST chunk of a turn only, _extract_first_clause()
#     also accepts a comma/clause boundary (، or ,) as a valid TTS
#     dispatch point, as long as the accumulated clause is at least
#     MIN_FIRST_CLAUSE_CHARS (15) long (so we don't synthesise tiny
#     2-word fragments with bad prosody). Every subsequent chunk in the
#     same turn uses the original full-sentence-only boundary
#     (SENTENCE_END_RE / _extract_sentences) for natural-sounding
#     speech. This cuts time-to-first-audio roughly in half for
#     long opening sentences while leaving the rest of the response's
#     prosody unaffected.
#
#   Issue 5 (Chinese/Hindi/other foreign-script words leaking into
#   Urdu speech):
#     The system prompt forbade Latin script but never explicitly
#     forbade OTHER non-Arabic scripts (Chinese/CJK, Devanagari/Hindi,
#     Cyrillic, Hangul, Thai, etc). There was also zero programmatic
#     safety net — whatever the LLM produced went straight to TTS,
#     which mispronounces or chokes on foreign Unicode mid-call.
#
#     Fix (two layers, this file is layer 2):
#       Layer 1 (prompt_builder.py): explicit multi-script prohibition
#       added to the system prompt with named scripts.
#       Layer 2 (here): strip_foreign_scripts() runs on EVERY sentence
#       immediately before _fire_tts() is called — a hard safety net
#       that removes CJK/Devanagari/Cyrillic/Hangul/Thai characters
#       regardless of what the LLM actually output. Latin characters
#       and digits are deliberately NOT stripped here, since they're
#       legitimate in Urdu call-center speech (times, prices, phone
#       numbers, acronyms) — those are handled by the existing
#       TTS-side transliterator. Only genuinely foreign SCRIPTS are
#       removed, never re-interpreted or transliterated (that's a TTS-
#       side concern); this is purely a "never let TTS choke on
#       garbage" guard.
#
# (v3 note retained for reference)
# v3 — arabic_clarify removed
#   _build_prompt() no longer passes arabic_clarification to build_system_prompt().
#   prompt_builder.build_system_prompt() no longer accepts that parameter.
#
# v2 — URDU-ONLY MODE
#   _fire_tts() uses session's locked language (always "ur") instead of
#   per-sentence content detection to pick TTS voice.

import os
import time
import logging
import asyncio
import re
import threading
import httpx
from typing import List, Optional, Dict, Any

from client import groq, fireworks
from utils import detect_language_from_content, SUPPORTED_LANGUAGES
from language_config import normalize_language
from prompt_builder import build_system_prompt
from summary import generate_summary
from greeting_handler import (
    classify_opening_message,
    get_short_reply,
    should_short_circuit,
)
from utils import (
    estimate_payload_tokens,
    strip_formatting,
)

logger = logging.getLogger("voice_pipeline")

TOKEN_THRESHOLD = 4000
DURATION_THRESHOLD_MINUTES = 30
MODEL_FAST = "llama-3.1-8b-instant"
MODEL_HEAVY = "llama-3.3-70b-versatile"
FALLBACK_MODEL_FIREWORKS = "accounts/fireworks/models/llama-v3p1-8b-instruct"

TTS_SPACE_URL = os.getenv("TTS_SPACE_URL", "https://voice-tts-tts.hf.space")
TTS_SYNTHESISE_TIMEOUT = 15.0

# Full-sentence boundary — used for every chunk AFTER the first in a turn.
SENTENCE_END_RE = re.compile(
    r'[^.!?\u06d4\u061f\u0964\n]+[.!?\u06d4\u061f\u0964]+(?:\s|$)'
    r'|[^.!?\u06d4\u061f\u0964\n]*\n',
    re.UNICODE,
)

# Issue 4: clause boundary (comma OR sentence end) — used ONLY for the
# first chunk of a turn, to reduce time-to-first-audio on long opening
# sentences. MIN_FIRST_CLAUSE_CHARS guards against dispatching tiny
# fragments with poor TTS prosody.
CLAUSE_END_RE = re.compile(
    r'[^.!?,\u060c\u06d4\u061f\u0964\n]+[.!?,\u060c\u06d4\u061f\u0964]+(?:\s|$)'
    r'|[^.!?,\u060c\u06d4\u061f\u0964\n]*\n',
    re.UNICODE,
)
MIN_FIRST_CLAUSE_CHARS = 15

# Issue 5: known foreign (non-Urdu, non-Latin) script ranges. Latin and
# digits are intentionally excluded — they're legitimate in Urdu call
# speech (times, prices, acronyms) and handled by the TTS-side
# transliterator, not stripped here.
FOREIGN_SCRIPT_RE = re.compile(
    r'[\u4E00-\u9FFF\u3400-\u4DBF'   # CJK Unified Ideographs + Ext A (Chinese)
    r'\u3040-\u309F\u30A0-\u30FF'     # Hiragana, Katakana (Japanese)
    r'\uAC00-\uD7A3'                  # Hangul syllables (Korean)
    r'\u0900-\u097F'                  # Devanagari (Hindi)
    r'\u0400-\u04FF'                  # Cyrillic (Russian etc.)
    r'\u0E00-\u0E7F'                  # Thai
    r']+'
)


def strip_foreign_scripts(text: str) -> str:
    """
    Issue 5: hard safety-net filter. Removes characters from scripts that
    have no place in Urdu-only TTS output (Chinese, Japanese, Korean,
    Hindi/Devanagari, Cyrillic, Thai) regardless of what the LLM produced.
    Latin letters and digits are NOT touched — legitimate in Urdu call
    speech (times, prices, acronyms) and handled by the TTS transliterator.
    Logs a warning whenever it actually strips something, for visibility.
    """
    if not text:
        return text
    cleaned = FOREIGN_SCRIPT_RE.sub('', text)
    cleaned = re.sub(r'\s{2,}', ' ', cleaned).strip()
    if cleaned != text.strip():
        logger.warning(
            f"[Issue 5] Stripped foreign-script characters from LLM output: "
            f"original={text[:80]!r} cleaned={cleaned[:80]!r}"
        )
    return cleaned


def _extract_sentences(buffer: str) -> tuple[list[str], str]:
    sentences: list[str] = []
    last_end = 0
    for m in SENTENCE_END_RE.finditer(buffer):
        sentence = buffer[m.start():m.end()].strip()
        if sentence:
            sentences.append(sentence)
        last_end = m.end()
    return sentences, buffer[last_end:]


def _extract_first_clause(buffer: str) -> tuple[Optional[str], str]:
    """
    Issue 4: for the FIRST chunk of a turn only. Finds the earliest
    point where accumulated text reaches a clause boundary (comma or
    sentence end) AND is at least MIN_FIRST_CLAUSE_CHARS long. Returns
    (clause_text, remainder) or (None, buffer) if no qualifying
    boundary has been reached yet.
    """
    matches = list(CLAUSE_END_RE.finditer(buffer))
    if not matches:
        return None, buffer

    for m in matches:
        candidate = buffer[0:m.end()].strip()
        if len(candidate) >= MIN_FIRST_CLAUSE_CHARS:
            return candidate, buffer[m.end():]

    return None, buffer


async def fire_tts(sentence: str, language: str, gender: str) -> tuple[Optional[bytes], str]:
    """Public alias — used by voice_orchestrator for bridge phrase synthesis."""
    return await _fire_tts(sentence, language, gender)


async def _fire_tts(sentence: str, language: str, gender: str) -> tuple[Optional[bytes], str]:
    """
    v2 (Urdu-only mode): tts_lang is always the session's language —
    no per-sentence content-based re-detection.

    v4 (Issue 5): sentence is passed through strip_foreign_scripts()
    before being sent to TTS — a hard safety net regardless of what
    upstream prompt/model behaviour produced.
    """
    sentence = sentence.strip()
    if not sentence:
        return None, normalize_language(language)

    # Issue 5: strip any foreign-script characters before synthesis.
    sentence = strip_foreign_scripts(sentence)
    if not sentence:
        # Entire "sentence" was foreign-script garbage with nothing left —
        # nothing meaningful to synthesise.
        return None, normalize_language(language)

    tts_lang = normalize_language(language)

    # DISABLED (Urdu-only mode) — per-sentence content-based language detection:
    # tts_lang = normalize_language(detect_language_from_content(sentence))

    try:
        async with httpx.AsyncClient(timeout=TTS_SYNTHESISE_TIMEOUT) as client:
            resp = await client.post(
                f"{TTS_SPACE_URL}/synthesise",
                json={"text": sentence, "language": tts_lang, "gender": gender},
            )
        if resp.status_code == 200:
            logger.info(f"TTS ✓  lang={tts_lang}  chars={len(sentence)}")
            return resp.content, tts_lang
        logger.warning(f"TTS Space returned HTTP {resp.status_code}: {resp.text[:120]}")
        return None, tts_lang
    except httpx.TimeoutException:
        logger.error(f"TTS request timed out ({TTS_SYNTHESISE_TIMEOUT}s): {sentence[:60]!r}")
        return None, tts_lang
    except Exception as exc:
        logger.error(f"TTS request error: {exc}")
        return None, tts_lang


def _start_groq_stream_thread(
    params: dict,
    q: asyncio.Queue,
    loop: asyncio.AbstractEventLoop,
    cancel_event: threading.Event,
) -> None:
    def _run() -> None:
        try:
            stream = groq.chat.completions.create(**{**params, "stream": True})
            for chunk in stream:
                if cancel_event.is_set():
                    return
                delta = chunk.choices[0].delta
                if delta and delta.content:
                    loop.call_soon_threadsafe(q.put_nowait, ("token", delta.content))
            loop.call_soon_threadsafe(q.put_nowait, ("done", None))
        except Exception as exc:
            if not cancel_event.is_set():
                loop.call_soon_threadsafe(q.put_nowait, ("error", str(exc)))

    threading.Thread(target=_run, daemon=True, name="groq-stream-worker").start()


def _select_primary_model(session, detected_language: str) -> str:
    ld = session.language_detector
    if detected_language == "ur" and (ld.is_locked() or ld.confidence >= 0.70):
        return MODEL_HEAVY
    return MODEL_FAST


def _fields_spec_dicts(session) -> List[Dict[str, Any]]:
    return list(session.caller_collection_fields or [])


async def _roll_summary(session, current_summary: str) -> str:
    if len(session.messages) <= 6:
        return current_summary or ""
    messages_to_summarize = session.messages[:-6]
    session.messages = session.messages[-6:]
    return await generate_summary(
        current_summary or "",
        messages_to_summarize,
        caller_info=session.caller_info.to_dict(),
        caller_fields_spec=_fields_spec_dicts(session),
    )


async def _query_rag(session, message: str) -> str:
    if not getattr(session, "rag_enabled", False):
        return ""
    company_id = getattr(session, "company_id", None)
    if not company_id:
        return ""

    agent_config_id = getattr(session, "agent_config_id", None)
    rag_tenant_id = f"{company_id}_a{agent_config_id}" if agent_config_id else company_id

    RAG_SPACE_URL = os.getenv("RAG_SPACE_URL", "")
    RAG_API_KEY   = os.getenv("RAG_API_KEY", "")
    if not RAG_SPACE_URL:
        logger.warning("[RAG] RAG_SPACE_URL not set — skipping live RAG query.")
        return ""

    try:
        async with httpx.AsyncClient(timeout=3.0) as client:
            tenant_ids = [rag_tenant_id]
            if agent_config_id and rag_tenant_id != company_id:
                tenant_ids.append(company_id)  # legacy fallback for pre-multi-agent KB

            for tenant in tenant_ids:
                resp = await client.post(
                    f"{RAG_SPACE_URL}/api/v1/query",
                    json={"company_id": tenant, "query": message, "top_k": 5},
                    headers={"X-RAG-API-Key": RAG_API_KEY},
                )

                if resp.status_code != 200:
                    logger.warning(f"[RAG] Query returned HTTP {resp.status_code} for tenant={tenant}")
                    continue

                rag_data = resp.json()
                if rag_data.get("needs_rag") and rag_data.get("context"):
                    logger.info(
                        f"[RAG] ✅ Context retrieved for tenant={tenant} "
                        f"mode={rag_data.get('search_mode')} "
                        f"chunks={len(rag_data.get('chunks', []))}"
                    )
                    return rag_data["context"]

                logger.info(f"[RAG] Skipped for tenant={tenant}: needs_rag={rag_data.get('needs_rag')}")

        return ""

    except httpx.TimeoutException:
        logger.warning(f"[RAG] Query timed out (3s) for tenant={rag_tenant_id} — continuing without context")
        return ""
    except Exception as rag_err:
        logger.warning(f"[RAG] Query failed (non-fatal) for tenant={rag_tenant_id}: {rag_err}")
        return ""


def _build_prompt(session, detected_language: str, input_text: str) -> str:
    greeting_phase = (
        session.turn_count <= 3
        and not session.language_detector.is_locked()
        and session.greeting_streak > 0
    )
    return build_system_prompt(
        company_name=session.company_name,
        agent_name=session.agent_name,
        agent_gender=session.agent_gender,
        language=detected_language,
        document_context=session.document_context,
        custom_rules=session.custom_rules,
        collected_caller_info=session.caller_info.to_dict(),
        input_text=input_text,
        # REMOVED: arabic_clarification (Urdu-only mode)
        caller_collection_fields=_fields_spec_dicts(session),
        turn_count=session.turn_count,
        greeting_phase=greeting_phase,
        company_name_ur=getattr(session, "company_name_ur", None),
        agent_name_ur=getattr(session, "agent_name_ur", None),
        agent_role=getattr(session, "agent_role", None),
        agent_role_ur=getattr(session, "agent_role_ur", None),
        overview_ur=getattr(session, "overview_ur", None),
    )


async def run_turn(
    session,
    message: str,
    *,
    cancel_event: threading.Event,
    on_token,
    on_audio,
    on_done,
    agent_gender: str = "male",
    current_summary: str = "",
    session_start: Optional[int] = None
) -> None:
    start_timestamp = time.time()

    detected_language = normalize_language(session.process_user_message(message))
    session.add_message("user", message)

    opening_tier = classify_opening_message(
        message,
        session.turn_count,
        session.language_detector.is_locked(),
    )
    short_circuit = should_short_circuit(opening_tier)
    if short_circuit:
        session.greeting_streak += 1
    else:
        session.greeting_streak = 0

    new_summary, rag_context = await asyncio.gather(
        _roll_summary(session, current_summary or ""),
        _query_rag(session, message),
    )

    if rag_context:
        session.document_context = rag_context

    system_prompt = _build_prompt(session, detected_language, message)

    optimized_payload = [{"role": "system", "content": system_prompt}]
    if new_summary:
        optimized_payload.append({
            "role": "system",
            "content": f"Summary of earlier conversation context for your memory bank: {new_summary}",
        })
    optimized_payload.extend(session.get_history_for_api())

    raw_session_start = session_start or 0
    if 0 < raw_session_start < 1e10:
        raw_session_start *= 1000
    duration_minutes = max(0.0, (time.time() * 1000 - raw_session_start) / 60000) if raw_session_start else 0
    estimated_tokens = estimate_payload_tokens(optimized_payload)
    max_response_tokens: Optional[int] = None

    if duration_minutes >= DURATION_THRESHOLD_MINUTES or estimated_tokens > TOKEN_THRESHOLD:
        logger.warning(
            f"[pipeline] Token Governor active. "
            f"Tokens={estimated_tokens}, Duration={duration_minutes:.1f}m"
        )
        cleaned_history = [
            {"role": m["role"], "content": strip_formatting(m["content"])}
            for m in session.get_history_for_api()
        ]
        concise_control = {
            "role": "system",
            "content": "Token Governor Active: Reply ultra-concisely. No markdown, no filler. 1-2 short sentences max.",
        }
        short_sys = [
            m for m in optimized_payload
            if m.get("role") == "system" and len(m.get("content", "")) < 400
        ]
        optimized_payload = [concise_control] + short_sys
        if new_summary:
            optimized_payload.append({
                "role": "system",
                "content": f"Summary of earlier conversation context for your memory bank: {new_summary}",
            })
        optimized_payload.extend(cleaned_history)
        max_response_tokens = 120

    primary_model = _select_primary_model(session, detected_language)
    inference_params: Dict[str, Any] = {
        "model": primary_model,
        "messages": optimized_payload,
        "temperature": 0.4,
        "top_p": 0.95,
    }
    if max_response_tokens:
        inference_params["max_tokens"] = max_response_tokens

    tts_gender = agent_gender if agent_gender in ("female", "male") else "female"

    loop = asyncio.get_running_loop()
    q: asyncio.Queue = asyncio.Queue()
    full_text_parts: list[str] = []
    buffer = ""
    tts_tasks: List[asyncio.Task] = []
    # Issue 4: track whether we've dispatched the first TTS chunk yet —
    # only the FIRST chunk of a turn uses clause-level (comma) splitting.
    first_chunk_dispatched = False

    async def _synthesise_and_emit(sentence: str):
        if cancel_event.is_set():
            return
        audio_bytes, tts_lang = await _fire_tts(sentence, detected_language, tts_gender)
        if audio_bytes and not cancel_event.is_set():
            if asyncio.iscoroutinefunction(on_audio):
                await on_audio(sentence, audio_bytes, tts_lang)
            else:
                on_audio(sentence, audio_bytes, tts_lang)

    try:
        if short_circuit:
            short_reply = get_short_reply(opening_tier, message, session.greeting_streak)
            full_text_parts.append(short_reply)

            if cancel_event.is_set():
                return

            if asyncio.iscoroutinefunction(on_token):
                await on_token(short_reply)
            else:
                on_token(short_reply)

            task = asyncio.create_task(_synthesise_and_emit(short_reply))
            tts_tasks.append(task)
            await task
        else:
            _start_groq_stream_thread(inference_params, q, loop, cancel_event)

            while not cancel_event.is_set():
                event_type, data = await q.get()

                if cancel_event.is_set():
                    break

                if event_type == "error":
                    if fireworks is not None:
                        logger.warning("[pipeline] Groq error -> Fireworks fallback")
                        try:
                            fb_params = {**inference_params, "model": FALLBACK_MODEL_FIREWORKS}
                            fb_resp = await loop.run_in_executor(
                                None,
                                lambda: fireworks.chat.completions.create(**fb_params),
                            )
                            if cancel_event.is_set():
                                return
                            fallback_text = fb_resp.choices[0].message.content.strip()
                            full_text_parts.append(fallback_text)

                            if asyncio.iscoroutinefunction(on_token):
                                await on_token(fallback_text)
                            else:
                                on_token(fallback_text)

                            fb_sentences, _ = _extract_sentences(fallback_text + " ")
                            for sent in fb_sentences:
                                if cancel_event.is_set():
                                    break
                                task = asyncio.create_task(_synthesise_and_emit(sent))
                                tts_tasks.append(task)

                            if tts_tasks:
                                await asyncio.gather(*tts_tasks, return_exceptions=True)
                        except Exception as fb_exc:
                            logger.error(f"[pipeline] Fallback failed: {fb_exc}")
                            raise fb_exc
                    else:
                        raise Exception(data)
                    break

                elif event_type == "done":
                    remainder = buffer.strip()
                    if remainder and not cancel_event.is_set():
                        full_text_parts.append(remainder)
                        if asyncio.iscoroutinefunction(on_token):
                            await on_token(remainder)
                        else:
                            on_token(remainder)
                        task = asyncio.create_task(_synthesise_and_emit(remainder))
                        tts_tasks.append(task)
                        first_chunk_dispatched = True

                    if tts_tasks:
                        await asyncio.gather(*tts_tasks, return_exceptions=True)
                    break

                else:
                    token: str = data
                    full_text_parts.append(token)
                    if asyncio.iscoroutinefunction(on_token):
                        await on_token(token)
                    else:
                        on_token(token)
                    buffer += token

                    # Issue 4: first chunk of the turn uses clause-level
                    # (comma) splitting to reduce time-to-first-audio on
                    # long opening sentences. Every chunk after that uses
                    # full-sentence-only splitting for natural prosody.
                    if not first_chunk_dispatched:
                        clause, buffer = _extract_first_clause(buffer)
                        if clause:
                            task = asyncio.create_task(_synthesise_and_emit(clause))
                            tts_tasks.append(task)
                            first_chunk_dispatched = True
                    else:
                        completed_sentences, buffer = _extract_sentences(buffer)
                        for sentence in completed_sentences:
                            if cancel_event.is_set():
                                break
                            task = asyncio.create_task(_synthesise_and_emit(sentence))
                            tts_tasks.append(task)

    except asyncio.CancelledError:
        logger.warning(f"[pipeline] asyncio.CancelledError caught. Setting cancel flag.")
        cancel_event.set()
    finally:
        for task in tts_tasks:
            if not task.done():
                task.cancel()

        if not cancel_event.is_set() or len(full_text_parts) > 0:
            full_text = "".join(full_text_parts).strip()
            session.add_message("assistant", full_text)

            elapsed_ms = int((time.time() - start_timestamp) * 1000)
            if not cancel_event.is_set():
                if asyncio.iscoroutinefunction(on_done):
                    await on_done(full_text, detected_language, new_summary, elapsed_ms)
                else:
                    on_done(full_text, detected_language, new_summary, elapsed_ms)