Spaces:

010Abid
/

the-brain

Running

the-brain / python-services /llm_server.py
1 commit ahead with correct commit i dont about this one
e9b447a 4 months ago
42.7 kB
# llm_server.py
# Main FastAPI server — Multilingual Groq LLM Call Center Voice Orchestrator
# Exposed on port 7860 on HF Spaces
#
# v3 — Issues 4 + 5 fixed (mirrors voice_pipeline.py v4)
#   Applied to the /chat/stream SSE path for consistency with the
#   primary voice path (the-brain's WS voice flow uses voice_pipeline.py,
#   not this file — but /chat/stream is a secondary entry point that
#   should behave identically):
#     Issue 4: first chunk of a turn uses clause-level (comma) TTS
#       splitting instead of waiting for a full sentence, reducing
#       time-to-first-audio on long opening sentences.
#     Issue 5: _strip_foreign_scripts() runs on every sentence before
#       TTS — hard safety net removing CJK/Devanagari/Cyrillic/Hangul/
#       Thai characters regardless of what the LLM produced.
#
# v2 — URDU-ONLY MODE (matches voice_pipeline.py)
#   _fire_tts() previously re-detected language PER SENTENCE via
#   detect_language_from_content(), and used that (not the session's
#   locked language) to pick the TTS voice. Since this product is
#   Urdu-only for now, tts_lang is now ALWAYS normalize_language(language)
#   — the session's locked language (always "ur"). The per-sentence
#   content-detection override is preserved as a comment for when
#   per-agent language selection is reintroduced. This only affects the
#   /chat/stream SSE path; the primary voice path is voice_pipeline.py.

import os
import time
import logging
import asyncio
import re
import base64
import json
import threading
import httpx
from typing import List, Optional, Dict, Any
from fastapi import FastAPI, HTTPException, WebSocket, WebSocketDisconnect
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse, StreamingResponse
from pydantic import BaseModel

# Internal Domain Layer Imports
from client import groq, fireworks
from conversation_manager import ConversationManager, ConversationSession
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,
)
from stt_client import SttClient
from voice_orchestrator import VoiceOrchestratorSession
import neon_client

# ── Logging System Initialization ───────────────────────────────────────────
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(name)s] %(levelname)s: %(message)s",
    datefmt="%Y-%m-%d %H:%M:%S",
)
logger = logging.getLogger("llm_server")

# Verify core execution environment states
if not os.getenv("GROQ_API_KEY"):
    logger.critical("CRITICAL: GROQ_API_KEY environment configuration variable is missing.")

# ── Architectural Threshold Variables ───────────────────────────────────────
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 Bridge Configuration ──────────────────────────────────────────
# URL is read from env so it can be overridden without a code deploy.
TTS_SPACE_URL = os.getenv("TTS_SPACE_URL", "https://voice-tts-tts.hf.space")
TTS_SYNTHESISE_TIMEOUT = 15.0  # seconds per sentence request

# Multi-script sentence boundary detector.
# Captures a run of non-terminator characters followed by one or more
# sentence-ending punctuation marks, then optional whitespace / end-of-string.
# Terminators covered:
#   .  !  ?          — English
#   \u06d4  \u061f   — Urdu (۔) / Arabic-Urdu (؟)
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 /chat/stream turn, to reduce time-to-first-audio on
# long opening sentences. Mirrors voice_pipeline.py's fix for the
# primary voice path.
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 — legitimate in Urdu call speech
# (times, prices, acronyms) and handled by the TTS-side transliterator.
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']+'
)

# ── FastAPI App Configuration ────────────────────────────────────────────────
app = FastAPI(
    title="The Brain — Voice LLM Call Center Agent",
    description="English/Urdu bilingual voice pipeline orchestration layer powered by Groq and Fireworks AI.",
    version="1.0.0",
)

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_methods=["GET", "POST", "OPTIONS"],
    allow_headers=["*"],
    expose_headers=["X-Session-Id", "X-Language", "X-Turn-Count"],
)

