pipeline.py

"""
pipeline.py
-----------
Core streaming overlap pipeline.

Flow
----
   Conversation history
          │
          ▼
   SmolLM2 (thread)  ──tokens──►  SentenceBuffer
                                        │
                              sentence complete?
                                        │ yes
                                        ▼
                               MMS-TTS (executor)          ← runs while LLM
                                        │                    keeps generating!
                                        ▼
                              ws.send_bytes(pcm)
                                        │
                              ◄─────────┘ repeat

The key insight: TTS synthesis for sentence N starts as soon as sentence N is
complete — the LLM does NOT wait. By the time the browser finishes playing
sentence N, sentence N+1 is already synthesised and queued, giving near-zero
inter-sentence gaps and a much lower perceived total latency.
"""

from __future__ import annotations

import asyncio
import logging
import time
from threading import Thread
from typing import AsyncGenerator

from fastapi import WebSocket
from transformers import TextIteratorStreamer  # type: ignore

from models import MAX_NEW_TOKENS, ModelManager
from sentence_buffer import SentenceBuffer

logger = logging.getLogger(__name__)

# Maximum conversation turns kept in context (system + N user/assistant pairs)
MAX_HISTORY_TURNS = 6
SYSTEM_PROMPT = (
    "You are a helpful, friendly voice assistant. "
    "Keep every reply concise — two or three short sentences at most. "
    "Speak naturally; avoid bullet points, markdown, or lists."
)


# --------------------------------------------------------------------------- #
#  Async LLM token stream                                                      #
# --------------------------------------------------------------------------- #
async def _llm_token_stream(
    models: ModelManager,
    messages: list[dict],
    loop: asyncio.AbstractEventLoop,
) -> AsyncGenerator[str, None]:
    """
    Async generator that yields text tokens from SmolLM2 as they are produced.

    Architecture
    ~~~~~~~~~~~~
    model.generate()  runs in  gen_thread  (CPU/GPU bound).
    TextIteratorStreamer bridges the sync world → async Queue via
    run_coroutine_threadsafe so the FastAPI event loop stays unblocked.
    """
    prompt = models.build_prompt(messages)
    enc = models.tokenize(prompt)

    streamer = TextIteratorStreamer(
        models.llm_tokenizer,
        skip_prompt=True,
        skip_special_tokens=True,
        timeout=60.0,
    )

    gen_kwargs = dict(
        input_ids=enc["input_ids"],
        attention_mask=enc.get("attention_mask"),
        streamer=streamer,
        max_new_tokens=MAX_NEW_TOKENS,
        do_sample=True,
        temperature=0.7,
        repetition_penalty=1.1,
        pad_token_id=models.llm_tokenizer.eos_token_id,
    )

    token_q: asyncio.Queue[str | None] = asyncio.Queue()

    # ── thread: start generation + drain streamer into async queue ──────── #
    def _generate_and_feed():
        gen_thread = Thread(
            target=models.llm.generate, kwargs=gen_kwargs, daemon=True
        )
        gen_thread.start()
        try:
            for tok in streamer:            # blocks until each token is ready
                asyncio.run_coroutine_threadsafe(token_q.put(tok), loop)
        except Exception as exc:
            asyncio.run_coroutine_threadsafe(token_q.put(None), loop)
            logger.error("LLM stream error: %s", exc)
            return
        asyncio.run_coroutine_threadsafe(token_q.put(None), loop)  # sentinel

    feed_thread = Thread(target=_generate_and_feed, daemon=True)
    feed_thread.start()

    # ── yield tokens from the queue ─────────────────────────────────────── #
    while True:
        tok = await token_q.get()
        if tok is None:
            break
        yield tok


# --------------------------------------------------------------------------- #
#  StreamingPipeline                                                           #
# --------------------------------------------------------------------------- #
class StreamingPipeline:
    """
    Orchestrates the full LLM → SentenceBuffer → TTS → WebSocket pipeline
    with streaming overlap.
    """

    def __init__(self, models: ModelManager):
        self.models = models
        self._executor = None   # Use default loop executor (ThreadPoolExecutor)

    async def process(
        self,
        conversation: list[dict],
        ws: WebSocket,
        loop: asyncio.AbstractEventLoop,
    ) -> str:
        """
        Stream an LLM response over `ws` as audio chunks.

        Returns the full text of the assistant turn (for appending to history).

        Parameters
        ----------
        conversation : list of {role, content} dicts including the new user turn.
        ws           : active WebSocket connection to send PCM bytes to.
        loop         : running event loop (passed for thread-safe queue ops).
        """
        # Trim history to avoid context overflow
        conversation = _trim_history(conversation)

        sentence_buf = SentenceBuffer(min_length=8)
        full_text = ""
        t_start = time.perf_counter()
        ttft_logged = False

        async for token in _llm_token_stream(self.models, conversation, loop):
            if not ttft_logged:
                ttft_ms = (time.perf_counter() - t_start) * 1000
                logger.info("LLM TTFT: %.0f ms", ttft_ms)
                ttft_logged = True

            full_text += token
            ready = sentence_buf.add(token)

            for sentence in ready:
                await self._synth_and_send(sentence, ws)

        # Flush any partial sentence left in buffer
        tail = sentence_buf.flush()
        if tail:
            await self._synth_and_send(tail, ws)

        total_ms = (time.perf_counter() - t_start) * 1000
        logger.info(
            "Pipeline done in %.0f ms | response: %d chars", total_ms, len(full_text)
        )
        return full_text.strip()

    # ------------------------------------------------------------------ #
    async def _synth_and_send(self, sentence: str, ws: WebSocket):
        """
        Synthesise one sentence and stream PCM bytes over the WebSocket.

        TTS runs in the default thread-pool executor so the event loop
        remains free to handle other coroutines (e.g. interruption frames
        arriving from the client) during synthesis.
        """
        sentence = sentence.strip()
        if not sentence:
            return

        logger.debug("TTS ← '%s'", sentence[:80])
        t0 = time.perf_counter()

        pcm: bytes | None = await asyncio.get_event_loop().run_in_executor(
            None, self.models.synthesize, sentence
        )

        if not pcm:
            return

        tts_ms = (time.perf_counter() - t0) * 1000
        logger.info("TTS (%.0f ms): %d bytes for '%s…'", tts_ms, len(pcm), sentence[:40])

        # Send in 4 kB chunks to give the client a smooth receive stream
        CHUNK = 4096
        for i in range(0, len(pcm), CHUNK):
            await ws.send_bytes(pcm[i : i + CHUNK])


# --------------------------------------------------------------------------- #
#  Helpers                                                                     #
# --------------------------------------------------------------------------- #
def _trim_history(messages: list[dict]) -> list[dict]:
    """Keep system message + last MAX_HISTORY_TURNS user/assistant pairs."""
    system = [m for m in messages if m["role"] == "system"]
    turns  = [m for m in messages if m["role"] != "system"]

    # Each "turn" = 1 user + 1 assistant message = 2 items
    max_items = MAX_HISTORY_TURNS * 2
    if len(turns) > max_items:
        turns = turns[-max_items:]

    return system + turns


def build_initial_conversation() -> list[dict]:
    """Start a fresh conversation with only the system prompt."""
    return [{"role": "system", "content": SYSTEM_PROMPT}]