Merge pull request #3759 from pipecat-ai/mb/gradium-context-update

Switch Gradium TTS to AudioContextWordTTSService for multiplexing

changelog/3759.performance.md

@@ -0,0 +1 @@
+- Switched `GradiumTTSService` from `InterruptibleWordTTSService` to `AudioContextWordTTSService`, eliminating websocket disconnect/reconnect on every interruption by using `client_req_id`-based multiplexing.

src/pipecat/services/gradium/tts.py

@@ -6,6 +6,7 @@
 
 import base64
 import json
+import uuid
 from typing import Any, AsyncGenerator, Mapping, Optional
 
 from loguru import logger
@@ -16,13 +17,14 @@ from pipecat.frames.frames import (
     EndFrame,
     ErrorFrame,
     Frame,
+    InterruptionFrame,
     StartFrame,
     TTSAudioRawFrame,
     TTSStartedFrame,
     TTSStoppedFrame,
 )
 from pipecat.processors.frame_processor import FrameDirection
-from pipecat.services.tts_service import InterruptibleWordTTSService
+from pipecat.services.tts_service import AudioContextWordTTSService
 from pipecat.utils.tracing.service_decorators import traced_tts
 
 try:
@@ -37,7 +39,7 @@ except ModuleNotFoundError as e:
 SAMPLE_RATE = 48000
 
 
-class GradiumTTSService(InterruptibleWordTTSService):
+class GradiumTTSService(AudioContextWordTTSService):
     """Text-to-Speech service using Gradium's websocket API."""
 
     class InputParams(BaseModel):
@@ -71,9 +73,9 @@ class GradiumTTSService(InterruptibleWordTTSService):
             params: Additional configuration parameters.
             **kwargs: Additional arguments passed to parent class.
         """
-        # Initialize with parent class settings for proper frame handling
         super().__init__(
             push_stop_frames=True,
+            push_text_frames=False,
             pause_frame_processing=True,
             sample_rate=SAMPLE_RATE,
             **kwargs,
@@ -95,7 +97,7 @@ class GradiumTTSService(InterruptibleWordTTSService):
 
         # State tracking
         self._receive_task = None
-        self._current_context_id: Optional[str] = None
+        self._context_id: Optional[str] = None
 
     def can_generate_metrics(self) -> bool:
         """Check if this service can generate processing metrics.
@@ -126,7 +128,10 @@ class GradiumTTSService(InterruptibleWordTTSService):
 
     def _build_msg(self, text: str = "") -> dict:
         """Build JSON message for Gradium API."""