# Global memory state allocations
manager = ConversationManager()


class CallerFieldSpec(BaseModel):
    key: str
    label: str
    label_ur: str = ""
    type: str = "text"


# ── Unified Pydantic Request Models ─────────────────────────────────────────
class StartSessionRequest(BaseModel):
    company_name: str = "Our Company"
    agent_name: str = "Sara"
    agent_gender: str = "male"           # Default changed to male
    custom_rules: List[str] = []
    document_context: str = ""          # RAG text block container
    caller_collection_fields: List[CallerFieldSpec] = []
    session_id: Optional[str] = None    # Generated if missing

class ChatRequest(BaseModel):
    session_id: str
    message: str
    company_name: Optional[str] = None
    agent_name: Optional[str] = None
    agent_gender: Optional[str] = "male"  # Re-supplied for emergency session fallback
    currentSummary: Optional[str] = ""
    sessionStart: Optional[int] = None   # Epoch milliseconds payload timestamp

class EndSessionRequest(BaseModel):
    session_id: str


# ── TTS Bridge Helpers ──────────────────────────────────────────────────────

def _extract_sentences(buffer: str) -> tuple[list[str], str]:
    """
    Scan the rolling token buffer for complete sentences using SENTENCE_END_RE.
    Returns (completed_sentences, leftover_text) so the caller can continue
    accumulating tokens into the leftover portion.
    """
    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. Mirrors voice_pipeline.py.
    """
    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


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. Mirrors voice_pipeline.py.
    """
    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


async def _fire_tts(sentence: str, language: str, gender: str) -> tuple[Optional[bytes], str]:
    """
    Asynchronously POST a single sentence to the deployed TTS Space.
    Returns (audio_bytes, tts_language) on success, (None, tts_language) on failure.
    v2 (Urdu-only mode): tts_lang is always the session's language
    (normalize_language(language)) — no per-sentence content-based
    re-detection. See module docstring.
    """
    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:
        return None, normalize_language(language)

    tts_lang = normalize_language(language)

    # DISABLED (Urdu-only mode) — per-sentence content-based language
    # detection. Re-enable for multi-language support:
    #
    # tts_lang = normalize_language(detect_language_from_content(sentence))
    # session_lang = normalize_language(language)
    # if tts_lang != session_lang:
    #     logger.info(
    #         "TTS lang from content: session=%s → sentence=%s preview=%r",
    #         session_lang,
    #         tts_lang,
    #         sentence[:60],
    #     )

    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,
) -> None:
    """
    Spawns a daemon thread that calls Groq's synchronous streaming API
    and forwards each content chunk onto the asyncio Queue via
    loop.call_soon_threadsafe so the async event generator can consume it.
    Events pushed onto the queue:
      ("token", str)   — a streamed content chunk from the LLM
      ("done",  None)  — stream finished cleanly
      ("error", str)   — exception message from the stream
    """
    def _run() -> None:
        try:
            stream = groq.chat.completions.create(**{**params, "stream": True})
            for chunk in stream:
                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:
            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: ConversationSession, 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: ConversationSession) -> List[Dict[str, Any]]:
    return list(session.caller_collection_fields or [])


async def _roll_summary(session: ConversationSession, 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),
    )


def _build_prompt(session: ConversationSession, 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),
    )


# ── REST API Router Endpoints ────────────────────────────────────────────────

@app.post("/session/start")
async def start_session(req: StartSessionRequest):
    """
    Allocates and registers a stateful conversation block for an incoming call.
    Returns tracking identifiers required for subsequent transactional steps.
    """
    fields = [f.model_dump() for f in req.caller_collection_fields]
    session = manager.create_session(
        company_name=req.company_name,
        agent_name=req.agent_name,
        agent_gender=req.agent_gender,   # Pass parameter
        custom_rules=req.custom_rules,
        document_context=req.document_context,
        caller_collection_fields=fields,
        session_id=req.session_id,
    )

    logger.info(f"Communication sequence initialized: {session.session_id}")
    return {
        "session_id": session.session_id,
        "company_name": req.company_name,
        "agent_name": req.agent_name,
        "status": "ready",
        "message": f"Agent {req.agent_name} online and monitoring streaming data frames."
    }


