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

pipecat/transports/network/fastapi_websocket.py

@@ -4,6 +4,12 @@
 # SPDX-License-Identifier: BSD 2-Clause License
 #
 
+"""FastAPI WebSocket transport implementation for Pipecat.
+
+This module provides WebSocket-based transport for real-time audio/video streaming
+using FastAPI and WebSocket connections. Supports binary and text serialization
+with configurable session timeouts and WAV header generation.
+"""
 
 import asyncio
 import io
@@ -32,7 +38,6 @@ from pipecat.serializers.base_serializer import FrameSerializer, FrameSerializer
 from pipecat.transports.base_input import BaseInputTransport
 from pipecat.transports.base_output import BaseOutputTransport
 from pipecat.transports.base_transport import BaseTransport, TransportParams
-from pipecat.utils.asyncio.watchdog_async_iterator import WatchdogAsyncIterator
 
 try:
     from fastapi import WebSocket
@@ -46,24 +51,48 @@ except ModuleNotFoundError as e:
 
 
 class FastAPIWebsocketParams(TransportParams):
+    """Configuration parameters for FastAPI WebSocket transport.
+
+    Parameters:
+        add_wav_header: Whether to add WAV headers to audio frames.
+        serializer: Frame serializer for encoding/decoding messages.
+        session_timeout: Session timeout in seconds, None for no timeout.
+    """
+
     add_wav_header: bool = False
     serializer: Optional[FrameSerializer] = None
     session_timeout: Optional[int] = None
 
 
 class FastAPIWebsocketCallbacks(BaseModel):
+    """Callback functions for WebSocket events.
+
+    Parameters:
+        on_client_connected: Called when a client connects to the WebSocket.
+        on_client_disconnected: Called when a client disconnects from the WebSocket.
+        on_session_timeout: Called when a session timeout occurs.
+    """
+
     on_client_connected: Callable[[WebSocket], Awaitable[None]]
     on_client_disconnected: Callable[[WebSocket], Awaitable[None]]
     on_session_timeout: Callable[[WebSocket], Awaitable[None]]
 
 
 class FastAPIWebsocketClient:
-    def __init__(
-        self,
-        websocket: WebSocket,
-        is_binary: bool,
-        callbacks: FastAPIWebsocketCallbacks,
-    ):
+    """WebSocket client wrapper for handling connections and message passing.
+
+    Manages WebSocket state, message sending/receiving, and connection lifecycle
+    with support for both binary and text message types.
+    """
+
+    def __init__(self, websocket: WebSocket, is_binary: bool, callbacks: FastAPIWebsocketCallbacks):
+        """Initialize the WebSocket client.
+
+        Args:
+            websocket: The FastAPI WebSocket connection.
+            is_binary: Whether to use binary message format.
+            callbacks: Event callback functions.
+        """
         self._websocket = websocket
         self._closing = False
         self._is_binary = is_binary
@@ -72,14 +101,29 @@ class FastAPIWebsocketClient:
         self._conversation_id = None
 
     async def setup(self, _: StartFrame):
+        """Set up the WebSocket client.
+
+        Args:
+            _: The start frame (unused).
+        """
         self._leave_counter += 1
         if _.metadata and "call_id" in _.metadata:
             self._conversation_id = _.metadata["call_id"]
 
     def receive(self) -> typing.AsyncIterator[bytes | str]:
+        """Get an async iterator for receiving WebSocket messages.
+
+        Returns:
+            An async iterator yielding bytes or strings based on message type.
+        """
         return self._websocket.iter_bytes() if self._is_binary else self._websocket.iter_text()
 
     async def send(self, data: str | bytes):
+        """Send data through the WebSocket connection.
+
+        Args:
+            data: The data to send (string or bytes).
+        """
         try:
             if self._can_send():
                 if self._is_binary:
@@ -110,6 +154,7 @@ class FastAPIWebsocketClient:
                 await self.trigger_client_disconnected()
 
     async def disconnect(self):
+        """Disconnect the WebSocket client."""
         self._leave_counter -= 1
         if self._leave_counter > 0:
             return
@@ -123,27 +168,47 @@ class FastAPIWebsocketClient:
                 await self.trigger_client_disconnected()
 
     async def trigger_client_disconnected(self):
+        """Trigger the client disconnected callback."""
         await self._callbacks.on_client_disconnected(self._websocket)
 
     async def trigger_client_connected(self):
+        """Trigger the client connected callback."""
         await self._callbacks.on_client_connected(self._websocket)
 
     async def trigger_client_timeout(self):
+        """Trigger the client timeout callback."""
         await self._callbacks.on_session_timeout(self._websocket)
 
     def _can_send(self):
+        """Check if data can be sent through the WebSocket."""
         return self.is_connected and not self.is_closing
 
     @property
     def is_connected(self) -> bool:
+        """Check if the WebSocket is currently connected.
+
+        Returns:
+            True if the WebSocket is in connected state.
+        """
         return self._websocket.client_state == WebSocketState.CONNECTED
 
     @property
     def is_closing(self) -> bool:
+        """Check if the WebSocket is currently closing.
+
+        Returns:
+            True if the WebSocket is in the process of closing.
+        """
         return self._closing
 
 
 class FastAPIWebsocketInputTransport(BaseInputTransport):
+    """Input transport for FastAPI WebSocket connections.
+
+    Handles incoming WebSocket messages, deserializes frames, and manages
+    connection monitoring with optional session timeouts.
+    """
+
     def __init__(
         self,
         transport: BaseTransport,
@@ -151,6 +216,14 @@ class FastAPIWebsocketInputTransport(BaseInputTransport):
         params: FastAPIWebsocketParams,
         **kwargs,
     ):
+        """Initialize the WebSocket input transport.
+
+        Args:
+            transport: The parent transport instance.
+            client: The WebSocket client wrapper.
+            params: Transport configuration parameters.
+            **kwargs: Additional arguments passed to parent class.
+        """
         super().__init__(params, **kwargs)
         self._transport = transport
         self._client = client
@@ -162,6 +235,11 @@ class FastAPIWebsocketInputTransport(BaseInputTransport):
         self._initialized = False
 
     async def start(self, frame: StartFrame):
+        """Start the input transport and begin message processing.
+
+        Args:
+            frame: The start frame containing initialization parameters.
+        """
         await super().start(frame)
 
         if self._initialized:
@@ -180,6 +258,7 @@ class FastAPIWebsocketInputTransport(BaseInputTransport):
         await self.set_transport_ready(frame)
 
     async def _stop_tasks(self):
+        """Stop all running tasks."""
         if self._monitor_websocket_task:
             await self.cancel_task(self._monitor_websocket_task)
             self._monitor_websocket_task = None
@@ -188,24 +267,34 @@ class FastAPIWebsocketInputTransport(BaseInputTransport):
             self._receive_task = None
 
     async def stop(self, frame: EndFrame):
+        """Stop the input transport and cleanup resources.
+
+        Args:
+            frame: The end frame signaling transport shutdown.
+        """
         await super().stop(frame)
         await self._stop_tasks()
         await self._client.disconnect()
 
     async def cancel(self, frame: CancelFrame):
+        """Cancel the input transport and stop all processing.
+
+        Args:
+            frame: The cancel frame signaling immediate cancellation.
+        """
         await super().cancel(frame)
         await self._stop_tasks()
         await self._client.disconnect()
 
     async def cleanup(self):
+        """Clean up transport resources."""
         await super().cleanup()
         await self._transport.cleanup()
 
     async def _receive_messages(self):
+        """Main message receiving loop for WebSocket messages."""
         try:
-            async for message in WatchdogAsyncIterator(
-                self._client.receive(), manager=self.task_manager
-            ):
+            async for message in self._client.receive():
                 if not self._params.serializer:
                     continue
 
@@ -230,6 +319,12 @@ class FastAPIWebsocketInputTransport(BaseInputTransport):
 
 
 class FastAPIWebsocketOutputTransport(BaseOutputTransport):
+    """Output transport for FastAPI WebSocket connections.
+
+    Handles outgoing frame serialization, audio streaming with timing simulation,
+    and WebSocket message transmission with optional WAV header generation.
+    """
+
     def __init__(
         self,
         transport: BaseTransport,
@@ -237,6 +332,14 @@ class FastAPIWebsocketOutputTransport(BaseOutputTransport):
         params: FastAPIWebsocketParams,
         **kwargs,
     ):
+        """Initialize the WebSocket output transport.
+
+        Args:
+            transport: The parent transport instance.
+            client: The WebSocket client wrapper.
+            params: Transport configuration parameters.
+            **kwargs: Additional arguments passed to parent class.
+        """
         super().__init__(params, **kwargs)
 
         self._transport = transport
@@ -255,6 +358,11 @@ class FastAPIWebsocketOutputTransport(BaseOutputTransport):
         self._initialized = False
 
     async def start(self, frame: StartFrame):
+        """Start the output transport and initialize timing.
+
+        Args:
+            frame: The start frame containing initialization parameters.
+        """
         await super().start(frame)
 
         if self._initialized:
