diff --git a/CHANGELOG.md b/CHANGELOG.md index 13577e342..e3976af2d 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -27,6 +27,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ### Fixed +- Fixed a `TavusVideoService` issue that was causing audio choppiness. + - Fixed an issue in `SmallWebRTCTransport` where an error was thrown if the client did not create a video transceiver. diff --git a/src/pipecat/services/tavus/video.py b/src/pipecat/services/tavus/video.py index cbc31a2ba..3699ba512 100644 --- a/src/pipecat/services/tavus/video.py +++ b/src/pipecat/services/tavus/video.py @@ -6,7 +6,9 @@ """This module implements Tavus as a sink transport layer""" +import asyncio import base64 +from typing import Optional import aiohttp from loguru import logger @@ -16,6 +18,7 @@ from pipecat.frames.frames import ( CancelFrame, EndFrame, Frame, + StartFrame, StartInterruptionFrame, TransportMessageUrgentFrame, TTSAudioRawFrame, @@ -50,6 +53,10 @@ class TavusVideoService(AIService): 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} @@ -78,45 +85,98 @@ class TavusVideoService(AIService): 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._create_send_task() + async def stop(self, frame: EndFrame): await super().stop(frame) await self._end_conversation() + await self._cancel_send_task() async def cancel(self, frame: CancelFrame): await super().cancel(frame) await self._end_conversation() + await self._cancel_send_task() - async def _end_conversation(self) -> None: + async def process_frame(self, frame: Frame, direction: FrameDirection): + await super().process_frame(frame, direction) + + 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) + elif isinstance(frame, TTSStoppedFrame): + await self._queue_audio(b"\x00\x00", self._sample_rate, done=True) + await self.stop_ttfb_metrics() + await self.stop_processing_metrics() + else: + await self.push_frame(frame, direction) + + async def _handle_interruptions(self): + await self._cancel_send_task() + await self._create_send_task() + await self._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() - async def _encode_audio_and_send(self, audio: bytes, in_rate: int, done: bool) -> None: + 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() + self._send_task = self.create_task(self._send_task_handler()) + + async def _cancel_send_task(self): + if self._send_task: + await self.cancel_task(self._send_task) + self._send_task = None + + 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 + # second of a 24000 sample rate would be 48000 bytes (16-bit samples and + # 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 + + audio_buffer = bytearray() + 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) + audio_buffer.clear() + else: + audio = await self._resampler.resample(audio, in_rate, self._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""" - if not done: - audio = await self._resampler.resample(audio, in_rate, self._sample_rate) 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) - async def process_frame(self, frame: Frame, direction: FrameDirection): - await super().process_frame(frame, direction) - if 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._encode_audio_and_send(frame.audio, frame.sample_rate, done=False) - elif isinstance(frame, TTSStoppedFrame): - await self._encode_audio_and_send(b"\x00", self._sample_rate, done=True) - await self.stop_ttfb_metrics() - await self.stop_processing_metrics() - elif isinstance(frame, StartInterruptionFrame): - await self._send_interrupt_message() - else: - await self.push_frame(frame, direction) - async def _send_interrupt_message(self) -> None: transport_frame = TransportMessageUrgentFrame( message={ @@ -127,7 +187,7 @@ class TavusVideoService(AIService): ) await self.push_frame(transport_frame) - async def _send_audio_message(self, audio_base64: str, done: bool) -> None: + async def _send_audio_message(self, audio_base64: str, done: bool): transport_frame = TransportMessageUrgentFrame( message={ "message_type": "conversation",