dv-pipecat-ai 0.0.74.dev770__py3-none-any.whl → 0.0.82.dev776__py3-none-any.whl

pipecat/transports/network/webrtc_connection.py

@@ -4,6 +4,13 @@
 # SPDX-License-Identifier: BSD 2-Clause License
 #
 
+"""Small WebRTC connection implementation for Pipecat.
+
+This module provides a WebRTC connection implementation using aiortc,
+with support for audio/video tracks, data channels, and signaling
+for real-time communication applications.
+"""
+
 import asyncio
 import json
 import time
@@ -32,39 +39,90 @@ except ModuleNotFoundError as e:
 SIGNALLING_TYPE = "signalling"
 AUDIO_TRANSCEIVER_INDEX = 0
 VIDEO_TRANSCEIVER_INDEX = 1
+SCREEN_VIDEO_TRANSCEIVER_INDEX = 2
 
 
 class TrackStatusMessage(BaseModel):
+    """Message for updating track enabled/disabled status.
+
+    Parameters:
+        type: Message type identifier.
+        receiver_index: Index of the track receiver to update.
+        enabled: Whether the track should be enabled or disabled.
+    """
+
     type: Literal["trackStatus"]
     receiver_index: int
     enabled: bool
 
 
 class RenegotiateMessage(BaseModel):
+    """Message requesting WebRTC renegotiation.
+
+    Parameters:
+        type: Message type identifier for renegotiation requests.
+    """
+
     type: Literal["renegotiate"] = "renegotiate"
 
 
 class PeerLeftMessage(BaseModel):
+    """Message indicating a peer has left the connection.
+
+    Parameters:
+        type: Message type identifier for peer departure.
+    """
+
     type: Literal["peerLeft"] = "peerLeft"
 
 
 class SignallingMessage:
+    """Union types for signaling message handling.
+
+    Parameters:
+        Inbound: Types of messages that can be received from peers.
+        outbound: Types of messages that can be sent to peers.
+    """
+
     Inbound = Union[TrackStatusMessage]  # in case we need to add new messages in the future
     outbound = Union[RenegotiateMessage]
 
 
 class SmallWebRTCTrack:
+    """Wrapper for WebRTC media tracks with enabled/disabled state management.
+
+    Provides additional functionality on top of aiortc MediaStreamTrack including
+    enable/disable control and frame discarding for audio and video streams.
+    """
+
     def __init__(self, track: MediaStreamTrack):
+        """Initialize the WebRTC track wrapper.
+
+        Args:
+            track: The underlying MediaStreamTrack to wrap.
+            index: The index of the track in the transceiver (0 for mic, 1 for cam, 2 for screen)
+        """
         self._track = track
         self._enabled = True
 
     def set_enabled(self, enabled: bool) -> None:
+        """Enable or disable the track.
+
+        Args:
+            enabled: Whether the track should be enabled for receiving frames.
+        """
         self._enabled = enabled
 
     def is_enabled(self) -> bool:
+        """Check if the track is currently enabled.
+
+        Returns:
+            True if the track is enabled for receiving frames.
+        """
         return self._enabled
 
     async def discard_old_frames(self):
