AudioBufferProcessor: use on_audio_data event handler to retrieve audio
This commit is contained in:
@@ -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,
|
||||
|
||||
@@ -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()
|
||||
|
||||
|
||||
@@ -1,3 +1,4 @@
|
||||
aiofiles
|
||||
python-dotenv
|
||||
fastapi[all]
|
||||
uvicorn
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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()
|
||||
|
||||
|
||||
@@ -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")
|
||||
|
||||
Reference in New Issue
Block a user