Merge pull request #1223 from pipecat-ai/aleix/audio-context-tts-service

audio context tts service and cartesia fixes

CHANGELOG.md

@@ -9,6 +9,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Added
 
+- Added new `AudioContextWordTTSService`. This is a TTS base class for TTS
+  services that handling multiple separate audio requests.
+
 - Added new frames `EmulateUserStartedSpeakingFrame` and
   `EmulateUserStoppedSpeakingFrame` which can be used to emulated VAD behavior
   without VAD being present or not being triggered.
@@ -102,6 +105,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Fixed
 
+- Fixed a `CartesiaTTSService` service issue that would cause audio overlapping
+  in some cases.
+
 - Fixed a websocket-based service issue (e.g. `CartesiaTTSService`) that was
   preventing a reconnection after the server disconnected cleanly, which was
   causing an inifite loop instead.

src/pipecat/services/ai_services.py

@@ -419,7 +419,7 @@ class WordTTSService(TTSService):
 
     async def start(self, frame: StartFrame):
         await super().start(frame)
-        await self._create_words_task()
+        self._create_words_task()
 
     async def stop(self, frame: EndFrame):
         await super().stop(frame)
@@ -439,7 +439,7 @@ class WordTTSService(TTSService):
         await super()._handle_interruption(frame, direction)
         self.reset_word_timestamps()
 
-    async def _create_words_task(self):
+    def _create_words_task(self):
         self._words_task = self.create_task(self._words_task_handler())
 
     async def _stop_words_task(self):
@@ -469,6 +469,115 @@ class WordTTSService(TTSService):
             self._words_queue.task_done()
 
 
+class AudioContextWordTTSService(WordTTSService):
+    """This services allow us to send multiple TTS request to the services. Each
+    request could be multiple sentences long which are grouped by context. For
+    this to work, the TTS service needs to support handling multiple requests at
+    once (i.e. multiple simultaneous contexts).
+
+    The audio received from the TTS will be played in context order. That is, if
+    we requested audio for a context "A" and then audio for context "B", the
+    audio from context ID "A" will be played first.
+
+    """
+
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+        self._contexts_queue = asyncio.Queue()
+        self._contexts: Dict[str, asyncio.Queue] = {}
+        self._audio_context_task = None
+
+    async def create_audio_context(self, context_id: str):
+        """Create a new audio context."""
+        await self._contexts_queue.put(context_id)
+        self._contexts[context_id] = asyncio.Queue()
+        logger.trace(f"{self} created audio context {context_id}")
+
+    async def append_to_audio_context(self, context_id: str, frame: TTSAudioRawFrame):
+        """Append audio to an existing context."""
+        if self.audio_context_available(context_id):
+            logger.trace(f"{self} appending audio {frame} to audio context {context_id}")
+            await self._contexts[context_id].put(frame)
+        else:
+            logger.warning(f"{self} unable to append audio to context {context_id}")
+
+    async def remove_audio_context(self, context_id: str):
+        """Remove an existing audio context."""
+        if self.audio_context_available(context_id):
+            # We just mark the audio context for deletion by appending
+            # None. Once we reach None while handling audio we know we can
+            # safely remove the context.
+            logger.trace(f"{self} marking audio context {context_id} for deletion")
+            await self._contexts[context_id].put(None)
+        else:
+            logger.warning(f"{self} unable to remove context {context_id}")
+
+    def audio_context_available(self, context_id: str) -> bool:
+        """Checks whether the given audio context is registered."""
+        return context_id in self._contexts
+
+    async def start(self, frame: StartFrame):
+        await super().start(frame)
+        self._create_audio_context_task()
+
+    async def stop(self, frame: EndFrame):
+        await super().stop(frame)
+        await self._stop_audio_context_task()
+
+    async def cancel(self, frame: CancelFrame):
+        await super().cancel(frame)
+        await self._stop_audio_context_task()
+
+    async def _handle_interruption(self, frame: StartInterruptionFrame, direction: FrameDirection):
+        await super()._handle_interruption(frame, direction)
+        await self._stop_audio_context_task()
+        self._create_audio_context_task()
+
+    def _create_audio_context_task(self):
+        self._contexts_queue = asyncio.Queue()
+        self._contexts: Dict[str, asyncio.Queue] = {}
+        self._audio_context_task = self.create_task(self._audio_context_task_handler())
+
+    async def _stop_audio_context_task(self):
+        if self._audio_context_task:
+            await self.cancel_task(self._audio_context_task)
+            self._audio_context_task = None
+
+    async def _audio_context_task_handler(self):
+        """In this task we process audio contexts in order."""
+        while True:
+            context_id = await self._contexts_queue.get()
+
+            # Process the audio context until the context doesn't have more
+            # audio available (i.e. we find None).
+            await self._handle_audio_context(context_id)
+
+            # We just finished processing the context, so we can safely remove it.
+            del self._contexts[context_id]
+            self._contexts_queue.task_done()
+
+            # Append some silence between sentences.
+            silence = b"\x00" * self.sample_rate
+            frame = TTSAudioRawFrame(audio=silence, sample_rate=self.sample_rate, num_channels=1)
+            await self.push_frame(frame)
+
+    async def _handle_audio_context(self, context_id: str):
+        # If we don't receive any audio during this time, we consider the context finished.
+        AUDIO_CONTEXT_TIMEOUT = 3.0
+        queue = self._contexts[context_id]
+        running = True
+        while running:
+            try:
+                frame = await asyncio.wait_for(queue.get(), timeout=AUDIO_CONTEXT_TIMEOUT)
+                if frame:
+                    await self.push_frame(frame)
+                running = frame is not None
+            except asyncio.TimeoutError:
+                # We didn't get audio, so let's consider this context finished.
+                logger.trace(f"{self} time out on audio context {context_id}")
+                break
+
+
 class STTService(AIService):
     """STTService is a base class for speech-to-text services."""
 

