BaseInputTransport: handle StopFrame and pause pushing frames

CHANGELOG.md

@@ -9,6 +9,11 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Added
 
+- `BaseInputTransport` now handles `StopFrame`. When a `StopFrame` is received
+  the transport will pause sending frames downstream until a new `StartFrame` is
+  received. This allows the transport to be reused (keeping the same connection)
+  in a different pipeline.
+
 - Updated AssemblyAI STT service to support their latest streaming
   speech-to-text model with improved transcription latency and endpointing.
 

src/pipecat/transports/base_input.py

@@ -24,9 +24,11 @@ from pipecat.frames.frames import (
     FilterUpdateSettingsFrame,
     Frame,
     InputAudioRawFrame,
+    InputImageRawFrame,
     MetricsFrame,
     StartFrame,
     StartInterruptionFrame,
+    StopFrame,
     StopInterruptionFrame,
     SystemFrame,
     UserStartedSpeakingFrame,
@@ -57,6 +59,11 @@ class BaseInputTransport(FrameProcessor):
         # if passthrough is enabled.
         self._audio_task = None
 
+        # If the transport is stopped with `StopFrame` we might still be
+        # receiving frames from the transport but we really don't want to push
+        # them downstream until we get another `StartFrame`.
+        self._paused = False
+
         if self._params.vad_enabled:
             import warnings
 
@@ -117,6 +124,8 @@ class BaseInputTransport(FrameProcessor):
         return self._params.turn_analyzer
 
     async def start(self, frame: StartFrame):
+        self._paused = False
+
         self._sample_rate = self._params.audio_in_sample_rate or frame.audio_in_sample_rate
 
         # Configure VAD analyzer.
@@ -133,28 +142,33 @@ class BaseInputTransport(FrameProcessor):
 
     async def stop(self, frame: EndFrame):
         # Cancel and wait for the audio input task to finish.
-        if self._audio_task and self._params.audio_in_enabled:
+        await self._cancel_audio_task()
-            await self.cancel_task(self._audio_task)
-            self._audio_task = None
         # Stop audio filter.
         if self._params.audio_in_filter:
             await self._params.audio_in_filter.stop()
 
+    async def pause(self, frame: StopFrame):
+        self._paused = True
+        # Cancel task so we clear the queue
+        await self._cancel_audio_task()
+        # Retart the task
+        self._create_audio_task()
+
     async def cancel(self, frame: CancelFrame):
         # Cancel and wait for the audio input task to finish.
-        if self._audio_task and self._params.audio_in_enabled:
+        await self._cancel_audio_task()
-            await self.cancel_task(self._audio_task)
-            self._audio_task = None
 
     async def set_transport_ready(self, frame: StartFrame):
         """To be called when the transport is ready to stream."""
         # Create audio input queue and task if needed.
-        if not self._audio_task and self._params.audio_in_enabled:
+        self._create_audio_task()
-            self._audio_in_queue = asyncio.Queue()
+
-            self._audio_task = self.create_task(self._audio_task_handler())
+    async def push_video_frame(self, frame: InputImageRawFrame):
+        if self._params.video_in_enabled and not self._paused:
+            await self.push_frame(frame)
 
     async def push_audio_frame(self, frame: InputAudioRawFrame):
-        if self._params.audio_in_enabled:
+        if self._params.audio_in_enabled and not self._paused:
             await self._audio_in_queue.put(frame)
 
     #
@@ -190,6 +204,9 @@ class BaseInputTransport(FrameProcessor):
             # finish and the task finishes when EndFrame is processed.
             await self.push_frame(frame, direction)
             await self.stop(frame)
+        elif isinstance(frame, StopFrame):
+            await self.push_frame(frame, direction)
+            await self.pause(frame)
         elif isinstance(frame, VADParamsUpdateFrame):
             if self.vad_analyzer:
                 self.vad_analyzer.set_params(frame.params)
@@ -231,6 +248,16 @@ class BaseInputTransport(FrameProcessor):
     # Audio input
     #
 
+    def _create_audio_task(self):
+        if not self._audio_task and self._params.audio_in_enabled:
+            self._audio_in_queue = asyncio.Queue()
+            self._audio_task = self.create_task(self._audio_task_handler())
+
+    async def _cancel_audio_task(self):
+        if self._audio_task:
+            await self.cancel_task(self._audio_task)
+            self._audio_task = None
+
     async def _vad_analyze(self, audio_frame: InputAudioRawFrame) -> VADState:
         state = VADState.QUIET
         if self.vad_analyzer:

src/pipecat/transports/network/small_webrtc.py

@@ -19,7 +19,6 @@ from pipecat.frames.frames import (
     EndFrame,
     Frame,
     InputAudioRawFrame,
-    InputImageRawFrame,
     OutputAudioRawFrame,
     OutputImageRawFrame,
     SpriteFrame,
@@ -232,7 +231,8 @@ class SmallWebRTCClient:
             frame_array = frame.to_ndarray(format=format_name)
             frame_rgb = self._convert_frame(frame_array, format_name)
 
-            image_frame = InputImageRawFrame(
+            image_frame = UserImageRawFrame(
+                user_id=self._webrtc_connection.pc_id,
                 image=frame_rgb.tobytes(),
                 size=(frame.width, frame.height),
                 format="RGB",
@@ -424,7 +424,7 @@ class SmallWebRTCInputTransport(BaseInputTransport):
         try:
             async for video_frame in self._client.read_video_frame():
                 if video_frame:
-                    await self.push_frame(video_frame)
+                    await self.push_video_frame(video_frame)
 
                     # Check if there are any pending image requests and create UserImageRawFrame
                     if self._image_requests:
@@ -438,7 +438,7 @@ class SmallWebRTCInputTransport(BaseInputTransport):
                                 format=video_frame.format,
                             )
                             # Push the frame to the pipeline
-                            await self.push_frame(image_frame)
+                            await self.push_video_frame(image_frame)
                             # Remove from pending requests
                             del self._image_requests[req_id]
 

src/pipecat/transports/services/daily.py

@@ -1006,14 +1006,14 @@ class DailyInputTransport(BaseInputTransport):
         await self._transport.cleanup()
 
     async def start(self, frame: StartFrame):
+        # Parent start.
+        await super().start(frame)
+
         if self._initialized:
             return
 
         self._initialized = True
 
-        # Parent start.
-        await super().start(frame)
-
         # Setup client.
         await self._client.start(frame)
 
@@ -1148,7 +1148,7 @@ class DailyInputTransport(BaseInputTransport):
                 format=video_frame.color_format,
             )
             frame.transport_source = video_source
-            await self.push_frame(frame)
+            await self.push_video_frame(frame)
             self._video_renderers[participant_id][video_source]["timestamp"] = curr_time
 
 
@@ -1183,14 +1183,14 @@ class DailyOutputTransport(BaseOutputTransport):
         await self._transport.cleanup()
 
     async def start(self, frame: StartFrame):
+        # Parent start.
+        await super().start(frame)
+
         if self._initialized:
             return
 
         self._initialized = True
 
-        # Parent start.
-        await super().start(frame)
-
         # Setup client.
         await self._client.start(frame)