+        """Discard old frames from the track queue to reduce latency."""
         remote_track = self._track
         if isinstance(remote_track, RemoteStreamTrack):
             if not hasattr(remote_track, "_queue") or not isinstance(
@@ -78,11 +136,24 @@ class SmallWebRTCTrack:
                 remote_track._queue.task_done()
 
     async def recv(self) -> Optional[Frame]:
-        if not self._enabled:
+        """Receive the next frame from the track.
+
+        Returns:
+            The next frame, except for video tracks, where it returns the frame only if the track is enabled, otherwise, returns None.
+        """
+        if not self._enabled and self._track.kind == "video":
             return None
         return await self._track.recv()
 
     def __getattr__(self, name):
+        """Forward attribute access to the underlying track.
+
+        Args:
+            name: The attribute name to access.
+
+        Returns:
+            The attribute value from the underlying track.
+        """
         # Forward other attribute/method calls to the underlying track
         return getattr(self._track, name)
 
@@ -92,7 +163,22 @@ IceServer = RTCIceServer
 
 
 class SmallWebRTCConnection(BaseObject):
+    """WebRTC connection implementation using aiortc.
+
+    Provides WebRTC peer connection functionality including ICE server configuration,
+    track management, data channel communication, and connection state handling
+    for real-time audio/video communication.
+    """
+
     def __init__(self, ice_servers: Optional[Union[List[str], List[IceServer]]] = None):
+        """Initialize the WebRTC connection.
+
+        Args:
+            ice_servers: List of ICE servers as URLs or IceServer objects.
+
+        Raises:
+            TypeError: If ice_servers contains mixed types or unsupported types.
+        """
         super().__init__()
         if not ice_servers:
             self.ice_servers: List[IceServer] = []
@@ -107,6 +193,7 @@ class SmallWebRTCConnection(BaseObject):
         self._track_getters = {
             AUDIO_TRANSCEIVER_INDEX: self.audio_input_track,
             VIDEO_TRANSCEIVER_INDEX: self.video_input_track,
+            SCREEN_VIDEO_TRANSCEIVER_INDEX: self.screen_video_input_track,
         }
 
         self._initialize()
@@ -126,13 +213,24 @@ class SmallWebRTCConnection(BaseObject):
 
     @property
     def pc(self) -> RTCPeerConnection:
+        """Get the underlying RTCPeerConnection.
+
+        Returns:
+            The aiortc RTCPeerConnection instance.
+        """
         return self._pc
 
     @property
     def pc_id(self) -> str:
+        """Get the peer connection identifier.
+
+        Returns:
+            The unique identifier for this peer connection.
+        """
         return self._pc_id
 
     def _initialize(self):
+        """Initialize the peer connection and associated components."""
         logger.debug("Initializing new peer connection")
         rtc_config = RTCConfiguration(iceServers=self.ice_servers)
 
@@ -147,6 +245,8 @@ class SmallWebRTCConnection(BaseObject):
         self._pending_app_messages = []
 
     def _setup_listeners(self):
+        """Set up event listeners for the peer connection."""
+
         @self._pc.on("datachannel")
         def on_datachannel(channel):
             self._data_channel = channel
@@ -208,6 +308,7 @@ class SmallWebRTCConnection(BaseObject):
                 await self._call_event_handler("track-ended", track)
 
     async def _create_answer(self, sdp: str, type: str):
+        """Create an SDP answer for the given offer."""
         offer = RTCSessionDescription(sdp=sdp, type=type)
         await self._pc.setRemoteDescription(offer)
 
@@ -223,9 +324,16 @@ class SmallWebRTCConnection(BaseObject):
         self._answer = self._pc.localDescription
 
     async def initialize(self, sdp: str, type: str):
+        """Initialize the connection with an SDP offer.
+
+        Args:
+            sdp: The SDP offer string.
+            type: The SDP type (usually "offer").
+        """
         await self._create_answer(sdp, type)
 
     async def connect(self):
+        """Connect the WebRTC peer connection and handle initial setup."""
         self._connect_invoked = True
         # If we already connected, trigger again the connected event
         if self.is_connected():
@@ -238,9 +346,23 @@ class SmallWebRTCConnection(BaseObject):
             video_input_track = self.video_input_track()
             if video_input_track:
                 await self.video_input_track().discard_old_frames()
-            self.ask_to_renegotiate()
+            screen_video_input_track = self.screen_video_input_track()
+            if screen_video_input_track:
+                await self.screen_video_input_track().discard_old_frames()
+            if video_input_track or screen_video_input_track:
+                # This prevents an issue where sometimes the WebRTC connection can be established
+                # before the bot is ready to receive video. When that happens, we can lose a couple
+                # of seconds of video before we received a key frame to finally start displaying it.
+                self.ask_to_renegotiate()
 
     async def renegotiate(self, sdp: str, type: str, restart_pc: bool = False):
+        """Renegotiate the WebRTC connection with new parameters.
+
+        Args:
+            sdp: The new SDP offer string.
+            type: The SDP type (usually "offer").
+            restart_pc: Whether to restart the peer connection entirely.
+        """
         logger.debug(f"Renegotiating {self._pc_id}")
 
         if restart_pc:
@@ -264,14 +386,24 @@ class SmallWebRTCConnection(BaseObject):
         asyncio.create_task(delayed_task())
 
     def force_transceivers_to_send_recv(self):
+        """Force all transceivers to bidirectional send/receive mode."""
         for transceiver in self._pc.getTransceivers():
-            transceiver.direction = "sendrecv"
+            # For now, we only support sendrecv for camera audio and video (the first two transceivers)
+            if transceiver.mid == "0" or transceiver.mid == "1":
+                transceiver.direction = "sendrecv"
+            else:
+                transceiver.direction = "recvonly"
             # logger.debug(
             #    f"Transceiver: {transceiver}, Mid: {transceiver.mid}, Direction: {transceiver.direction}"
             # )
             # logger.debug(f"Sender track: {transceiver.sender.track}")
 
     def replace_audio_track(self, track):
+        """Replace the audio track in the first transceiver.
+
+        Args:
+            track: The new audio track to use for sending.
+        """
         logger.debug(f"Replacing audio track {track.kind}")
         # Transceivers always appear in creation-order for both peers
         # For now we are only considering that we are going to have 02 transceivers,
@@ -283,6 +415,11 @@ class SmallWebRTCConnection(BaseObject):
             logger.warning("Audio transceiver not found. Cannot replace audio track.")
 
     def replace_video_track(self, track):
+        """Replace the video track in the second transceiver.
+
+        Args:
+            track: The new video track to use for sending.
+        """
         logger.debug(f"Replacing video track {track.kind}")
         # Transceivers always appear in creation-order for both peers
         # For now we are only considering that we are going to have 02 transceivers,
@@ -293,11 +430,29 @@ class SmallWebRTCConnection(BaseObject):
         else:
             logger.warning("Video transceiver not found. Cannot replace video track.")
 
+    def replace_screen_video_track(self, track):
+        """Replace the screen video track in the second transceiver.
+
+        Args:
+            track: The new screen video track to use for sending.
+        """
+        logger.debug(f"Replacing screen video track {track.kind}")
+        # Transceivers always appear in creation-order for both peers
+        # For now we are only considering that we are going to have 02 transceivers,
+        # one for audio and one for video
+        transceivers = self._pc.getTransceivers()
+        if len(transceivers) > 2 and transceivers[2].sender:
+            transceivers[2].sender.replaceTrack(track)
+        else:
+            logger.warning("Screen video transceiver not found. Cannot replace screen video track.")
+
     async def disconnect(self):
+        """Disconnect from the WebRTC peer connection."""
         self.send_app_message({"type": SIGNALLING_TYPE, "message": PeerLeftMessage().model_dump()})
         await self._close()
 
     async def _close(self):
+        """Close the peer connection and cleanup resources."""
         if self._pc:
             await self._pc.close()
         self._message_queue.clear()
@@ -305,6 +460,12 @@ class SmallWebRTCConnection(BaseObject):
         self._track_map = {}
 
     def get_answer(self):
+        """Get the SDP answer for the current connection.
+
+        Returns:
+            Dictionary containing SDP answer, type, and peer connection ID,
+            or None if no answer is available.
+        """
         if not self._answer:
             return None
 
@@ -315,6 +476,7 @@ class SmallWebRTCConnection(BaseObject):
         }
 
     async def _handle_new_connection_state(self):