src/pipecat/services/cartesia.py

@@ -27,7 +27,7 @@ from pipecat.frames.frames import (
     TTSStoppedFrame,
 )
 from pipecat.processors.frame_processor import FrameDirection
-from pipecat.services.ai_services import TTSService, WordTTSService
+from pipecat.services.ai_services import AudioContextWordTTSService, TTSService
 from pipecat.services.websocket_service import WebsocketService
 from pipecat.transcriptions.language import Language
 
@@ -75,7 +75,7 @@ def language_to_cartesia_language(language: Language) -> Optional[str]:
     return result
 
 
-class CartesiaTTSService(WordTTSService, WebsocketService):
+class CartesiaTTSService(AudioContextWordTTSService, WebsocketService):
     class InputParams(BaseModel):
         language: Optional[Language] = Language.EN
         speed: Optional[Union[str, float]] = ""
@@ -105,7 +105,7 @@ class CartesiaTTSService(WordTTSService, WebsocketService):
         # if we're interrupted. Cartesia gives us word-by-word timestamps. We
         # can use those to generate text frames ourselves aligned with the
         # playout timing of the audio!
-        WordTTSService.__init__(
+        AudioContextWordTTSService.__init__(
             self,
             aggregate_sentences=True,
             push_text_frames=False,
@@ -191,12 +191,12 @@ class CartesiaTTSService(WordTTSService, WebsocketService):
         self._receive_task = self.create_task(self._receive_task_handler(self.push_error))
 
     async def _disconnect(self):
-        await self._disconnect_websocket()
-
         if self._receive_task:
             await self.cancel_task(self._receive_task)
             self._receive_task = None
 
+        await self._disconnect_websocket()
+
     async def _connect_websocket(self):
         try:
             logger.debug("Connecting to Cartesia")
@@ -239,21 +239,19 @@ class CartesiaTTSService(WordTTSService, WebsocketService):
         logger.trace(f"{self}: flushing audio")
         msg = self._build_msg(text="", continue_transcript=False)
         await self._websocket.send(msg)
+        self._context_id = None
 
     async def _receive_messages(self):
         async for message in self._get_websocket():
             msg = json.loads(message)
-            if not msg or msg["context_id"] != self._context_id:
+            if not msg or not self.audio_context_available(msg["context_id"]):
                 continue
             if msg["type"] == "done":
                 await self.stop_ttfb_metrics()
-                # Unset _context_id but not the _context_id_start_timestamp
-                # because we are likely still playing out audio and need the
-                # timestamp to set send context frames.
-                self._context_id = None
                 await self.add_word_timestamps(
                     [("TTSStoppedFrame", 0), ("LLMFullResponseEndFrame", 0), ("Reset", 0)]
                 )
+                await self.remove_audio_context(msg["context_id"])
             elif msg["type"] == "timestamps":
                 await self.add_word_timestamps(
                     list(zip(msg["word_timestamps"]["words"], msg["word_timestamps"]["start"]))
@@ -266,12 +264,13 @@ class CartesiaTTSService(WordTTSService, WebsocketService):
                     sample_rate=self.sample_rate,
                     num_channels=1,
                 )
-                await self.push_frame(frame)
+                await self.append_to_audio_context(msg["context_id"], frame)
             elif msg["type"] == "error":
                 logger.error(f"{self} error: {msg}")
                 await self.push_frame(TTSStoppedFrame())
                 await self.stop_all_metrics()
                 await self.push_error(ErrorFrame(f"{self} error: {msg['error']}"))
+                self._context_id = None
             else:
                 logger.error(f"{self} error, unknown message type: {msg}")
 
@@ -299,6 +298,7 @@ class CartesiaTTSService(WordTTSService, WebsocketService):
                 await self.start_ttfb_metrics()
                 yield TTSStartedFrame()
                 self._context_id = str(uuid.uuid4())
+                await self.create_audio_context(self._context_id)
 
             msg = self._build_msg(text=text or " ")  # Text must contain at least one character
 

src/pipecat/services/rime.py

@@ -7,7 +7,7 @@
 import base64
 import json
 import uuid
-from typing import AsyncGenerator, Optional, Union
+from typing import AsyncGenerator, Optional
 
 import aiohttp
 from loguru import logger
@@ -28,7 +28,7 @@ from pipecat.frames.frames import (
     TTSStoppedFrame,
 )
 from pipecat.processors.frame_processor import FrameDirection
-from pipecat.services.ai_services import TTSService, WordTTSService
+from pipecat.services.ai_services import AudioContextWordTTSService, TTSService
 from pipecat.services.websocket_service import WebsocketService
 from pipecat.transcriptions.language import Language
 
@@ -58,7 +58,7 @@ def language_to_rime_language(language: Language) -> str:
     return LANGUAGE_MAP.get(language, "eng")
 
 
-class RimeTTSService(WordTTSService, WebsocketService):
+class RimeTTSService(AudioContextWordTTSService, WebsocketService):
     """Text-to-Speech service using Rime's websocket API.
 
     Uses Rime's websocket JSON API to convert text to speech with word-level timing
@@ -95,7 +95,7 @@ class RimeTTSService(WordTTSService, WebsocketService):
             params: Additional configuration parameters.
         """
         # Initialize with parent class settings for proper frame handling
-        WordTTSService.__init__(
+        AudioContextWordTTSService.__init__(
             self,
             aggregate_sentences=True,
             push_text_frames=False,
@@ -249,12 +249,18 @@ class RimeTTSService(WordTTSService, WebsocketService):
 
         return word_pairs
 
+    async def flush_audio(self):
+        if not self._context_id or not self._websocket:
+            return
+        logger.trace(f"{self}: flushing audio")
+        self._context_id = None
+
     async def _receive_messages(self):
         """Process incoming websocket messages."""
         async for message in self._get_websocket():
             msg = json.loads(message)
 
-            if not msg or msg["contextId"] != self._context_id:
+            if not msg or not self.audio_context_available(msg["contextId"]):
                 continue
 
             if msg["type"] == "chunk":
@@ -266,7 +272,7 @@ class RimeTTSService(WordTTSService, WebsocketService):
                     sample_rate=self.sample_rate,
                     num_channels=1,
                 )
-                await self.push_frame(frame)
+                await self.append_to_audio_context(msg["contextId"], frame)
 
             elif msg["type"] == "timestamps":
                 # Process word timing information
@@ -288,6 +294,7 @@ class RimeTTSService(WordTTSService, WebsocketService):
                 await self.push_frame(TTSStoppedFrame())
                 await self.stop_all_metrics()
                 await self.push_error(ErrorFrame(f"{self} error: {msg['message']}"))
+                self._context_id = None
 
     async def push_frame(self, frame: Frame, direction: FrameDirection = FrameDirection.DOWNSTREAM):
         """Push frame and handle end-of-turn conditions."""
@@ -329,6 +336,7 @@ class RimeTTSService(WordTTSService, WebsocketService):
                     self._started = True
                     self._cumulative_time = 0
                     self._context_id = str(uuid.uuid4())
+                    await self.create_audio_context(self._context_id)
 
                 msg = self._build_msg(text=text)
                 await self._get_websocket().send(json.dumps(msg))