@@ -269,20 +377,37 @@ class FastAPIWebsocketOutputTransport(BaseOutputTransport):
         await self.set_transport_ready(frame)
 
     async def stop(self, frame: EndFrame):
+        """Stop the output transport and cleanup resources.
+
+        Args:
+            frame: The end frame signaling transport shutdown.
+        """
         await super().stop(frame)
         await self._write_frame(frame)
         await self._client.disconnect()
 
     async def cancel(self, frame: CancelFrame):
+        """Cancel the output transport and stop all processing.
+
+        Args:
+            frame: The cancel frame signaling immediate cancellation.
+        """
         await super().cancel(frame)
         await self._write_frame(frame)
         await self._client.disconnect()
 
     async def cleanup(self):
+        """Clean up transport resources."""
         await super().cleanup()
         await self._transport.cleanup()
 
     async def process_frame(self, frame: Frame, direction: FrameDirection):
+        """Process outgoing frames with special handling for interruptions.
+
+        Args:
+            frame: The frame to process.
+            direction: The direction of frame flow in the pipeline.
+        """
         await super().process_frame(frame, direction)
 
         if isinstance(frame, StartInterruptionFrame):
@@ -290,15 +415,20 @@ class FastAPIWebsocketOutputTransport(BaseOutputTransport):
             self._next_send_time = 0
 
     async def send_message(self, frame: TransportMessageFrame | TransportMessageUrgentFrame):
+        """Send a transport message frame.
+
+        Args:
+            frame: The transport message frame to send.
+        """
         await self._write_frame(frame)
 
     async def write_audio_frame(self, frame: OutputAudioRawFrame):
-        if self._client.is_closing:
-            return
+        """Write an audio frame to the WebSocket with timing simulation.
 
-        if not self._client.is_connected:
-            # Simulate audio playback with a sleep.
-            await self._write_audio_sleep()
+        Args:
+            frame: The output audio frame to write.
+        """
+        if self._client.is_closing or not self._client.is_connected:
             return
 
         frame = OutputAudioRawFrame(
@@ -327,6 +457,7 @@ class FastAPIWebsocketOutputTransport(BaseOutputTransport):
         await self._write_audio_sleep()
 
     async def _write_frame(self, frame: Frame):
+        """Serialize and send a frame through the WebSocket."""
         if not self._params.serializer:
             return
 
@@ -338,6 +469,7 @@ class FastAPIWebsocketOutputTransport(BaseOutputTransport):
             logger.error(f"{self} exception sending data: {e.__class__.__name__} ({e})")
 
     async def _write_audio_sleep(self):
+        """Simulate audio playback timing with appropriate delays."""
         # Simulate a clock.
         current_time = time.monotonic()
         sleep_duration = max(0, self._next_send_time - current_time)
@@ -349,6 +481,12 @@ class FastAPIWebsocketOutputTransport(BaseOutputTransport):
 
 
 class FastAPIWebsocketTransport(BaseTransport):
+    """FastAPI WebSocket transport for real-time audio/video streaming.
+
+    Provides bidirectional WebSocket communication with frame serialization,
+    session management, and event handling for client connections and timeouts.
+    """
+
     def __init__(
         self,
         websocket: WebSocket,
@@ -356,6 +494,14 @@ class FastAPIWebsocketTransport(BaseTransport):
         input_name: Optional[str] = None,
         output_name: Optional[str] = None,
     ):
+        """Initialize the FastAPI WebSocket transport.
+
+        Args:
+            websocket: The FastAPI WebSocket connection.
+            params: Transport configuration parameters.
+            input_name: Optional name for the input processor.
+            output_name: Optional name for the output processor.
+        """
         super().__init__(input_name=input_name, output_name=output_name)
 
         self._params = params
@@ -385,16 +531,29 @@ class FastAPIWebsocketTransport(BaseTransport):
         self._register_event_handler("on_session_timeout")
 
     def input(self) -> FastAPIWebsocketInputTransport:
+        """Get the input transport processor.
+
+        Returns:
+            The WebSocket input transport instance.
+        """
         return self._input
 
     def output(self) -> FastAPIWebsocketOutputTransport:
+        """Get the output transport processor.
+
+        Returns:
+            The WebSocket output transport instance.
+        """
         return self._output
 
     async def _on_client_connected(self, websocket):
+        """Handle client connected event."""
         await self._call_event_handler("on_client_connected", websocket)
 
     async def _on_client_disconnected(self, websocket):
+        """Handle client disconnected event."""
         await self._call_event_handler("on_client_disconnected", websocket)
 
     async def _on_session_timeout(self, websocket):
+        """Handle session timeout event."""
         await self._call_event_handler("on_session_timeout", websocket)