+        """Handle changes in the peer connection state."""
         state = self._pc.connectionState
         if state == "connected" and not self._connect_invoked:
             # We are going to wait until the pipeline is ready before triggering the event
@@ -328,7 +490,12 @@ class SmallWebRTCConnection(BaseObject):
     # Despite the fact that aiortc provides this listener, they don't have a status for "disconnected"
     # So, there is no advantage in looking at self._pc.connectionState
     # That is why we are trying to keep our own state
-    def is_connected(self):
+    def is_connected(self) -> bool:
+        """Check if the WebRTC connection is currently active.
+
+        Returns:
+            True if the connection is active and receiving data.
+        """
         # If the small webrtc transport has never invoked to connect
         # we are acting like if we are not connected
         if not self._connect_invoked:
@@ -342,12 +509,17 @@ class SmallWebRTCConnection(BaseObject):
         return (time.time() - self._last_received_time) < 3
 
     def audio_input_track(self):
+        """Get the audio input track wrapper.
+
+        Returns:
+            SmallWebRTCTrack wrapper for the audio track, or None if unavailable.
+        """
         if self._track_map.get(AUDIO_TRANSCEIVER_INDEX):
             return self._track_map[AUDIO_TRANSCEIVER_INDEX]
 
         # Transceivers always appear in creation-order for both peers
