Refactoring the TavusVideoService to allow to work with any transport.

src/pipecat/services/tavus/video.py

@@ -7,10 +7,11 @@
 """This module implements Tavus as a sink transport layer"""
 
 import asyncio
-import base64
+import time
 from typing import Optional
 
 import aiohttp
+from daily.daily import AudioData, VideoFrame
 from loguru import logger
 
 from pipecat.audio.utils import create_default_resampler
@@ -18,19 +19,38 @@ from pipecat.frames.frames import (
     CancelFrame,
     EndFrame,
     Frame,
+    OutputAudioRawFrame,
+    OutputImageRawFrame,
     StartFrame,
     StartInterruptionFrame,
-    TransportMessageUrgentFrame,
     TTSAudioRawFrame,
     TTSStartedFrame,
     TTSStoppedFrame,
 )
-from pipecat.processors.frame_processor import FrameDirection
+from pipecat.processors.frame_processor import FrameDirection, FrameProcessorSetup
 from pipecat.services.ai_service import AIService
+from pipecat.transports.services.tavus import TavusCallbacks, TavusParams, TavusTransportClient
 
 
 class TavusVideoService(AIService):
-    """Class to send base64 encoded audio to Tavus"""
+    """
+    Service class that proxies audio to Tavus and receives both audio and video in return.
+
+    It uses the `TavusTransportClient` to manage the session and handle communication. When
+    audio is sent, Tavus responds with both audio and video streams, which are then routed
+    through Pipecat’s media pipeline.
+
+    In use cases such as with `DailyTransport`, this results in two distinct virtual rooms:
+        - **Tavus room**: Contains the Tavus Avatar and the Pipecat Bot.
+        - **User room**: Contains the Pipecat Bot and the user.
+
+    Args:
+        api_key (str): Tavus API key used for authentication.
+        replica_id (str): ID of the Tavus voice replica to use for speech synthesis.
+        persona_id (str): ID of the Tavus persona. Defaults to "pipecat0" to use the Pipecat TTS voice.
+        session (aiohttp.ClientSession): Async HTTP session used for communication with Tavus.
+        **kwargs: Additional arguments passed to the parent `AIService` class.
+    """
 
     def __init__(
         self,
@@ -39,54 +59,101 @@ class TavusVideoService(AIService):
         replica_id: str,
         persona_id: str = "pipecat0",  # Use `pipecat0` so that your TTS voice is used in place of the Tavus persona
         session: aiohttp.ClientSession,
-        sample_rate: int = 16000,
         **kwargs,
     ) -> None:
         super().__init__(**kwargs)
         self._api_key = api_key
+        self._session = session
         self._replica_id = replica_id
         self._persona_id = persona_id
-        self._session = session
+
-        self._sample_rate = sample_rate
+        self._other_participant_has_joined = False
+        self._client: Optional[TavusTransportClient] = None
 
         self._conversation_id: str
-
         self._resampler = create_default_resampler()
 
         self._audio_buffer = bytearray()
         self._queue = asyncio.Queue()
         self._send_task: Optional[asyncio.Task] = None
 
-    async def initialize(self) -> str:
+    async def setup(self, setup: FrameProcessorSetup):
-        url = "https://tavusapi.com/v2/conversations"
+        await super().setup(setup)
-        headers = {"Content-Type": "application/json", "x-api-key": self._api_key}
+        callbacks = TavusCallbacks(
-        payload = {
+            on_participant_joined=self._on_participant_joined,
-            "replica_id": self._replica_id,
+            on_participant_left=self._on_participant_left,
-            "persona_id": self._persona_id,
+        )
-        }
+        self._client = TavusTransportClient(
-        async with self._session.post(url, headers=headers, json=payload) as r:
+            bot_name="Pipecat",
-            r.raise_for_status()
+            callbacks=callbacks,
-            response_json = await r.json()
+            api_key=self._api_key,
+            replica_id=self._replica_id,
+            persona_id=self._persona_id,
+            session=self._session,
+            params=TavusParams(
+                audio_out_enabled=True,
+                microphone_out_enabled=False,
+                audio_in_enabled=True,
+                video_in_enabled=True,
+                video_out_enabled=True,
+            ),
+        )
+        await self._client.setup(setup)
 
-        logger.debug(f"TavusVideoService joined {response_json['conversation_url']}")
+    async def cleanup(self):
-        self._conversation_id = response_json["conversation_id"]
+        await super().cleanup()
-        return response_json["conversation_url"]
+        await self._client.cleanup()
+        self._client = None
+
+    async def _on_participant_left(self, participant, reason):
+        participant_id = participant["id"]
+        logger.info(f"Participant left {participant_id}, reason: {reason}")
+
+    async def _on_participant_joined(self, participant):
+        participant_id = participant["id"]
+        logger.info(f"Participant joined {participant_id}")
+        if not self._other_participant_has_joined:
+            self._other_participant_has_joined = True
+            await self._client.capture_participant_video(
+                participant_id, self._on_participant_video_frame, 30
+            )
+            await self._client.capture_participant_audio(
+                participant_id=participant_id,
+                callback=self._on_participant_audio_data,
+                sample_rate=self._client.out_sample_rate,
+            )
+
+    async def _on_participant_video_frame(
+        self, participant_id: str, video_frame: VideoFrame, video_source: str
+    ):
+        frame = OutputImageRawFrame(
+            image=video_frame.buffer,
+            size=(video_frame.width, video_frame.height),
+            format=video_frame.color_format,
+        )
+        frame.transport_source = video_source
+        await self.push_frame(frame)
+
+    async def _on_participant_audio_data(
+        self, participant_id: str, audio: AudioData, audio_source: str
+    ):
+        frame = OutputAudioRawFrame(
+            audio=audio.audio_frames,
+            sample_rate=audio.sample_rate,
+            num_channels=audio.num_channels,
+        )
+        frame.transport_source = audio_source
+        await self.push_frame(frame)
 
     def can_generate_metrics(self) -> bool:
         return True
 
     async def get_persona_name(self) -> str:
