Merge pull request #1898 from pipecat-ai/tavus_transport_fixes

Fixing TavusTransport with some TTS services.

src/pipecat/services/tavus/video.py

@@ -24,13 +24,14 @@ from pipecat.frames.frames import (
     StartFrame,
     StartInterruptionFrame,
     TTSAudioRawFrame,
-    TTSStartedFrame,
-    TTSStoppedFrame,
 )
 from pipecat.processors.frame_processor import FrameDirection, FrameProcessorSetup
 from pipecat.services.ai_service import AIService
 from pipecat.transports.services.tavus import TavusCallbacks, TavusParams, TavusTransportClient
 
+# Using the same values that we do in the BaseOutputTransport
+BOT_VAD_STOP_SECS = 0.35
+
 
 class TavusVideoService(AIService):
     """
@@ -169,16 +170,8 @@ class TavusVideoService(AIService):
         if isinstance(frame, StartInterruptionFrame):
             await self._handle_interruptions()
             await self.push_frame(frame, direction)
-        elif isinstance(frame, TTSStartedFrame):
-            await self.start_processing_metrics()
-            await self.start_ttfb_metrics()
-            self._current_idx_str = str(frame.id)
         elif isinstance(frame, TTSAudioRawFrame):
-            await self._queue_audio(frame.audio, frame.sample_rate, done=False)
+            await self._queue.put(frame)
-        elif isinstance(frame, TTSStoppedFrame):
-            await self._queue_audio(b"\x00\x00", self._client.in_sample_rate, done=True)
-            await self.stop_ttfb_metrics()
-            await self.stop_processing_metrics()
         else:
             await self.push_frame(frame, direction)
 
@@ -191,9 +184,6 @@ class TavusVideoService(AIService):
         await self._client.stop()
         self._other_participant_has_joined = False
 
-    async def _queue_audio(self, audio: bytes, in_rate: int, done: bool):
-        await self._queue.put((audio, in_rate, done))
-
     async def _create_send_task(self):
         if not self._send_task:
             self._queue = asyncio.Queue()
@@ -204,15 +194,6 @@ class TavusVideoService(AIService):
             await self.cancel_task(self._send_task)
             self._send_task = None
 
-    # TODO (Filipi): this should be all that is needed use this Microphone Echo mode
-    # https://docs.tavus.io/sections/conversational-video-interface/layers-and-modes-overview#microphone-echo
-    # This would allow us to send an audio stream for the replica to repeat
-    # Checking with Tavus what is the right way to create the Persona to make it work
-    # async def _send_task_handler(self):
-    #    while True:
-    #        (audio, in_rate, done) = await self._queue.get()
-    #        await self._client.write_raw_audio_frames(audio)
-
     async def _send_task_handler(self):
         # Daily app-messages have a 4kb limit and also a rate limit of 20
         # messages per second. Below, we only consider the rate limit because 1
@@ -221,25 +202,28 @@ class TavusVideoService(AIService):
         # limit (even including base64 encoding). For a sample rate of 16000,
         # that would be 32000 / 20 = 1600.
         sample_rate = self._client.out_sample_rate
+        # 50 ms of audio
         MAX_CHUNK_SIZE = int((sample_rate * 2) / 20)
 
         audio_buffer = bytearray()
+        current_idx_str = None
+        silence = b"\x00" * MAX_CHUNK_SIZE
+        samples_sent = 0
+        start_time = None
+
+        while True:
+            try:
+                frame = await asyncio.wait_for(self._queue.get(), timeout=BOT_VAD_STOP_SECS)
+                if isinstance(frame, TTSAudioRawFrame):
+                    # starting the new inference
+                    if current_idx_str is None:
+                        current_idx_str = str(frame.id)
                         samples_sent = 0
                         start_time = time.time()
 
-        while True:
+                    audio = await self._resampler.resample(
-            (audio, in_rate, done) = await self._queue.get()
+                        frame.audio, frame.sample_rate, sample_rate
-
-            if done:
-                # Send any remaining audio.
-                if len(audio_buffer) > 0:
-                    await self._client.encode_audio_and_send(
-                        bytes(audio_buffer), done, self._current_idx_str
                     )
-                await self._client.encode_audio_and_send(audio, done, self._current_idx_str)
-                audio_buffer.clear()
-            else:
-                audio = await self._resampler.resample(audio, in_rate, sample_rate)
                     audio_buffer.extend(audio)
                     while len(audio_buffer) >= MAX_CHUNK_SIZE:
                         chunk = audio_buffer[:MAX_CHUNK_SIZE]
@@ -248,11 +232,22 @@ class TavusVideoService(AIService):
                         # Compute wait time for synchronization
                         wait = start_time + (samples_sent / sample_rate) - time.time()
                         if wait > 0:
+                            logger.trace(f"TavusVideoService _send_task_handler wait: {wait}")
                             await asyncio.sleep(wait)
 
                         await self._client.encode_audio_and_send(
-                        bytes(chunk), done, self._current_idx_str
+                            bytes(chunk), False, current_idx_str
                         )
 
                         # Update timestamp based on number of samples sent
                         samples_sent += len(chunk) // 2  # 2 bytes per sample (16-bit)
+            except asyncio.TimeoutError:
+                # Bot has stopped speaking
+                # Send any remaining audio.
+                if len(audio_buffer) > 0:
+                    await self._client.encode_audio_and_send(
+                        bytes(audio_buffer), False, current_idx_str
+                    )
+                await self._client.encode_audio_and_send(silence, True, current_idx_str)
+                audio_buffer.clear()
+                current_idx_str = None

src/pipecat/transports/services/tavus.py

@@ -11,6 +11,8 @@ from pydantic import BaseModel
 
 from pipecat.audio.utils import create_default_resampler
 from pipecat.frames.frames import (
+    BotStartedSpeakingFrame,
+    BotStoppedSpeakingFrame,
     CancelFrame,
     EndFrame,
     Frame,
@@ -20,8 +22,6 @@ from pipecat.frames.frames import (
     StartInterruptionFrame,
     TransportMessageFrame,
     TransportMessageUrgentFrame,
-    TTSStartedFrame,
-    TTSStoppedFrame,
 )
 from pipecat.processors.frame_processor import FrameDirection, FrameProcessor, FrameProcessorSetup
 from pipecat.transports.base_input import BaseInputTransport
@@ -365,7 +365,8 @@ class TavusOutputTransport(BaseOutputTransport):
         self._client = client
         self._params = params
         self._samples_sent = 0
-        self._start_time = time.time()
+        self._start_time = None
+        self._current_idx_str: Optional[str] = None
 
     async def setup(self, setup: FrameProcessorSetup):
         await super().setup(setup)
@@ -377,8 +378,6 @@ class TavusOutputTransport(BaseOutputTransport):
 
     async def start(self, frame: StartFrame):
         await super().start(frame)
-        self._samples_sent = 0
-        self._start_time = time.time()
         await self._client.start(frame)
         await self.set_transport_ready(frame)
 
@@ -394,25 +393,42 @@ class TavusOutputTransport(BaseOutputTransport):
         logger.info(f"TavusOutputTransport sending message {frame}")
         await self._client.send_message(frame)
 
+    async def push_frame(self, frame: Frame, direction: FrameDirection = FrameDirection.DOWNSTREAM):
+        # The BotStartedSpeakingFrame and BotStoppedSpeakingFrame are created inside BaseOutputTransport
+        # so TavusOutputTransport never receives these frames.
+        # This is a workaround, so we can more reliably be aware when the bot has started or stopped speaking
+        if direction == FrameDirection.DOWNSTREAM:
+            if isinstance(frame, BotStartedSpeakingFrame):
+                if self._current_idx_str is not None:
+                    logger.warning("TavusOutputTransport self._current_idx_str is already defined!")
+                self._current_idx_str = str(frame.id)
+                self._start_time = time.time()
+                self._samples_sent = 0
+            elif isinstance(frame, BotStoppedSpeakingFrame):
+                silence = b"\x00" * self.audio_chunk_size
+                await self._client.encode_audio_and_send(silence, True, self._current_idx_str)
+                self._current_idx_str = None
+        await super().push_frame(frame, direction)
+
     async def process_frame(self, frame: Frame, direction: FrameDirection):
         await super().process_frame(frame, direction)
         if isinstance(frame, StartInterruptionFrame):
             await self._handle_interruptions()
-        elif isinstance(frame, TTSStartedFrame):
-            self._current_idx_str = str(frame.id)
-        elif isinstance(frame, TTSStoppedFrame):
-            logger.debug(f"TAVUS: {self}: stopped speaking")
-            await self._client.encode_audio_and_send(b"\x00\x00", True, self._current_idx_str)
 
     async def _handle_interruptions(self):
         await self._client.send_interrupt_message()
 
     async def write_raw_audio_frames(self, frames: bytes, destination: Optional[str] = None):
         # Compute wait time for synchronization
-        wait = self._start_time + (self._samples_sent / self._sample_rate) - time.time()
+        wait = self._start_time + (self._samples_sent / self.sample_rate) - time.time()
         if wait > 0:
+            logger.trace(f"TavusOutputTransport write_raw_audio_frames wait: {wait}")
             await asyncio.sleep(wait)
 
+        if self._current_idx_str is None:
+            logger.warning("TavusOutputTransport self._current_idx_str not defined yet!")
+            return
+
         await self._client.encode_audio_and_send(frames, False, self._current_idx_str)
 
         # Update timestamp based on number of samples sent