diff --git a/src/pipecat_getstream/transport.py b/src/pipecat_getstream/transport.py index f8ead86..aaae2a4 100644 --- a/src/pipecat_getstream/transport.py +++ b/src/pipecat_getstream/transport.py @@ -23,9 +23,7 @@ from getstream.video.rtc.pb.stream.video.sfu.models.models_pb2 import TrackType from getstream.video.rtc.tracks import SubscriptionConfig, TrackSubscriptionConfig from loguru import logger -from pipecat.audio.utils import create_stream_resampler from pipecat.frames.frames import ( - AudioRawFrame, CancelFrame, EndFrame, EndTaskFrame, @@ -895,7 +893,11 @@ def __init__( self._audio_in_task: asyncio.Task | None = None self._video_in_task: asyncio.Task | None = None - self._resampler = create_stream_resampler() + # One resampler per participant, so their audio history and pts do not mix. + self._resamplers: dict[str, av.AudioResampler] = {} + self._transport.add_event_handler( + "on_audio_track_unsubscribed", self._on_audio_track_unsubscribed + ) self._initialized = False @@ -970,22 +972,10 @@ async def _audio_in_task_handler(self): async for audio_data in audio_iterator: if audio_data: pcm_data, participant_id = audio_data - pipecat_audio_frame = await self._convert_stream_audio_to_pipecat( - pcm_data - ) - - if len(pipecat_audio_frame.audio) == 0: - continue - input_audio_frame = UserAudioRawFrame( - user_id=participant_id, - audio=pipecat_audio_frame.audio, - sample_rate=pipecat_audio_frame.sample_rate, - num_channels=pipecat_audio_frame.num_channels, - ) - pts_seconds = pcm_data.pts_seconds - if pts_seconds is not None: - input_audio_frame.pts = seconds_to_nanoseconds(pts_seconds) - await self.push_audio_frame(input_audio_frame) + for input_audio_frame in self._convert_stream_audio_to_pipecat( + pcm_data, participant_id + ): + await self.push_audio_frame(input_audio_frame) async def _video_in_task_handler(self): """Handle incoming video frames from participants.""" @@ -1005,39 +995,52 @@ async def _video_in_task_handler(self): ) await self.push_video_frame(input_video_frame) - async def _convert_stream_audio_to_pipecat( - self, pcm_data: PcmData - ) -> AudioRawFrame: - """Convert Stream Video PcmData to Pipecat AudioRawFrame. + def _convert_stream_audio_to_pipecat( + self, pcm_data: PcmData, participant_id: str + ) -> list[UserAudioRawFrame]: + """Convert Stream Video PcmData to Pipecat UserAudioRawFrames. - Handles int16/float32 conversion and resampling. + Handles s16 conversion and resampling. The resampler holds back about + 1 ms of audio, and each frame's pts is the time of its first sample. Args: pcm_data: The PcmData from Stream Video SDK. + participant_id: The participant who sent the audio. Returns: - Converted AudioRawFrame for the pipeline. + Converted frames for the pipeline, possibly none. """ - samples = pcm_data.samples - - # Convert float32 to int16 if needed - if samples.dtype == np.float32: - samples = (samples * 32767).astype(np.int16) - elif samples.dtype != np.int16: - samples = samples.astype(np.int16) - - raw_bytes = samples.tobytes() - - # Resample to transport input sample rate - audio_data = await self._resampler.resample( - raw_bytes, pcm_data.sample_rate, self.sample_rate - ) + resampler = self._resamplers.get(participant_id) + if resampler is None: + resampler = av.AudioResampler(format="s16", rate=self.sample_rate) + self._resamplers[participant_id] = resampler + + av_frame = pcm_data.to_av_frame() + if pcm_data.time_base is None: + av_frame.pts = None + else: + av_frame.time_base = Fraction(pcm_data.time_base).limit_denominator() + + frames = [] + for resampled in resampler.resample(av_frame): + frame = UserAudioRawFrame( + user_id=participant_id, + audio=resampled.to_ndarray().tobytes(), + sample_rate=resampled.sample_rate, + num_channels=len(resampled.layout.channels), + ) + if resampled.pts is not None and resampled.time_base is not None: + frame.pts = seconds_to_nanoseconds( + float(resampled.pts * resampled.time_base) + ) + frames.append(frame) + return frames - return AudioRawFrame( - audio=audio_data, - sample_rate=self.sample_rate, - num_channels=pcm_data.channels if hasattr(pcm_data, "channels") else 1, - ) + def _on_audio_track_unsubscribed( + self, _transport: BaseTransport, participant_id: str + ): + """Drop the participant's resampler, because a new audio track restarts its pts.""" + self._resamplers.pop(participant_id, None) class GetstreamOutputTransport(BaseOutputTransport): diff --git a/src/tests/conftest.py b/src/tests/conftest.py index ce6569f..46f42be 100644 --- a/src/tests/conftest.py +++ b/src/tests/conftest.py @@ -139,17 +139,18 @@ def _factory(user_id: str, session_id: str, track_type: int): @pytest.fixture() def make_pcm_data(make_participant): - """Factory that creates a 20ms silent 48kHz PcmData chunk from a participant.""" + """Factory that creates a 48kHz PcmData chunk from a participant (20ms of silence by default).""" def _factory( user_id: str, session_id: str = "session-1", pts: int | None = None, + samples: np.ndarray | None = None, ) -> PcmData: return PcmData( sample_rate=48000, format="s16", - samples=np.zeros(960, dtype=np.int16), + samples=np.zeros(960, dtype=np.int16) if samples is None else samples, pts=pts, time_base=1 / 48000, participant=make_participant(user_id, session_id), @@ -189,37 +190,52 @@ async def on_pipeline_started(_worker, _frame): @pytest.fixture() -async def input_transport(create_callbacks, run_pipeline): - """A started GetstreamInputTransport's client and a downstream frame queue. +def create_input_transport(run_pipeline): + """Factory that starts a GetstreamInputTransport and returns its client and a downstream frame queue. The client skips the SFU join, so audio can be fed in with `client._on_audio()` and read back from the queue as the transport pushes it downstream. + The client reports its events to the transport, as the real client does. """ - params = GetstreamParams(audio_in_enabled=True, audio_in_sample_rate=48000) - transport = GetstreamTransport( - api_key="test-key", - token="test-token", - call_type="default", - call_id="test-call", - user_id="bot-user", - params=params, - ) - client = _OfflineClient( - api_key="test-key", - token="test-token", - call_type="default", - call_id="test-call", - user_id="bot-user", - params=params, - callbacks=create_callbacks(), - transport_name="test-transport", - ) - input_t = GetstreamInputTransport(transport, client, params) - - received: asyncio.Queue = asyncio.Queue() - sink = QueuedFrameProcessor( - queue=received, queue_direction=FrameDirection.DOWNSTREAM - ) - await run_pipeline(input_t, sink) - - return client, received + + async def _factory( + sample_rate: int, + ) -> tuple[GetstreamTransportClient, asyncio.Queue]: + params = GetstreamParams( + audio_in_enabled=True, audio_in_sample_rate=sample_rate + ) + transport = GetstreamTransport( + api_key="test-key", + token="test-token", + call_type="default", + call_id="test-call", + user_id="bot-user", + params=params, + ) + client = _OfflineClient( + api_key="test-key", + token="test-token", + call_type="default", + call_id="test-call", + user_id="bot-user", + params=params, + callbacks=transport._client._callbacks, + transport_name="test-transport", + ) + input_t = GetstreamInputTransport(transport, client, params) + + received: asyncio.Queue = asyncio.Queue() + sink = QueuedFrameProcessor( + queue=received, queue_direction=FrameDirection.DOWNSTREAM + ) + await run_pipeline(input_t, sink) + + return client, received + + return _factory + + +@pytest.fixture() +async def input_transport(create_input_transport): + """A started 48kHz GetstreamInputTransport's client and a downstream frame queue.""" + return await create_input_transport(48000) diff --git a/src/tests/test_getstream_transport.py b/src/tests/test_getstream_transport.py index ff148ee..2b46eb1 100644 --- a/src/tests/test_getstream_transport.py +++ b/src/tests/test_getstream_transport.py @@ -7,6 +7,7 @@ """ import asyncio +import contextlib import os import time import uuid @@ -117,6 +118,120 @@ async def test_emitted_user_audio_frames_carry_pts( assert [frame.user_id for frame in frames] == ["user-A"] * 3 assert [frame.pts for frame in frames] == [0, 20_000_000, 40_000_000] + async def test_resampled_frames_carry_pts_of_their_first_sample( + self, create_input_transport, make_pcm_data + ): + """After resampling, each frame's pts is the time of its first sample.""" + client, received = await create_input_transport(16000) + # Each input sample holds its own 48kHz index, so its value gives its time. + ramp = np.arange(9600, dtype=np.int16) + for start in range(0, 9600, 960): + client._on_audio( + make_pcm_data("user-A", pts=start, samples=ramp[start : start + 960]) + ) + + frames = [] + with contextlib.suppress(TimeoutError): + while True: + frames.append(await asyncio.wait_for(received.get(), timeout=0.5)) + + assert frames + for frame in frames: + first_sample = int(np.frombuffer(frame.audio, dtype=np.int16)[0]) + first_sample_ns = first_sample * 1_000_000_000 // 48000 + assert abs(frame.pts - first_sample_ns) <= 1_000_000 + + async def test_resampled_audio_keeps_the_end_of_speech_after_a_pause( + self, create_input_transport, make_pcm_data + ): + """The end of speech before a pause is delivered when the next audio arrives.""" + client, received = await create_input_transport(16000) + ramp = np.arange(9600, dtype=np.int16) + for start in range(0, 9600, 960): + client._on_audio( + make_pcm_data("user-A", pts=start, samples=ramp[start : start + 960]) + ) + await asyncio.sleep(0.5) + client._on_audio(make_pcm_data("user-A", pts=9600)) + + frames = [] + with contextlib.suppress(TimeoutError): + while True: + frames.append(await asyncio.wait_for(received.get(), timeout=0.5)) + + audio = np.concatenate( + [np.frombuffer(frame.audio, dtype=np.int16) for frame in frames] + ) + # The last speech sample is 9599; allow 1 ms (48 input samples) for the filter edge. + assert audio.max() >= 9599 - 48 + + async def test_participants_are_resampled_separately( + self, create_input_transport, make_pcm_data + ): + """Resampled audio of one participant contains no samples of another.""" + client, received = await create_input_transport(16000) + ramp = np.arange(9600, dtype=np.int16) + for start in range(0, 9600, 960): + client._on_audio( + make_pcm_data("user-A", pts=start, samples=ramp[start : start + 960]) + ) + client._on_audio( + make_pcm_data( + "user-B", pts=start, samples=np.full(960, -10000, dtype=np.int16) + ) + ) + + frames = [] + with contextlib.suppress(TimeoutError): + while True: + frames.append(await asyncio.wait_for(received.get(), timeout=0.5)) + + audio_a = np.concatenate( + [ + np.frombuffer(frame.audio, dtype=np.int16) + for frame in frames + if frame.user_id == "user-A" + ] + ) + assert audio_a.min() > -1000 + + async def test_audio_after_participant_rejoins_has_no_old_audio( + self, create_input_transport, make_pcm_data, make_participant_event + ): + """After a participant leaves, their next audio has correct pts and no audio from before.""" + client, received = await create_input_transport(16000) + ramp = np.arange(9600, dtype=np.int16) + for start in range(0, 9600, 960): + client._on_audio( + make_pcm_data( + "user-A", pts=48000 + start, samples=ramp[start : start + 960] + ) + ) + with contextlib.suppress(TimeoutError): + while True: + await asyncio.wait_for(received.get(), timeout=0.5) + + client._on_participant_left(make_participant_event("user-A")) + await asyncio.sleep(0.1) + # The new track of the rejoined participant starts its pts from 0 again. + for start in range(0, 9600, 960): + client._on_audio( + make_pcm_data( + "user-A", pts=start, samples=np.full(960, -10000, dtype=np.int16) + ) + ) + + frames = [] + with contextlib.suppress(TimeoutError): + while True: + frames.append(await asyncio.wait_for(received.get(), timeout=0.5)) + + audio = np.concatenate( + [np.frombuffer(frame.audio, dtype=np.int16) for frame in frames] + ) + assert frames[0].pts == 0 + assert audio.max() < 1000 + class TestGetstreamTransport: """Constructor validation for credential modes (api_secret vs token)."""