-        return {"text": text, "type": "text"}
+        msg = {"text": text, "type": "text"}
+        if self._context_id:
+            msg["client_req_id"] = self._context_id
+        return msg
 
     async def start(self, frame: StartFrame):
         """Start the service and establish websocket connection.
@@ -197,6 +202,7 @@ class GradiumTTSService(InterruptibleWordTTSService):
                 "type": "setup",
                 "output_format": "pcm",
                 "voice_id": self._voice_id,
+                "close_ws_on_eos": False,
             }
             if self._json_config is not None:
                 setup_msg["json_config"] = self._json_config
@@ -234,18 +240,35 @@ class GradiumTTSService(InterruptibleWordTTSService):
 
     async def flush_audio(self):
         """Flush any pending audio synthesis."""
-        if not self._websocket:
+        if not self._context_id or not self._websocket:
             return
         try:
-            msg = {"type": "end_of_stream"}
+            msg = {"type": "end_of_stream", "client_req_id": self._context_id}
             await self._websocket.send(json.dumps(msg))
+            self._context_id = None
         except ConnectionClosedOK:
             logger.debug(f"{self}: connection closed normally during flush")
         except Exception as e:
             logger.error(f"{self} exception: {e}")
 
+    async def _handle_interruption(self, frame: InterruptionFrame, direction: FrameDirection):
+        """Handle interruption by resetting context state.
+
+        The parent AudioContextTTSService._handle_interruption() cancels the audio context
+        task and creates a new one. We reset _context_id so the next run_tts() creates a
+        fresh context. No websocket reconnection needed — audio from the old client_req_id
+        will be silently dropped since the audio context no longer exists.
+
+        Args:
+            frame: The interruption frame.
+            direction: The direction of the frame.
+        """
+        await super()._handle_interruption(frame, direction)
+        await self.stop_all_metrics()
+        self._context_id = None
+
     async def _receive_messages(self):
-        """Process incoming websocket messages."""
+        """Process incoming websocket messages, demultiplexing by client_req_id."""
         # TODO(laurent): This should not be necessary as it should happen when
         # receiving the messages but this does not seem to always be the case
         # and that may lead to a busy polling loop.
@@ -253,41 +276,49 @@ class GradiumTTSService(InterruptibleWordTTSService):
             raise ConnectionClosedOK(None, None)
         async for message in self._get_websocket():
             msg = json.loads(message)
+            ctx_id = msg.get("client_req_id")
 
             if msg["type"] == "audio":
-                # Process audio chunk
+                if not ctx_id or not self.audio_context_available(ctx_id):
+                    continue
                 await self.stop_ttfb_metrics()
                 await self.start_word_timestamps()
                 frame = TTSAudioRawFrame(
                     audio=base64.b64decode(msg["audio"]),
                     sample_rate=self.sample_rate,
                     num_channels=1,
-                    context_id=self._current_context_id,
+                    context_id=ctx_id,
                 )
-                await self.push_frame(frame)
+                await self.append_to_audio_context(ctx_id, frame)
 
             elif msg["type"] == "text":
-                if self._current_context_id:
+                if ctx_id and self.audio_context_available(ctx_id):
-                    await self.add_word_timestamps(
+                    await self.add_word_timestamps([(msg["text"], msg["start_s"])], ctx_id)
-                        [(msg["text"], msg["start_s"])], self._current_context_id
+
-                    )
             elif msg["type"] == "end_of_stream":
-                await self.push_frame(TTSStoppedFrame())
+                if ctx_id and self.audio_context_available(ctx_id):
+                    await self.add_word_timestamps([("TTSStoppedFrame", 0), ("Reset", 0)], ctx_id)
+                    await self.remove_audio_context(ctx_id)
                 await self.stop_all_metrics()
 
             elif msg["type"] == "error":
-                await self.push_frame(TTSStoppedFrame())
+                await self.push_frame(TTSStoppedFrame(context_id=ctx_id))
                 await self.stop_all_metrics()
-                await self.push_error(error_msg=f"Error: {msg['message']}")
+                await self.push_error(error_msg=f"Error: {msg.get('message', msg)}")
 
-    async def push_frame(self, frame: Frame, direction: FrameDirection = FrameDirection.DOWNSTREAM):
+    def create_context_id(self) -> str:
-        """Push frame and handle end-of-turn conditions.
+        """Generate a unique context ID for a TTS request in case we don't have one already in progress.
 
-        Args:
+        Returns:
-            frame: The frame to push.
+            A unique string identifier for the TTS context.
-            direction: The direction to push the frame.
         """
-        await super().push_frame(frame, direction)
+        # If a context ID does not exist, create a new one.
+        # If an ID exists, continue using the current ID.
+        # When interruptions happens, user speech results in
+        # an interruption, which resets the context ID.
+        if not self._context_id:
+            return str(uuid.uuid4())
+        return self._context_id
 
     @traced_tts
     async def run_tts(self, text: str, context_id: str) -> AsyncGenerator[Frame, None]:
@@ -300,16 +331,18 @@ class GradiumTTSService(InterruptibleWordTTSService):
         Yields:
             Frame: Audio frames containing the synthesized speech.
         """
-        _state = self._websocket.state if self._websocket is not None else None
+        logger.debug(f"{self}: Generating TTS [{text}]")
-        logger.debug(f"{self}: Generating TTS [{text}] {_state}")
         try:
             if not self._websocket or self._websocket.state is State.CLOSED:
                 self._websocket = None
                 await self._connect()
 
             try:
-                self._current_context_id = context_id
+                if not self._context_id:
-                yield TTSStartedFrame(context_id=context_id)
+                    await self.start_ttfb_metrics()
+                    yield TTSStartedFrame(context_id=context_id)
+                    self._context_id = context_id
+                    await self.create_audio_context(self._context_id)
 
                 msg = self._build_msg(text=text)
                 await self._get_websocket().send(json.dumps(msg))