@app.post("/chat")
async def chat(req: ChatRequest):
    """
    Processes voice transcript inputs, executes dynamic RAG context assembly,
    evaluates historical payload token depth, and executes low-latency model inference.
    """
    if not req.message or not req.message.strip():
        raise HTTPException(status_code=400, detail="Transaction rejected: Content payload string cannot be empty.")
    if len(req.message) > 2000:
        raise HTTPException(status_code=400, detail="Transaction rejected: Maximum allowed buffer exceeds 2000 characters.")

    # 1. Fetch target state context, fallback to automatic creation if missing
    session = manager.get_session(req.session_id)
    if not session:
        session = manager.create_session(
            company_name=req.company_name or "Our Company",
            agent_name=req.agent_name or "Sara",
            agent_gender=req.agent_gender or "male",  # Use request gender, not hardcoded female
            session_id=req.session_id
        )
        logger.warning(f"Target memory space missing. Initialized emergency backup stack instance: {session.session_id}")

    start_timestamp = time.time()

    # 2. Update metadata tracking matrices
    detected_language = normalize_language(session.process_user_message(req.message))
    session.add_message("user", req.message)

    opening_tier = classify_opening_message(
        req.message,
        session.turn_count,
        session.language_detector.is_locked(),
    )

    # 3. Greeting short-circuit — no full LLM intro on bare greetings
    if should_short_circuit(opening_tier):
        session.greeting_streak += 1
        short_reply = get_short_reply(opening_tier, req.message, session.greeting_streak)
        session.add_message("assistant", short_reply)
        elapsed_ms = int((time.time() - start_timestamp) * 1000)
        logger.info(
            f"Greeting short-circuit. ID={session.session_id} | tier={opening_tier} | "
            f"Latency={elapsed_ms}ms"
        )
        return JSONResponse({
            "session_id": session.session_id,
            "provider": "template-greeting",
            "response": short_reply,
            "language": detected_language,
            "language_locked": session.language_detector.is_locked(),
            "turn_count": session.turn_count,
            "caller_info": session.caller_info.to_dict(),
            "activeSummary": req.currentSummary or "",
            "elapsed_ms": elapsed_ms,
            "tokens_used": 0,
        })

    session.greeting_streak = 0

    # 4. Rolling window summary
    new_summary = await _roll_summary(session, req.currentSummary or "")

    # 5. Build system prompt
    system_prompt = _build_prompt(session, detected_language, req.message)

    # 5. Build standard layout configurations
    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())

    # 6. Evaluate Token Governor Constraints
    raw_session_start = req.sessionStart or 0
    # Guard: clients should send epoch milliseconds (>1e12). If seconds were sent
    # instead (value < 1e10), convert to avoid a 1000x error that would immediately
    # trigger the Token Governor on every single request.
    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"Token Governor triggered. Window size: {estimated_tokens} tokens. Call duration: {duration_minutes:.2f} mins.")
        
        # Purge markdown/formatting configurations to preserve strict token structures
        cleaned_history = [
            {"role": m["role"], "content": strip_formatting(m["content"])} 
            for m in session.get_history_for_api()
        ]
        
        concise_controls = {
            "role": "system", 
            "content": "Token Governor Active: Reply ultra-concisely. No markdown, no filler words. Limit output strictly to 1-2 short conversational sentences."
        }
        
        # Isolate baseline identity parameters to decrease structural overhead parameters
        short_system_messages = [
            m for m in optimized_payload 
            if m.get("role") == "system" and len(m.get("content", "")) < 400
        ]
        
        optimized_payload = [concise_controls] + short_system_messages
        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

    # 7. Model Router (Urdu 70B when locked or confident)
    primary_model = _select_primary_model(session, detected_language)

    inference_parameters: Dict[str, Any] = {
        "model": primary_model,
        "messages": optimized_payload,
        "temperature": 0.4,  # Marginally cooled to further guarantee script structure adherence
        "top_p": 0.95
    }
    if max_response_tokens:
        inference_parameters["max_tokens"] = max_response_tokens

    # 8. Main Inference Execution Loop with Resilient High-Availability Fallback
    provider_tag = f"Groq ({primary_model})"
    loop = asyncio.get_running_loop()
    
    try:
        response = await loop.run_in_executor(
            None,
            lambda: groq.chat.completions.create(**inference_parameters)
        )
        assistant_message = response.choices[0].message.content.strip()
        tokens_used = response.usage.total_tokens

    except Exception as groq_error:
        # Check for HTTP/API Status Code 429 Rate Limit Errors
        status_code = getattr(groq_error, "status_code", None) or getattr(groq_error, "status", None)
        if status_code == 429 and fireworks is not None:
            logger.warning("Groq rate limited — falling back to Fireworks")
            
            inference_parameters["model"] = FALLBACK_MODEL_FIREWORKS
            provider_tag = "Fireworks AI (Fallback)"
            
            fallback_response = await loop.run_in_executor(
                None,
                lambda: fireworks.chat.completions.create(**inference_parameters)
            )
            assistant_message = fallback_response.choices[0].message.content.strip()
            tokens_used = fallback_response.usage.total_tokens if (hasattr(fallback_response, "usage") and fallback_response.usage) else 0
        else:
            logger.error(f"Unmanaged runtime error during network execution loop: {str(groq_error)}")
            raise HTTPException(status_code=502, detail=f"Upstream provider failure: {str(groq_error)}")

    # 9. Register output sequence history
    session.add_message("assistant", assistant_message)

    elapsed_ms = int((time.time() - start_timestamp) * 1000)
    logger.info(f"Execution wrapped. Metrics: ID={session.session_id} | Provider={provider_tag} | Latency={elapsed_ms}ms | Language={detected_language}")

    return JSONResponse({
        "session_id": session.session_id,
        "provider": provider_tag,
        "response": assistant_message,
        "language": detected_language,
        "language_locked": session.language_detector.is_locked(),
        "turn_count": session.turn_count,
        "caller_info": session.caller_info.to_dict(),
        "activeSummary": new_summary,
        "elapsed_ms": elapsed_ms,
        "tokens_used": tokens_used
    })


