""" 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}]