feat(smallestai): use Waves continuous WebSocket streaming protocol for TTS

livekit-plugins/livekit-plugins-smallestai/livekit/plugins/smallestai/tts.py

@@ -55,6 +55,7 @@ class _TTSOptions:
     language: LanguageCode
     output_format: TTSEncoding | str
     word_timestamps: bool
+    max_buffer_flush_ms: int
     base_url: str
     ws_url: str
 
@@ -71,6 +72,7 @@ def __init__(
         language: str = "en",
         output_format: TTSEncoding | str = "pcm",
         word_timestamps: bool = False,
+        max_buffer_flush_ms: int = 0,
         base_url: str = SMALLEST_BASE_URL,
         ws_url: str = SMALLEST_WS_URL,
         http_session: aiohttp.ClientSession | None = None,
@@ -98,6 +100,14 @@ def __init__(
                 only; HTTP synthesize() returns raw audio without word events. Disabled by
                 default. Supported on base-queue English + Hindi voices (meher, devansh,
                 kartik, maithili, liam, avery); other voices silently emit no word events.
+            max_buffer_flush_ms: Server-side buffer bound (milliseconds) for the continuous
+                WebSocket streaming protocol. As text tokens are streamed in with
+                ``continue: true``, the server accumulates them and forces partial audio
+                output once this many milliseconds of text have buffered, without waiting for
+                an explicit flush. Applies to WebSocket streaming only. ``0`` (default)
+                disables time-based forced flushing, so audio for a segment is produced when
+                the segment's end-of-input flush is sent. Raise it (e.g. 200-400) to trade a
+                little latency for fewer, larger audio chunks.
             base_url: Base URL for the Smallest AI HTTP API.
             ws_url: WebSocket URL for low-latency streaming synthesis.
             http_session: An existing aiohttp ClientSession to use.
@@ -130,6 +140,7 @@ def __init__(
             language=LanguageCode(language),
             output_format=output_format,
             word_timestamps=word_timestamps,
+            max_buffer_flush_ms=max_buffer_flush_ms,
             base_url=base_url,
             ws_url=ws_url,
         )
@@ -180,6 +191,7 @@ def update_options(
         language: NotGivenOr[str] = NOT_GIVEN,
         output_format: NotGivenOr[TTSEncoding | str] = NOT_GIVEN,
         word_timestamps: NotGivenOr[bool] = NOT_GIVEN,
+        max_buffer_flush_ms: NotGivenOr[int] = NOT_GIVEN,
     ) -> None:
         """Update TTS options."""
         if is_given(model):
@@ -197,6 +209,8 @@ def update_options(
         if is_given(word_timestamps):
             self._opts.word_timestamps = word_timestamps
             self._capabilities.aligned_transcript = word_timestamps
+        if is_given(max_buffer_flush_ms):
+            self._opts.max_buffer_flush_ms = max_buffer_flush_ms
 
     def synthesize(
         self,
@@ -272,106 +286,129 @@ async def _run(self, output_emitter: tts.AudioEmitter) -> None:
 
 
 class SynthesizeStream(tts.SynthesizeStream):
-    """WebSocket-based streaming synthesis — primary path used by the agent pipeline."""
+    """WebSocket-based streaming synthesis — primary path used by the agent pipeline.
+
+    Uses the Waves continuous-streaming protocol: text tokens are forwarded to the
+    server as they arrive (``continue: true``) rather than buffered locally, and an
+    explicit ``flush: true`` message closes each segment. This lets synthesis begin
+    before the full segment text is known, lowering time-to-first-byte.
+    """
 
     def __init__(self, *, tts: TTS, conn_options: APIConnectOptions) -> None:
         super().__init__(tts=tts, conn_options=conn_options)
         self._tts: TTS = tts
         self._opts = replace(tts._opts)
 
     async def _run(self, output_emitter: tts.AudioEmitter) -> None:
+        request_id = utils.shortuuid()
         output_emitter.initialize(
-            request_id=utils.shortuuid(),
+            request_id=request_id,
             sample_rate=self._opts.sample_rate,
             num_channels=NUM_CHANNELS,
             mime_type="audio/pcm",
             stream=True,
         )
+        # One stream instance = one segment (the agent creates a new stream per segment).
+        output_emitter.start_segment(segment_id=request_id)
 
         try:
-            text_buffer = ""
-            async for data in self._input_ch:
-                if isinstance(data, str):
-                    text_buffer += data
-                elif isinstance(data, self._FlushSentinel):
-                    if text_buffer.strip():
-                        await self._run_ws(text_buffer.strip(), output_emitter)
-                    text_buffer = ""
+            async with self._tts._pool.connection(timeout=self._conn_options.timeout) as ws:
+                self._acquire_time = self._tts._pool.last_acquire_time
+                self._connection_reused = self._tts._pool.last_connection_reused
+
+                send_task = asyncio.create_task(self._send_task(ws))
+                recv_task = asyncio.create_task(self._recv_task(ws, output_emitter))
+                try:
+                    await asyncio.gather(send_task, recv_task)
+                finally:
+                    await utils.aio.gracefully_cancel(send_task, recv_task)
         except asyncio.TimeoutError:
             raise APITimeoutError() from None
         except aiohttp.ClientResponseError as e:
             raise APIStatusError(
-                message=e.message, status_code=e.status, request_id=None, body=None
+                message=e.message, status_code=e.status, request_id=request_id, body=None
             ) from None
         except APIStatusError:
             raise
         except Exception as e:
             raise APIConnectionError() from e
 
-    async def _run_ws(self, text: str, output_emitter: tts.AudioEmitter) -> None:
-        segment_id = utils.shortuuid()
-        output_emitter.start_segment(segment_id=segment_id)
+        output_emitter.end_segment()
 
+    def _base_payload(self) -> dict[str, Any]:
         payload: dict[str, Any] = {
             "model": self._opts.model,
             "voice_id": self._opts.voice_id,
-            "text": text,
             "sample_rate": self._opts.sample_rate,
             "speed": self._opts.speed,
             "language": self._opts.language.language
             if isinstance(self._opts.language, LanguageCode)
             else self._opts.language,
+            "max_buffer_flush_ms": self._opts.max_buffer_flush_ms,
         }
-
         if self._opts.word_timestamps:
             payload["word_timestamps"] = True
-
-        async with self._tts._pool.connection(timeout=self._conn_options.timeout) as ws:
-            self._acquire_time = self._tts._pool.last_acquire_time
-            self._connection_reused = self._tts._pool.last_connection_reused
+        return payload
+
+    async def _send_task(self, ws: aiohttp.ClientWebSocketResponse) -> None:
+        # Forward each token as it arrives (continuous streaming) instead of buffering
+        # the whole segment, so the server can begin synthesis before the text is
+        # complete. The flush sentinel (or end of input) closes the segment.
+        async for data in self._input_ch:
+            if isinstance(data, self._FlushSentinel):
+                break
+            if not data.strip():
+                continue
+            token_pkt = {**self._base_payload(), "text": data, "continue": True, "flush": False}
             self._mark_started()
-            await ws.send_str(json.dumps(payload))
-
-            while True:
-                msg = await ws.receive(timeout=self._conn_options.timeout)
-
-                if msg.type in (
-                    aiohttp.WSMsgType.CLOSE,
-                    aiohttp.WSMsgType.CLOSED,
-                    aiohttp.WSMsgType.CLOSING,
-                ):
-                    raise APIStatusError(
-                        "SmallestAI WebSocket closed unexpectedly",
-                        status_code=ws.close_code or -1,
-                        body=str(msg.data),
-                    )
+            await ws.send_str(json.dumps(token_pkt))
+
+        # Flush the server-side buffer to force out any remaining audio for the segment.
+        flush_pkt = {**self._base_payload(), "text": "", "continue": False, "flush": True}
+        await ws.send_str(json.dumps(flush_pkt))
 
-                if msg.type != aiohttp.WSMsgType.TEXT:
-                    continue
-
-                event = json.loads(msg.data)
-                status = event.get("status")
-
-                if status == "chunk":
-                    audio_b64 = event.get("data", {}).get("audio")
-                    if audio_b64:
-                        output_emitter.push(base64.b64decode(audio_b64))
-                elif status == "word_timestamp":
-                    data = event.get("data", {})
-                    word = data.get("word")
-                    start = data.get("start")
-                    end = data.get("end")
-                    if word is not None and start is not None and end is not None:
-                        output_emitter.push_timed_transcript(
-                            TimedString(text=word, start_time=start, end_time=end)
-                        )
-                elif status == "complete":
-                    output_emitter.end_segment()
-                    break
-                elif status == "error":
-                    raise APIConnectionError(
-                        f"SmallestAI TTS error: {event.get('message', 'unknown error')}"
+    async def _recv_task(
+        self, ws: aiohttp.ClientWebSocketResponse, output_emitter: tts.AudioEmitter
+    ) -> None:
+        while True:
+            msg = await ws.receive(timeout=self._conn_options.timeout)
+
+            if msg.type in (
+                aiohttp.WSMsgType.CLOSE,
+                aiohttp.WSMsgType.CLOSED,
+                aiohttp.WSMsgType.CLOSING,
+            ):
+                raise APIStatusError(
+                    "SmallestAI WebSocket closed unexpectedly",
+                    status_code=ws.close_code or -1,
+                    body=str(msg.data),
+                )
+
+            if msg.type != aiohttp.WSMsgType.TEXT:
+                continue
+
+            event = json.loads(msg.data)
+            status = event.get("status")
+
+            if status == "chunk":
+                audio_b64 = event.get("data", {}).get("audio")
+                if audio_b64:
+                    output_emitter.push(base64.b64decode(audio_b64))
+            elif status == "word_timestamp":
+                data = event.get("data", {})
+                word = data.get("word")
+                start = data.get("start")
+                end = data.get("end")
+                if word is not None and start is not None and end is not None:
+                    output_emitter.push_timed_transcript(
+                        TimedString(text=word, start_time=start, end_time=end)
                     )
+            elif status == "complete":
+                break
+            elif status == "error":
+                raise APIConnectionError(
+                    f"SmallestAI TTS error: {event.get('message', 'unknown error')}"
+                )
 
 
 def _to_smallest_options(opts: _TTSOptions) -> dict[str, Any]: