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

pipecat/transports/network/websocket_server.py

@@ -4,6 +4,13 @@
 # SPDX-License-Identifier: BSD 2-Clause License
 #
 
+"""WebSocket server transport implementation for Pipecat.
+
+This module provides WebSocket server transport functionality for real-time
+audio and data streaming, including client connection management, session
+handling, and frame serialization.
+"""
+
 import asyncio
 import io
 import time
@@ -32,6 +39,8 @@ from pipecat.transports.base_transport import BaseTransport, TransportParams
 
 try:
     import websockets
+    from websockets.asyncio.server import serve as websocket_serve
+    from websockets.protocol import State
 except ModuleNotFoundError as e:
     logger.error(f"Exception: {e}")
     logger.error("In order to use websockets, you need to `pip install pipecat-ai[websocket]`.")
@@ -39,12 +48,29 @@ except ModuleNotFoundError as e:
 
 
 class WebsocketServerParams(TransportParams):
+    """Configuration parameters for WebSocket server transport.
+
+    Parameters:
+        add_wav_header: Whether to add WAV headers to audio frames.
+        serializer: Frame serializer for message encoding/decoding.
+        session_timeout: Timeout in seconds for client sessions.
+    """
+
     add_wav_header: bool = False
     serializer: Optional[FrameSerializer] = None
     session_timeout: Optional[int] = None
 
 
 class WebsocketServerCallbacks(BaseModel):
+    """Callback functions for WebSocket server events.
+
+    Parameters:
+        on_client_connected: Called when a client connects to the server.
+        on_client_disconnected: Called when a client disconnects from the server.
+        on_session_timeout: Called when a client session times out.
+        on_websocket_ready: Called when the WebSocket server is ready to accept connections.
+    """
+
     on_client_connected: Callable[[websockets.WebSocketServerProtocol], Awaitable[None]]
     on_client_disconnected: Callable[[websockets.WebSocketServerProtocol], Awaitable[None]]
     on_session_timeout: Callable[[websockets.WebSocketServerProtocol], Awaitable[None]]
@@ -52,6 +78,12 @@ class WebsocketServerCallbacks(BaseModel):
 
 
 class WebsocketServerInputTransport(BaseInputTransport):
+    """WebSocket server input transport for receiving client data.
+
+    Handles incoming WebSocket connections, message processing, and client
+    session management including timeout monitoring and connection lifecycle.
+    """
+
     def __init__(
         self,
         transport: BaseTransport,
@@ -61,6 +93,16 @@ class WebsocketServerInputTransport(BaseInputTransport):
         callbacks: WebsocketServerCallbacks,
         **kwargs,
     ):
+        """Initialize the WebSocket server input transport.
+
+        Args:
+            transport: The parent transport instance.
+            host: Host address to bind the WebSocket server to.
+            port: Port number to bind the WebSocket server to.
+            params: WebSocket server configuration parameters.
+            callbacks: Callback functions for WebSocket events.
+            **kwargs: Additional arguments passed to parent class.
+        """
         super().__init__(params, **kwargs)
 
         self._transport = transport
@@ -82,6 +124,11 @@ class WebsocketServerInputTransport(BaseInputTransport):
         self._initialized = False
 
     async def start(self, frame: StartFrame):
+        """Start the WebSocket server and initialize components.
+
+        Args:
+            frame: The start frame containing initialization parameters.
+        """
         await super().start(frame)
 
         if self._initialized:
@@ -96,16 +143,26 @@ class WebsocketServerInputTransport(BaseInputTransport):
         await self.set_transport_ready(frame)
 
     async def stop(self, frame: EndFrame):
+        """Stop the WebSocket server and cleanup resources.
+
+        Args:
+            frame: The end frame signaling transport shutdown.
+        """
         await super().stop(frame)
         self._stop_server_event.set()
         if self._monitor_task:
             await self.cancel_task(self._monitor_task)
             self._monitor_task = None
         if self._server_task:
-            await self.wait_for_task(self._server_task)
+            await self._server_task
             self._server_task = None
 
     async def cancel(self, frame: CancelFrame):
+        """Cancel the WebSocket server and stop all processing.
+
+        Args:
+            frame: The cancel frame signaling immediate cancellation.
+        """
         await super().cancel(frame)
         if self._monitor_task:
             await self.cancel_task(self._monitor_task)
