wx44wx/pipecat
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Refactoring the TavusVideoService to allow to work with any transport.

Browse Source
This commit is contained in:
Filipi Fuchter
2025-05-23 23:02:49 -03:00
parent 86e6841569
commit e0d1381f87
1 changed files with 128 additions and 72 deletions
Download Patch File Download Diff File Expand all files Collapse all files

200
src/pipecat/services/tavus/video.py
View File

@@ -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)