Use a PyAV resampler per participant for input audio

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):

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)

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)."""