@@ -115,16 +172,19 @@ class WebsocketServerInputTransport(BaseInputTransport):
             self._server_task = None
 
     async def cleanup(self):
+        """Cleanup resources and parent transport."""
         await super().cleanup()
         await self._transport.cleanup()
 
     async def _server_task_handler(self):
+        """Handle WebSocket server startup and client connections."""
         logger.info(f"Starting websocket server on {self._host}:{self._port}")
-        async with websockets.serve(self._client_handler, self._host, self._port) as server:
+        async with websocket_serve(self._client_handler, self._host, self._port) as server:
             await self._callbacks.on_websocket_ready()
             await self._stop_server_event.wait()
 
-    async def _client_handler(self, websocket: websockets.WebSocketServerProtocol, path):
+    async def _client_handler(self, websocket: websockets.WebSocketServerProtocol):
+        """Handle individual client connections and message processing."""
         logger.info(f"New client connection from {websocket.remote_address}")
         if self._websocket:
             await self._websocket.close()
@@ -170,12 +230,10 @@ class WebsocketServerInputTransport(BaseInputTransport):
     async def _monitor_websocket(
         self, websocket: websockets.WebSocketServerProtocol, session_timeout: int
     ):
-        """Wait for session_timeout seconds, if the websocket is still open,
-        trigger timeout event.
-        """
+        """Monitor WebSocket connection for session timeout."""
         try:
             await asyncio.sleep(session_timeout)
-            if not websocket.closed:
+            if websocket.state is not State.CLOSED:
                 await self._callbacks.on_session_timeout(websocket)
         except asyncio.CancelledError:
             logger.info(f"Monitoring task cancelled for: {websocket.remote_address}")
@@ -183,7 +241,20 @@ class WebsocketServerInputTransport(BaseInputTransport):
 
 
 class WebsocketServerOutputTransport(BaseOutputTransport):
+    """WebSocket server output transport for sending data to clients.
+
+    Handles outgoing frame serialization, audio streaming with timing control,
+    and client connection management for WebSocket communication.
+    """
+
     def __init__(self, transport: BaseTransport, params: WebsocketServerParams, **kwargs):
+        """Initialize the WebSocket server output transport.
+
+        Args:
+            transport: The parent transport instance.
+            params: WebSocket server configuration parameters.
+            **kwargs: Additional arguments passed to parent class.
+        """
         super().__init__(params, **kwargs)
 
         self._transport = transport
@@ -203,12 +274,22 @@ class WebsocketServerOutputTransport(BaseOutputTransport):
         self._initialized = False
 
     async def set_client_connection(self, websocket: Optional[websockets.WebSocketServerProtocol]):
+        """Set the active client WebSocket connection.
+
+        Args:
+            websocket: The WebSocket connection to set as active, or None to clear.
+        """
         if self._websocket:
             await self._websocket.close()
             logger.warning("Only one client allowed, using new connection")
         self._websocket = websocket
 
     async def start(self, frame: StartFrame):
+        """Start the output transport and initialize components.
+
+        Args:
+            frame: The start frame containing initialization parameters.
+        """
         await super().start(frame)
 
         if self._initialized:
@@ -222,18 +303,35 @@ class WebsocketServerOutputTransport(BaseOutputTransport):
         await self.set_transport_ready(frame)
 
     async def stop(self, frame: EndFrame):
+        """Stop the output transport and send final frame.
+
+        Args:
+            frame: The end frame signaling transport shutdown.
+        """
         await super().stop(frame)
         await self._write_frame(frame)
 
     async def cancel(self, frame: CancelFrame):
+        """Cancel the output transport and send cancellation frame.
+
+        Args:
+            frame: The cancel frame signaling immediate cancellation.
+        """
         await super().cancel(frame)
         await self._write_frame(frame)
 
     async def cleanup(self):
+        """Cleanup resources and parent transport."""
         await super().cleanup()
         await self._transport.cleanup()
 
     async def process_frame(self, frame: Frame, direction: FrameDirection):
+        """Process frames and handle interruption timing.
+
+        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):
@@ -241,12 +339,20 @@ class WebsocketServerOutputTransport(BaseOutputTransport):
             self._next_send_time = 0
 
     async def send_message(self, frame: TransportMessageFrame | TransportMessageUrgentFrame):
+        """Send a transport message frame to the client.
+
+        Args:
+            frame: The transport message frame to send.
+        """
         await self._write_frame(frame)
 
     async def write_audio_frame(self, frame: OutputAudioRawFrame):
+        """Write an audio frame to the WebSocket client with timing control.
+
+        Args:
+            frame: The output audio frame to write.
+        """
         if not self._websocket:
-            # Simulate audio playback with a sleep.
-            await self._write_audio_sleep()
             return
 
         frame = OutputAudioRawFrame(
@@ -275,6 +381,7 @@ class WebsocketServerOutputTransport(BaseOutputTransport):
         await self._write_audio_sleep()
 
     async def _write_frame(self, frame: Frame):
+        """Serialize and send a frame to the WebSocket client."""
         if not self._params.serializer:
             return
 
@@ -286,6 +393,7 @@ class WebsocketServerOutputTransport(BaseOutputTransport):
             logger.error(f"{self} exception sending data: {e.__class__.__name__} ({e})")
 
     async def _write_audio_sleep(self):
+        """Simulate audio device timing by sleeping between audio chunks."""
         # Simulate a clock.
         current_time = time.monotonic()
         sleep_duration = max(0, self._next_send_time - current_time)
@@ -297,6 +405,13 @@ class WebsocketServerOutputTransport(BaseOutputTransport):
 
 
 class WebsocketServerTransport(BaseTransport):
+    """WebSocket server transport for bidirectional real-time communication.
+
+    Provides a complete WebSocket server implementation with separate input and
+    output transports, client connection management, and event handling for
+    real-time audio and data streaming applications.
+    """
+
     def __init__(
         self,
         params: WebsocketServerParams,
@@ -305,6 +420,15 @@ class WebsocketServerTransport(BaseTransport):
         input_name: Optional[str] = None,
         output_name: Optional[str] = None,
     ):
+        """Initialize the WebSocket server transport.
+
+        Args:
+            params: WebSocket server configuration parameters.
+            host: Host address to bind the server to. Defaults to "localhost".
+            port: Port number to bind the server to. Defaults to 8765.
+            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._host = host
         self._port = port
@@ -328,6 +452,11 @@ class WebsocketServerTransport(BaseTransport):
         self._register_event_handler("on_websocket_ready")
 
     def input(self) -> WebsocketServerInputTransport:
+        """Get the input transport for receiving client data.
+
+        Returns:
+            The WebSocket server input transport instance.
+        """
         if not self._input:
             self._input = WebsocketServerInputTransport(
                 self, self._host, self._port, self._params, self._callbacks, name=self._input_name
@@ -335,6 +464,11 @@ class WebsocketServerTransport(BaseTransport):
         return self._input
 
     def output(self) -> WebsocketServerOutputTransport:
+        """Get the output transport for sending data to clients.
+
+        Returns:
+            The WebSocket server output transport instance.
+        """
         if not self._output:
             self._output = WebsocketServerOutputTransport(
                 self, self._params, name=self._output_name
@@ -342,6 +476,7 @@ class WebsocketServerTransport(BaseTransport):
         return self._output
 
     async def _on_client_connected(self, websocket):
+        """Handle client connection events."""
         if self._output:
             await self._output.set_client_connection(websocket)
             await self._call_event_handler("on_client_connected", websocket)
@@ -349,6 +484,7 @@ class WebsocketServerTransport(BaseTransport):
             logger.error("A WebsocketServerTransport output is missing in the pipeline")
 
     async def _on_client_disconnected(self, websocket):
+        """Handle client disconnection events."""
         if self._output:
             await self._output.set_client_connection(None)
             await self._call_event_handler("on_client_disconnected", websocket)
@@ -356,7 +492,9 @@ class WebsocketServerTransport(BaseTransport):
             logger.error("A WebsocketServerTransport output is missing in the pipeline")
 
     async def _on_session_timeout(self, websocket):
+        """Handle client session timeout events."""
         await self._call_event_handler("on_session_timeout", websocket)
 
     async def _on_websocket_ready(self):
+        """Handle WebSocket server ready events."""
         await self._call_event_handler("on_websocket_ready")