AudioBufferProcessor: use on_audio_data event handler to retrieve audio

This commit is contained in:
Aleix Conchillo Flaqué
2024-12-03 14:10:50 -08:00
parent a6a4910931
commit 322dd0cea1
6 changed files with 97 additions and 65 deletions

View File

@@ -102,7 +102,6 @@ async def main():
audio_buffer_processor=audio_buffer_processor,
aiohttp_session=session,
api_key=os.getenv("CANONICAL_API_KEY"),
api_url=os.getenv("CANONICAL_API_URL"),
call_id=str(uuid.uuid4()),
assistant="pipecat-chatbot",
assistant_speaks_first=True,

View File

@@ -4,7 +4,9 @@
# SPDX-License-Identifier: BSD 2-Clause License
#
import aiofiles
import asyncio
import io
import os
import sys
@@ -32,15 +34,17 @@ logger.remove(0)
logger.add(sys.stderr, level="DEBUG")
async def save_audio(audiobuffer):
if audiobuffer.has_audio():
merged_audio = audiobuffer.merge_audio_buffers()
async def save_audio(audio: bytes, sample_rate: int, num_channels: int):
if len(audio) > 0:
filename = f"conversation_recording{datetime.datetime.now().strftime('%Y%m%d_%H%M%S')}.wav"
with wave.open(filename, "wb") as wf:
wf.setnchannels(2)
wf.setsampwidth(2)
wf.setframerate(audiobuffer._sample_rate)
wf.writeframes(merged_audio)
with io.BytesIO() as buffer:
with wave.open(buffer, "wb") as wf:
wf.setsampwidth(2)
wf.setnchannels(num_channels)
wf.setframerate(sample_rate)
wf.writeframes(audio)
async with aiofiles.open(filename, "wb") as file:
await file.write(buffer.getvalue())
print(f"Merged audio saved to {filename}")
else:
print("No audio data to save")
@@ -106,7 +110,9 @@ async def main():
context = OpenAILLMContext(messages)
context_aggregator = llm.create_context_aggregator(context)
audiobuffer = AudioBufferProcessor()
# Save audio every 10 seconds.
audiobuffer = AudioBufferProcessor(buffer_size=480000)
pipeline = Pipeline(
[
transport.input(), # microphone
@@ -121,6 +127,10 @@ async def main():
task = PipelineTask(pipeline, PipelineParams(allow_interruptions=True))
@audiobuffer.event_handler("on_audio_data")
async def on_audio_data(buffer, audio, sample_rate, num_channels):
await save_audio(audio, sample_rate, num_channels)
@transport.event_handler("on_first_participant_joined")
async def on_first_participant_joined(transport, participant):
await transport.capture_participant_transcription(participant["id"])
@@ -130,7 +140,6 @@ async def main():
async def on_participant_left(transport, participant, reason):
print(f"Participant left: {participant}")
await task.queue_frame(EndFrame())
await save_audio(audiobuffer)
runner = PipelineRunner()

View File

@@ -1,3 +1,4 @@
aiofiles
python-dotenv
fastapi[all]
uvicorn

View File

@@ -4,9 +4,6 @@
# SPDX-License-Identifier: BSD 2-Clause License
#
import wave
from io import BytesIO
from pipecat.audio.utils import interleave_stereo_audio, mix_audio, resample_audio
from pipecat.frames.frames import (
Frame,
@@ -17,43 +14,59 @@ from pipecat.processors.frame_processor import FrameDirection, FrameProcessor
class AudioBufferProcessor(FrameProcessor):
"""This processor buffers audio raw frames (input and output) that can later
be obtained as an in-memory WAV. You can provide the desired output
`sample_rate` and incoming audio frames will resampled to match it. Also,
you can provide the number of channels, 1 for mono and 2 for stereo. With
mono audio user and bot audio will be mixed, in the case of stereo the left
channel will be used for the user's audio and the right channel for the bot.
"""This processor buffers audio raw frames (input and output). The mixed
audio can be obtained by calling `get_audio()` (if `buffer_size` is 0) or by
registering an "on_audio_data" event handler. The event handler will be
called every time `buffer_size` is reached.
You can provide the desired output `sample_rate` and incoming audio frames
will resampled to match it. Also, you can provide the number of channels, 1
for mono and 2 for stereo. With mono audio user and bot audio will be mixed,
in the case of stereo the left channel will be used for the user's audio and
the right channel for the bot.
"""
def __init__(self, *, sample_rate: int = 24000, num_channels: int = 1, **kwargs):
def __init__(
self, *, sample_rate: int = 24000, num_channels: int = 1, buffer_size: int = 0, **kwargs
):
super().__init__(**kwargs)
self._sample_rate = sample_rate
self._num_channels = num_channels
self._buffer_size = buffer_size
self._user_audio_buffer = bytearray()
self._bot_audio_buffer = bytearray()
def _buffer_has_audio(self, buffer: bytearray) -> bool:
return buffer is not None and len(buffer) > 0
self._register_event_handler("on_audio_data")
@property
def sample_rate(self) -> int:
return self._sample_rate
@property
def num_channels(self) -> int:
return self._num_channels
def has_audio(self) -> bool:
return self._buffer_has_audio(self._user_audio_buffer) and self._buffer_has_audio(
self._bot_audio_buffer
)
def reset_audio_buffer(self):
self._user_audio_buffer = bytearray()
self._bot_audio_buffer = bytearray()
def merge_audio_buffers(self) -> bytes:
if self._num_channels == 1:
return self._merge_mono()
return mix_audio(bytes(self._user_audio_buffer), bytes(self._bot_audio_buffer))
elif self._num_channels == 2:
return self._merge_stereo()
return interleave_stereo_audio(
bytes(self._user_audio_buffer), bytes(self._bot_audio_buffer)
)
else:
return b""
def reset_audio_buffers(self):
self._user_audio_buffer = bytearray()
self._bot_audio_buffer = bytearray()
async def process_frame(self, frame: Frame, direction: FrameDirection):
await super().process_frame(frame, direction)
@@ -65,30 +78,25 @@ class AudioBufferProcessor(FrameProcessor):
if len(self._user_audio_buffer) > len(self._bot_audio_buffer):
silence = b"\x00" * len(resampled)
self._bot_audio_buffer.extend(silence)
# If the bot is speaking, include all audio from the bot.
if isinstance(frame, OutputAudioRawFrame):
elif isinstance(frame, OutputAudioRawFrame):
resampled = resample_audio(frame.audio, frame.sample_rate, self._sample_rate)
self._bot_audio_buffer.extend(resampled)
def _merge_mono(self) -> bytes:
with BytesIO() as buffer:
with wave.open(buffer, "wb") as wf:
wf.setnchannels(1)
wf.setsampwidth(2)
wf.setframerate(self._sample_rate)
mixed = mix_audio(bytes(self._user_audio_buffer), bytes(self._bot_audio_buffer))
wf.writeframes(mixed)
return buffer.getvalue()
if self._buffer_size > 0 and len(self._user_audio_buffer) > self._buffer_size:
await self._call_on_audio_data_handler()
def _merge_stereo(self) -> bytes:
with BytesIO() as buffer:
with wave.open(buffer, "wb") as wf:
wf.setnchannels(2)
wf.setsampwidth(2)
wf.setframerate(self._sample_rate)
stereo = interleave_stereo_audio(
bytes(self._user_audio_buffer), bytes(self._bot_audio_buffer)
)
wf.writeframes(stereo)
return buffer.getvalue()
await self.push_frame(frame, direction)
async def _call_on_audio_data_handler(self):
if not self.has_audio():
return
merged_audio = self.merge_audio_buffers()
await self._call_event_handler(
"on_audio_data", merged_audio, self._sample_rate, self._num_channels
)
self.reset_audio_buffers()
def _buffer_has_audio(self, buffer: bytearray) -> bool:
return buffer is not None and len(buffer) > 0

View File

@@ -676,6 +676,7 @@ class RTVIProcessor(FrameProcessor):
await self.push_frame(frame, direction)
async def cleanup(self):
await super().cleanup()
if self._pipeline:
await self._pipeline.cleanup()

View File

@@ -5,13 +5,16 @@
#
import aiohttp
import io
import os
import uuid
import wave
from datetime import datetime
from typing import Dict, List, Tuple
from pipecat.frames.frames import CancelFrame, EndFrame, Frame
from pipecat.processors.audio import audio_buffer_processor
from pipecat.processors.audio.audio_buffer_processor import AudioBufferProcessor
from pipecat.processors.frame_processor import FrameDirection
from pipecat.services.ai_services import AIService
@@ -81,9 +84,11 @@ class CanonicalMetricsService(AIService):
self._output_dir = output_dir
async def stop(self, frame: EndFrame):
await super().stop(frame)
await self._process_audio()
async def cancel(self, frame: CancelFrame):
await super().cancel(frame)
await self._process_audio()
async def process_frame(self, frame: Frame, direction: FrameDirection):
@@ -91,23 +96,32 @@ class CanonicalMetricsService(AIService):
await self.push_frame(frame, direction)
async def _process_audio(self):
pipeline = self._audio_buffer_processor
if pipeline.has_audio():
os.makedirs(self._output_dir, exist_ok=True)
filename = self._get_output_filename()
wave_data = pipeline.merge_audio_buffers()
audio_buffer_processor = self._audio_buffer_processor
if not audio_buffer_processor.has_audio():
return
os.makedirs(self._output_dir, exist_ok=True)
filename = self._get_output_filename()
audio = audio_buffer_processor.merge_audio_buffers()
with io.BytesIO() as buffer:
with wave.open(buffer, "wb") as wf:
wf.setsampwidth(2)
wf.setnchannels(audio_buffer_processor.num_channels)
wf.setframerate(audio_buffer_processor.sample_rate)
wf.writeframes(audio)
async with aiofiles.open(filename, "wb") as file:
await file.write(wave_data)
await file.write(buffer.getvalue())
try:
await self._multipart_upload(filename)
pipeline.reset_audio_buffer()
await aiofiles.os.remove(filename)
except FileNotFoundError:
pass
except Exception as e:
logger.error(f"Failed to upload recording: {e}")
try:
await self._multipart_upload(filename)
await aiofiles.os.remove(filename)
audio_buffer_processor.reset_audio_buffers()
except FileNotFoundError:
pass
except Exception as e:
logger.error(f"Failed to upload recording: {e}")
def _get_output_filename(self):
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")