@app.post("/chat/stream")
async def chat_stream(req: ChatRequest):
    """
    Streaming LLM → TTS bridge endpoint using Server-Sent Events (SSE).
    Architecture
    ------------
    1. Groq streaming runs in a daemon thread (_start_groq_stream_thread).
    2. Tokens are forwarded onto an asyncio.Queue so the async generator
       can consume them without blocking the event loop.
    3. Each token is appended to a rolling buffer. When SENTENCE_END_RE
       finds a complete sentence, _fire_tts() immediately fires an async
       httpx POST to the TTS Space — while the LLM is still generating
       the rest of the response (the "live speaking" effect).
    4. On Groq failure the endpoint transparently falls back to Fireworks
       (non-streaming) before surfacing an error to the client.
    SSE event schema (each event is a JSON object on a `data:` line):
      {"type": "token",  "text": "..."}                         — raw LLM token
      {"type": "audio",  "sentence": "...", "audio_b64": "...", "lang": "..."}
                                                                — base64 audio/mpeg
      {"type": "done",   "full_text": "...", "language": "...", ...}
                                                                — final metadata
      {"type": "error",  "detail": "..."}                        — fatal error
    """
    if not req.message or not req.message.strip():
        raise HTTPException(status_code=400, detail="Content payload cannot be empty.")
    if len(req.message) > 2000:
        raise HTTPException(status_code=400, detail="Maximum payload exceeds 2000 characters.")

    # ── 1. Session resolution (mirrors /chat) ─────────────────────────────────
    session = manager.get_session(req.session_id)
    if not session:
        session = manager.create_session(
            company_name=req.company_name or "Our Company",
            agent_name=req.agent_name or "Sara",
            agent_gender=req.agent_gender or "male",  # Use request gender, not hardcoded female
            session_id=req.session_id,
        )
        logger.warning(f"[stream] Emergency session created: {session.session_id}")

    # ── 2. Language detection + greeting gate ───────────────────────────────
    detected_language = normalize_language(session.process_user_message(req.message))
    session.add_message("user", req.message)

    opening_tier = classify_opening_message(
        req.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 = await _roll_summary(session, req.currentSummary or "")
    system_prompt = _build_prompt(session, detected_language, req.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())

    # ── 3. Token Governor (mirrors /chat) ─────────────────────────────────────
    raw_session_start = req.sessionStart 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"[stream] Token Governor triggered. "
            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

    # ── 4. Model selection ────────────────────────────────────────────────────
    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

    # Resolve TTS gender from session (TTS Space accepts "female" | "male")
    tts_gender = session.agent_gender if session.agent_gender in ("female", "male") else "female"
    start_timestamp = time.time()

    # ── 5. SSE event generator ────────────────────────────────────────────────
    async def event_generator():
        loop = asyncio.get_running_loop()
        q: asyncio.Queue = asyncio.Queue()
        full_text_parts: list[str] = []
        buffer = ""
        # Issue 4: first chunk of the turn uses clause-level splitting
        # to reduce time-to-first-audio; subsequent chunks use full
        # sentence-only splitting for natural prosody.
        first_chunk_dispatched = False

        def _sse(data: dict) -> str:
            """Format a dict as an SSE data frame."""
            return f"data: {json.dumps(data, ensure_ascii=False)}\n\n"

        async def _synthesise_and_emit(sentence: str) -> Optional[str]:
            """
            Fire TTS for one sentence.
            Returns a ready-to-yield SSE audio event string, or None if TTS failed.
            """
            audio_bytes, tts_lang = await _fire_tts(sentence, detected_language, tts_gender)
            if audio_bytes:
                return _sse({
                    "type": "audio",
                    "sentence": sentence,
                    "audio_b64": base64.b64encode(audio_bytes).decode("utf-8"),
                    "lang": tts_lang,
                })
            return None

        # Greeting short-circuit path — template reply + TTS, no Groq
        if short_circuit:
            short_reply = get_short_reply(opening_tier, req.message, session.greeting_streak)
            full_text_parts.append(short_reply)
            yield _sse({"type": "token", "text": short_reply})
            audio_event = await _synthesise_and_emit(short_reply)
            if audio_event:
                yield audio_event
        else:
            _start_groq_stream_thread(inference_params, q, loop)

        try:
            if short_circuit:
                pass  # skip Groq loop
            while not short_circuit:
                event_type, data = await q.get()

                # ── Groq stream error: attempt Fireworks non-stream fallback ──
                if event_type == "error":
                    if fireworks is not None:
                        logger.warning(
                            "[stream] Groq stream error — activating Fireworks fallback"
                        )
                        try:
                            fb_params = {**inference_params, "model": FALLBACK_MODEL_FIREWORKS}
                            fb_loop = asyncio.get_running_loop()
                            fb_resp = await fb_loop.run_in_executor(
                                None,
                                lambda: fireworks.chat.completions.create(**fb_params),
                            )
                            fallback_text = fb_resp.choices[0].message.content.strip()
                            full_text_parts.append(fallback_text)
                            yield _sse({"type": "token", "text": fallback_text})
                            # Dispatch TTS for each sentence in the fallback block
                            fb_sentences, _ = _extract_sentences(fallback_text + " ")
                            for sent in fb_sentences:
                                audio_event = await _synthesise_and_emit(sent)
                                if audio_event:
                                    yield audio_event
                        except Exception as fb_exc:
                            logger.error(f"[stream] Fireworks fallback failed: {fb_exc}")
                            yield _sse({"type": "error", "detail": str(fb_exc)})
                            return
                    else:
                        yield _sse({"type": "error", "detail": data})
                        return
                    break

                # ── Stream finished: flush any trailing buffer content ─────────
                elif event_type == "done":
                    remainder = buffer.strip()
                    if remainder:
                        full_text_parts.append(remainder)
                        yield _sse({"type": "token", "text": remainder})
                        audio_event = await _synthesise_and_emit(remainder)
                        if audio_event:
                            yield audio_event
                        first_chunk_dispatched = True
                    break

                # ── Normal token: accumulate + detect sentence boundaries ──────
                else:
                    token: str = data
                    full_text_parts.append(token)
                    yield _sse({"type": "token", "text": token})
                    buffer += token

                    # Issue 4: first chunk uses clause-level (comma)
                    # splitting; every chunk after that uses full-sentence
                    # boundaries only.
                    if not first_chunk_dispatched:
                        clause, buffer = _extract_first_clause(buffer)
                        if clause:
                            audio_event = await _synthesise_and_emit(clause)
                            if audio_event:
                                yield audio_event
                            first_chunk_dispatched = True
                    else:
                        completed_sentences, buffer = _extract_sentences(buffer)
                        for sentence in completed_sentences:
                            audio_event = await _synthesise_and_emit(sentence)
                            if audio_event:
                                yield audio_event

        except asyncio.CancelledError:
            logger.warning(
                f"[stream] Client disconnected mid-stream: {session.session_id}"
            )
            return

        # ── 6. Finalise session state ─────────────────────────────────────────
        full_text = "".join(full_text_parts).strip()
        session.add_message("assistant", full_text)

        elapsed_ms = int((time.time() - start_timestamp) * 1000)
        provider_label = "template-greeting" if short_circuit else primary_model
        logger.info(
            f"[stream] Complete. "
            f"ID={session.session_id} | Model={provider_label} "
            f"| Latency={elapsed_ms}ms | Lang={detected_language}"
        )

        # Final metadata event — same shape as /chat JSON response
        yield _sse({
            "type": "done",
            "session_id": session.session_id,
            "full_text": full_text,
            "language": detected_language,
            "language_locked": session.language_detector.is_locked(),
            "turn_count": session.turn_count,
            "caller_info": session.caller_info.to_dict(),
            "activeSummary": new_summary,
            "elapsed_ms": elapsed_ms,
        })

    return StreamingResponse(
        event_generator(),
        media_type="text/event-stream",
        headers={
            "Cache-Control": "no-cache",
            "X-Accel-Buffering": "no",   # Disable Nginx response buffering
            "Connection": "keep-alive",
        },
    )


@app.post("/session/end")
async def end_session(req: EndSessionRequest):
    """Closes down session trackers and outputs permanent call summaries."""
    summary = manager.end_session(req.session_id)
    if not summary:
        return {
            "status": "already_ended",
            "summary": None,
        }
    return {
        "status": "ended",
        "summary": summary,
    }


@app.get("/health")
async def health():
    """Runs data sweeps and returns systems heartbeat indicators."""
    manager.cleanup_expired()
    
    neon_connected = False
    if neon_client._pool is not None:
        try:
            async with neon_client._pool.acquire() as conn:
                await conn.execute("SELECT 1")
                neon_connected = True
        except Exception:
            pass

    # Dynamically check RAG Space reachability instead of hardcoding True
    rag_ready = False
    RAG_SPACE_URL = os.getenv("RAG_SPACE_URL", "")
    if RAG_SPACE_URL:
        try:
            async with httpx.AsyncClient(timeout=2.0) as client:
                r = await client.get(f"{RAG_SPACE_URL}/health")
                rag_ready = r.status_code == 200
        except Exception:
            pass

    return {
        "status": "ok",
        "version": "1.0.0",
        "active_sessions": manager.active_count(),
        "supported_languages": list(SUPPORTED_LANGUAGES),
        "rag_ready": rag_ready,
        "neon_connected": neon_connected
    }


@app.get("/")
async def root():
    return {
        "service": "The Brain — Voice LLM Core Infrastructure Server Module",
        "status": "online",
        "documentation": "/docs"
    }


@app.on_event("startup")
async def startup_event():
    await neon_client.init_db_pool()


@app.on_event("shutdown")
async def shutdown_event():
    await neon_client.close_db_pool()


@app.websocket("/ws/voice")
async def websocket_voice(websocket: WebSocket):
    await websocket.accept()
    logger.info("New client WebSocket connection request to /ws/voice")

    try:
        start_message = await websocket.receive_json()

        is_setup = start_message.get("event") == "setup"
        is_session_start = start_message.get("type") == "session_start"

        if not is_setup and not is_session_start:
            logger.error("First WebSocket message was not type session_start or event setup")
            await websocket.send_json({"type": "error", "detail": "Session must be initialized with setup or session_start."})
            await websocket.close()
            return

        session_id = start_message.get("session_id")
        if not session_id:
            logger.error("Handshake message missing session_id")
            await websocket.send_json({"type": "error", "detail": "session_id is required."})
            await websocket.close()
            return

        stt_client = SttClient()

        if is_setup:
            company_id = start_message.get("company_id")
            line_token = start_message.get("line_token")

            if not company_id:
                logger.error("Setup handshake missing company_id")
                await websocket.send_json({"type": "error", "detail": "company_id is required."})
                await websocket.close()
                return

            # Create a placeholder session first to construct the orchestrator
            session = ConversationSession(session_id, "Loading...", "Loading...")
            orchestrator = VoiceOrchestratorSession(websocket, session, stt_client)
            orchestrator.is_configured = False

            # Fire database fetch and STT pre-connect concurrently
            config = await neon_client.fetch_tenant_config(company_id, line_token)

            if not config:
                logger.error(f"Tenant config lookup failed: company_id={company_id}, line_token={line_token}")
                await websocket.send_json({"type": "error", "detail": "Company or line not found or inactive."})
                await stt_client.close()
                await websocket.close()
                return

            # Verify line token matching if present in phone line constraints
            if line_token and config.get("line_token") != line_token:
                logger.error(f"Line token mismatch: company_id={company_id}")
                await websocket.send_json({"type": "error", "detail": "Invalid line token configuration."})
                await stt_client.close()
                await websocket.close()
                return

            # Hydrate session from database configuration
            session = ConversationSession.from_tenant_config(config, session_id)
            manager.sessions[session_id] = session
            orchestrator.session = session

            orchestrator.is_configured = True
            await websocket.send_json({
                "type": "configured",
                "company_id": company_id,
                "agent_name": session.agent_name
            })
        else:
            company_name = start_message.get("company_name", "Our Company")
            agent_name = start_message.get("agent_name", "Sara")
            agent_gender = start_message.get("agent_gender", "male")

            session = manager.get_session(session_id)
            if not session:
                session = manager.create_session(
                    company_name=company_name,
                    agent_name=agent_name,
                    agent_gender=agent_gender,
                    session_id=session_id
                )
            orchestrator = VoiceOrchestratorSession(websocket, session, stt_client)
            orchestrator.is_configured = True

        # 4. Start background loops
        await orchestrator.start()

        # 5. Enter ingress loop to process client messages
        await orchestrator.run_ingress_loop()
        
    except WebSocketDisconnect:
        logger.info("Client WebSocket disconnected from /ws/voice")
    except Exception as e:
        logger.exception(f"Unhandled error in voice WebSocket connection: {e}")
        try:
            await websocket.close()
        except Exception:
            pass