-        url = f"https://tavusapi.com/v2/personas/{self._persona_id}"
+        return await self._client.get_persona_name()
-        headers = {"Content-Type": "application/json", "x-api-key": self._api_key}
-        async with self._session.get(url, headers=headers) as r:
-            r.raise_for_status()
-            response_json = await r.json()
-
-        logger.debug(f"TavusVideoService persona grabbed {response_json}")
-        return response_json["persona_name"]
 
     async def start(self, frame: StartFrame):
         await super().start(frame)
+        await self._client.start(frame)
         await self._create_send_task()
 
     async def stop(self, frame: EndFrame):
@@ -112,7 +179,7 @@ class TavusVideoService(AIService):
         elif isinstance(frame, TTSAudioRawFrame):
             await self._queue_audio(frame.audio, frame.sample_rate, done=False)
         elif isinstance(frame, TTSStoppedFrame):
-            await self._queue_audio(b"\x00\x00", self._sample_rate, done=True)
+            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:
@@ -121,13 +188,11 @@ class TavusVideoService(AIService):
     async def _handle_interruptions(self):
         await self._cancel_send_task()
         await self._create_send_task()
-        await self._send_interrupt_message()
+        await self._client.send_interrupt_message()
 
     async def _end_conversation(self):
-        url = f"https://tavusapi.com/v2/conversations/{self._conversation_id}/end"
+        await self._client.stop()
-        headers = {"Content-Type": "application/json", "x-api-key": self._api_key}
+        self._other_participant_has_joined = False
-        async with self._session.post(url, headers=headers) as r:
-            r.raise_for_status()
 
     async def _queue_audio(self, audio: bytes, in_rate: int, done: bool):
         await self._queue.put((audio, in_rate, done))
@@ -142,6 +207,15 @@ 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
@@ -149,57 +223,39 @@ class TavusVideoService(AIService):
         # 1 channel). So, that is 48000 / 20 = 2400, which is below the 4kb
         # limit (even including base64 encoding). For a sample rate of 16000,
         # that would be 32000 / 20 = 1600.
-        MAX_CHUNK_SIZE = int((self._sample_rate * 2) / 20)
+        sample_rate = self._client.out_sample_rate
-        SLEEP_TIME = 1 / 20
+        MAX_CHUNK_SIZE = int((sample_rate * 2) / 20)
 
         audio_buffer = bytearray()
+        samples_sent = 0
+        start_time = time.time()
+
         while True:
             (audio, in_rate, done) = await self._queue.get()
 
             if done:
                 # Send any remaining audio.
                 if len(audio_buffer) > 0:
-                    await self._encode_audio_and_send(bytes(audio_buffer), done)
+                    await self._client.encode_audio_and_send(
-                await self._encode_audio_and_send(audio, done)
+                        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, self._sample_rate)
+                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]
                     audio_buffer = audio_buffer[MAX_CHUNK_SIZE:]
-                    await self._encode_audio_and_send(bytes(chunk), done)
-                    await asyncio.sleep(SLEEP_TIME)
 
-    async def _encode_audio_and_send(self, audio: bytes, done: bool):
+                    # Compute wait time for synchronization
-        """Encodes audio to base64 and sends it to Tavus"""
+                    wait = start_time + (samples_sent / sample_rate) - time.time()
-        audio_base64 = base64.b64encode(audio).decode("utf-8")
+                    if wait > 0:
-        logger.trace(f"{self}: sending {len(audio)} bytes")
+                        await asyncio.sleep(wait)
-        await self._send_audio_message(audio_base64, done=done)
 
-    async def _send_interrupt_message(self) -> None:
+                    await self._client.encode_audio_and_send(
-        transport_frame = TransportMessageUrgentFrame(
+                        bytes(chunk), done, self._current_idx_str
-            message={
+                    )
-                "message_type": "conversation",
-                "event_type": "conversation.interrupt",
-                "conversation_id": self._conversation_id,
-            }
-        )
-        await self.push_frame(transport_frame)
 
-    async def _send_audio_message(self, audio_base64: str, done: bool):
+                    # Update timestamp based on number of samples sent
-        transport_frame = TransportMessageUrgentFrame(
+                    samples_sent += len(chunk) // 2  # 2 bytes per sample (16-bit)
-            message={
-                "message_type": "conversation",
-                "event_type": "conversation.echo",
-                "conversation_id": self._conversation_id,
-                "properties": {
-                    "modality": "audio",
-                    "inference_id": self._current_idx_str,
-                    "audio": audio_base64,
-                    "done": done,
-                    "sample_rate": self._sample_rate,
-                },
-            }
-        )
-        await self.push_frame(transport_frame)