From e0d1381f8777eee8fd0e6a25661874f5ae9a90a3 Mon Sep 17 00:00:00 2001 From: Filipi Fuchter Date: Fri, 23 May 2025 23:02:49 -0300 Subject: [PATCH] Refactoring the TavusVideoService to allow to work with any transport. --- src/pipecat/services/tavus/video.py | 200 ++++++++++++++++++---------- 1 file changed, 128 insertions(+), 72 deletions(-) diff --git a/src/pipecat/services/tavus/video.py b/src/pipecat/services/tavus/video.py index 3699ba512..2d7ba6cd4 100644 --- a/src/pipecat/services/tavus/video.py +++ b/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: - url = "https://tavusapi.com/v2/conversations" - headers = {"Content-Type": "application/json", "x-api-key": self._api_key} - payload = { - "replica_id": self._replica_id, - "persona_id": self._persona_id, - } - async with self._session.post(url, headers=headers, json=payload) as r: - r.raise_for_status() - response_json = await r.json() + async def setup(self, setup: FrameProcessorSetup): + await super().setup(setup) + callbacks = TavusCallbacks( + on_participant_joined=self._on_participant_joined, + on_participant_left=self._on_participant_left, + ) + self._client = TavusTransportClient( + bot_name="Pipecat", + callbacks=callbacks, + 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']}") - self._conversation_id = response_json["conversation_id"] - return response_json["conversation_url"] + async def cleanup(self): + await super().cleanup() + 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}" - 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"] + return await self._client.get_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" - headers = {"Content-Type": "application/json", "x-api-key": self._api_key} - async with self._session.post(url, headers=headers) as r: - r.raise_for_status() + 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)) @@ -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) - SLEEP_TIME = 1 / 20 + sample_rate = self._client.out_sample_rate + 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._encode_audio_and_send(audio, done) + 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, 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): - """Encodes audio to base64 and sends it to Tavus""" - audio_base64 = base64.b64encode(audio).decode("utf-8") - logger.trace(f"{self}: sending {len(audio)} bytes") - await self._send_audio_message(audio_base64, done=done) + # Compute wait time for synchronization + wait = start_time + (samples_sent / sample_rate) - time.time() + if wait > 0: + await asyncio.sleep(wait) - async def _send_interrupt_message(self) -> None: - transport_frame = TransportMessageUrgentFrame( - message={ - "message_type": "conversation", - "event_type": "conversation.interrupt", - "conversation_id": self._conversation_id, - } - ) - await self.push_frame(transport_frame) + await self._client.encode_audio_and_send( + bytes(chunk), done, self._current_idx_str + ) - async def _send_audio_message(self, audio_base64: str, done: bool): - transport_frame = TransportMessageUrgentFrame( - 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) + # Update timestamp based on number of samples sent + samples_sent += len(chunk) // 2 # 2 bytes per sample (16-bit)