-        # For now we are only considering that we are going to have 02 transceivers,
-        # one for audio and one for video
+        # For support 3 receivers in the following order:
+        #   audio, video, screenVideo
         transceivers = self._pc.getTransceivers()
         if len(transceivers) == 0 or not transceivers[AUDIO_TRANSCEIVER_INDEX].receiver:
             logger.warning("No audio transceiver is available")
@@ -359,12 +531,17 @@ class SmallWebRTCConnection(BaseObject):
         return audio_track
 
     def video_input_track(self):
+        """Get the video input track wrapper.
+
+        Returns:
+            SmallWebRTCTrack wrapper for the video track, or None if unavailable.
+        """
         if self._track_map.get(VIDEO_TRANSCEIVER_INDEX):
             return self._track_map[VIDEO_TRANSCEIVER_INDEX]
 
         # Transceivers always appear in creation-order for both peers
-        # For now we are only considering that we are going to have 02 transceivers,
-        # one for audio and one for video
+        # For support 3 receivers in the following order:
+        #   audio, video, screenVideo
         transceivers = self._pc.getTransceivers()
         if len(transceivers) <= 1 or not transceivers[VIDEO_TRANSCEIVER_INDEX].receiver:
             logger.warning("No video transceiver is available")
@@ -375,7 +552,34 @@ class SmallWebRTCConnection(BaseObject):
         self._track_map[VIDEO_TRANSCEIVER_INDEX] = video_track
         return video_track
 
+    def screen_video_input_track(self):
+        """Get the screen video input track wrapper.
+
+        Returns:
+            SmallWebRTCTrack wrapper for the screen video track, or None if unavailable.
+        """
+        if self._track_map.get(SCREEN_VIDEO_TRANSCEIVER_INDEX):
+            return self._track_map[SCREEN_VIDEO_TRANSCEIVER_INDEX]
+
+        # Transceivers always appear in creation-order for both peers
+        # For support 3 receivers in the following order:
+        #   audio, video, screenVideo
+        transceivers = self._pc.getTransceivers()
+        if len(transceivers) <= 2 or not transceivers[SCREEN_VIDEO_TRANSCEIVER_INDEX].receiver:
+            logger.warning("No screen video transceiver is available")
+            return None
+
+        track = transceivers[SCREEN_VIDEO_TRANSCEIVER_INDEX].receiver.track
+        video_track = SmallWebRTCTrack(track) if track else None
+        self._track_map[SCREEN_VIDEO_TRANSCEIVER_INDEX] = video_track
+        return video_track
+
     def send_app_message(self, message: Any):
+        """Send an application message through the data channel.
+
+        Args:
+            message: The message to send (will be JSON serialized).
+        """
         json_message = json.dumps(message)
         if self._data_channel and self._data_channel.readyState == "open":
             self._data_channel.send(json_message)
@@ -384,6 +588,7 @@ class SmallWebRTCConnection(BaseObject):
             self._message_queue.append(json_message)
 
     def ask_to_renegotiate(self):
+        """Request renegotiation of the WebRTC connection."""
         if self._renegotiation_in_progress:
             return
 
@@ -393,6 +598,7 @@ class SmallWebRTCConnection(BaseObject):
         )
 
     def _handle_signalling_message(self, message):
+        """Handle incoming signaling messages."""
         logger.debug(f"Signalling message received: {message}")
         inbound_adapter = TypeAdapter(SignallingMessage.Inbound)
         signalling_message = inbound_adapter.validate_python(message)