#
# Copyright (c) 2024-2026, Daily
#
# SPDX-License-Identifier: BSD 2-Clause License
#
"""The harness's WebSocket transport and recorder.
A live transport carries continuous audio both ways, speech or silence,
and VADs, turn detection, and streaming STT rely on that cadence. The eval
bot sends and receives audio only while a TTS produces it, so both edges
here reshape it: the output paces the user's audio to the bot one 40 ms
frame per tick, silence in between, and the input re-emits the bot's audio
at the same cadence, filling its gaps, so the harness's VAD finds the
bot's real turn ends.
The recording does not use the paced streams, whose jitter would make it
stutter: :class:`EvalClientRecorder` is fed the raw audio on both edges and
lays each side out on its own timeline.
"""
import asyncio
import time
import wave
from pathlib import Path
from pipecat.audio.utils import create_stream_resampler, mix_audio
from pipecat.frames.frames import (
CancelFrame,
EndFrame,
Frame,
InputAudioRawFrame,
InputTransportMessageFrame,
InterruptionFrame,
OutputAudioRawFrame,
StartFrame,
)
from pipecat.processors.frame_processor import FrameDirection
from pipecat.transports.websocket.client import (
WebsocketClientInputTransport,
WebsocketClientOutputTransport,
)
from pipecat.transports.websocket.rtvi_client import RTVIClientTransport
# One audio frame per tick, on both edges. 40ms matches the bot's output chunking
# (``audio_out_10ms_chunks`` defaults to 4), so the receive side re-emits the bot's
# chunks 1:1 rather than splitting them, and it's still fine-grained enough for the
# VAD/turn models (well under their start/stop windows).
FRAME_S = 0.04
async def _sleep_to_next_tick(next_send: float) -> float:
"""Sleep to the next tick and return the one after.
When the loop has fallen behind, the next tick is measured from now rather
than firing a burst of catch-up frames, which would compress the audio.
"""
sleep_s = max(0.0, next_send - time.monotonic())
await asyncio.sleep(sleep_s)
return time.monotonic() + FRAME_S if sleep_s == 0 else next_send + FRAME_S
[docs]
class EvalClientRecorder:
"""Builds the conversation recording from the raw audio, not the paced streams.
Python cannot hold the 40 ms pacing tick precisely, and a recording of the
paced streams stutters. So each side is recorded as it was produced or
received, laid out on its own timeline with silence only where a real
pause was, and mixed to mono at :meth:`write`. Audio the bot sent past an
interruption is dropped, as a real client would drop it.
"""
# A chunk arriving later than its side's playout position by more than this
# is a real pause between turns; anything shorter is jitter within a
# contiguous turn and stays gapless.
GAP_S = 0.2
[docs]
def __init__(self, sample_rate: int):
"""Initialize the recorder.
Args:
sample_rate: Output sample rate; both sides are resampled to it on write.
"""
self._rate = sample_rate
self._user = _RecorderTrack()
self._bot = _RecorderTrack()
[docs]
def add_user(self, audio: bytes, in_rate: int) -> None:
"""Record a chunk of the user's TTS audio (raw, before pacing to the bot)."""
self._user.add(audio, in_rate)
[docs]
def add_bot(self, audio: bytes, in_rate: int) -> None:
"""Record a chunk of the bot's audio (raw, before the gap-fill loop)."""
self._bot.add(audio, in_rate)
[docs]
def drop_bot_tail(self, num_bytes: int) -> None:
"""Drop the last ``num_bytes`` of the bot's audio: sent, but never played."""
self._bot.drop_tail(num_bytes)
[docs]
def drop_user_tail(self, num_bytes: int) -> None:
"""Drop the last ``num_bytes`` of the user's audio: synthesized, but never sent."""
self._user.drop_tail(num_bytes)
[docs]
def has_audio(self) -> bool:
"""Whether any audio has been recorded on either side."""
return self._user.first is not None or self._bot.first is not None
[docs]
async def write(self, path: str) -> bool:
"""Resample both sides to a common rate, align, mix to mono, and write a WAV.
Returns:
True if a file was written, False if nothing was recorded.
"""
firsts = [t.first for t in (self._user, self._bot) if t.first is not None]
if not firsts:
return False
start = min(firsts)
user = await self._user.rendered(self._rate, start)
bot = await self._bot.rendered(self._rate, start)
mixed = mix_audio(user, bot)
out = Path(path)
out.parent.mkdir(parents=True, exist_ok=True)
with wave.open(str(out), "wb") as wf:
wf.setnchannels(1)
wf.setsampwidth(2)
wf.setframerate(self._rate)
wf.writeframes(mixed)
return True
class _RecorderTrack:
"""One side of the recording: raw audio chunks and when each arrived.
Chunks are laid out back to back. Silence is inserted only when a chunk
arrives later than the playout position by more than the gap threshold,
since within a turn both sources run ahead of real time.
"""
def __init__(self):
self._chunks: list[tuple[float, bytes]] = []
self._rate: int | None = None
@property
def first(self) -> float | None:
"""When the first chunk arrived, or ``None`` before any audio."""
return self._chunks[0][0] if self._chunks else None
def add(self, audio: bytes, in_rate: int) -> None:
"""Record a raw audio chunk as of now."""
if not audio:
return
self._rate = in_rate
self._chunks.append((time.monotonic(), audio))
def drop_tail(self, num_bytes: int) -> None:
"""Remove the last ``num_bytes`` of audio, across as many chunks as it spans."""
while num_bytes > 0 and self._chunks:
arrived, chunk = self._chunks[-1]
if len(chunk) <= num_bytes:
self._chunks.pop()
num_bytes -= len(chunk)
else:
self._chunks[-1] = (arrived, chunk[: len(chunk) - num_bytes])
num_bytes = 0
async def rendered(self, out_rate: int, start: float) -> bytes:
"""Lay the track out, resample to ``out_rate``, and prepend its silence since ``start``."""
if self.first is None or self._rate is None:
return b""
audio = self._layout()
if self._rate != out_rate:
audio = await create_stream_resampler().resample(audio, self._rate, out_rate)
lead = int((self.first - start) * out_rate * 2) & ~1
return b"\x00" * lead + audio
def _layout(self) -> bytes:
"""Concatenate the chunks, with silence for each real pause."""
assert self._rate is not None
out = bytearray()
position = self._chunks[0][0] # playout position, as a wall-clock time
for arrived, chunk in self._chunks:
late = arrived - position
if late > EvalClientRecorder.GAP_S:
out.extend(b"\x00" * (int(late * self._rate * 2) & ~1))
position = arrived
elif late > 0:
position = arrived
out.extend(chunk)
position += len(chunk) / (self._rate * 2)
return bytes(out)
[docs]
class EvalClientOutputTransport(WebsocketClientOutputTransport):
"""Streams the user's audio to the bot as a continuous real-time stream.
A real-time task sends one 40 ms frame per tick: queued TTS audio when
there is some, silence otherwise, so the bot's VAD and turn detection see
the silence they need to end a turn. Runs only when audio output is on.
"""
[docs]
def __init__(self, *args, recorder: "EvalClientRecorder | None" = None, **kwargs):
"""Initialize the transport and its (lazily started) send stream."""
super().__init__(*args, **kwargs)
self._pending = bytearray()
self._send_task = None
self._recorder = recorder
[docs]
async def start(self, frame: StartFrame):
"""Start the transport and, in audio mode, the real-time send stream."""
await super().start(frame)
if self._params.audio_out_enabled and self._send_task is None:
self._send_task = self.create_task(self._send_task_handler())
[docs]
async def stop(self, frame: EndFrame):
"""Stop the send stream, then the transport."""
await self._cancel_send_task()
await super().stop(frame)
[docs]
async def cancel(self, frame: CancelFrame):
"""Cancel the send stream, then the transport."""
await self._cancel_send_task()
await super().cancel(frame)
[docs]
async def process_frame(self, frame: Frame, direction: FrameDirection):
"""Pass a frame on; an interruption drops the user audio not yet sent, and the recorder drops the same bytes."""
if isinstance(frame, InterruptionFrame) and self._pending:
if self._recorder is not None:
self._recorder.drop_user_tail(len(self._pending))
self._pending.clear()
await super().process_frame(frame, direction)
[docs]
async def write_audio_frame(self, frame: OutputAudioRawFrame) -> bool:
"""Queue the user audio for the send task.
Returns False so the media sender does not push this un-paced frame
downstream; the send task pushes the paced frames instead.
"""
if self._session.is_closing or not self._session.is_connected:
return False
# Record the raw TTS audio here (gapless), not the paced frames the send
# task emits: pacing underruns/jitter would stutter the recording.
if self._recorder is not None:
self._recorder.add_user(frame.audio, frame.sample_rate)
self._pending.extend(frame.audio)
return False
async def _send_task_handler(self):
"""Send one ~40ms frame every tick: queued audio, or silence when idle."""
chunk = int(self.sample_rate * FRAME_S) * 2 # 16-bit mono
silence = b"\x00" * chunk
next_send = time.monotonic()
while True:
if len(self._pending) >= chunk:
pcm = bytes(self._pending[:chunk])
del self._pending[:chunk]
elif self._pending:
# Pad the utterance's final partial chunk to a full frame.
pcm = bytes(self._pending) + silence[len(self._pending) :]
self._pending.clear()
else:
pcm = silence
frame = OutputAudioRawFrame(
audio=pcm,
sample_rate=self.sample_rate,
num_channels=self._params.audio_out_channels,
)
await self._send_frame(frame)
# Push every frame (audio and silence) downstream at this paced cadence:
# the harness recorder aligns tracks by wall-clock, so a continuous
# stream keeps the user turn at the right time (pushing only audio would
# leave the track idle and the recorder would misplace it).
await self.push_frame(frame)
next_send = await _sleep_to_next_tick(next_send)
async def _send_frame(self, frame: OutputAudioRawFrame):
"""Serialize and send one frame (raw PCM, via the RTVI serializer)."""
if self._session.is_closing or not self._session.is_connected:
return
await self._write_frame(frame)
async def _cancel_send_task(self):
if self._send_task is not None:
await self.cancel_task(self._send_task)
self._send_task = None
[docs]
class EvalClientTransport(RTVIClientTransport):
"""The harness's RTVI client transport, with audio edges that behave like a live transport.
The input fills the gaps in the bot's audio and the output paces the
user's; both feed the recorder the raw audio when one is given.
"""
[docs]
def __init__(self, *args, recorder: "EvalClientRecorder | None" = None, **kwargs):
"""Initialize the transport, optionally wiring a recorder to both edges.
Args:
recorder: Optional :class:`EvalClientRecorder` fed the raw audio on both
edges; ``None`` disables recording.
*args: Forwarded to :class:`~pipecat.transports.websocket.rtvi_client.RTVIClientTransport`.
**kwargs: Forwarded to the parent transport.
"""
super().__init__(*args, **kwargs)
self._recorder = recorder
[docs]
def output(self) -> WebsocketClientOutputTransport:
"""Return the stream-shaping output transport."""
if not self._output:
self._output = EvalClientOutputTransport(
self, self._session, self._params, recorder=self._recorder
)
return self._output