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

pipecat/transports/network/websocket_client.py

@@ -4,6 +4,13 @@
 # SPDX-License-Identifier: BSD 2-Clause License
 #
 
+"""WebSocket client transport implementation for Pipecat.
+
+This module provides a WebSocket client transport that enables bidirectional
+communication over WebSocket connections, with support for audio streaming,
+frame serialization, and connection management.
+"""
+
 import asyncio
 import io
 import time
@@ -13,6 +20,7 @@ from typing import Awaitable, Callable, Optional
 import websockets
 from loguru import logger
 from pydantic.main import BaseModel
+from websockets.asyncio.client import connect as websocket_connect
 
 from pipecat.frames.frames import (
     CancelFrame,
@@ -34,17 +42,38 @@ from pipecat.utils.asyncio.task_manager import BaseTaskManager
 
 
 class WebsocketClientParams(TransportParams):
+    """Configuration parameters for WebSocket client transport.
+
+    Parameters:
+        add_wav_header: Whether to add WAV headers to audio frames.
+        serializer: Frame serializer for encoding/decoding messages.
+    """
+
     add_wav_header: bool = True
     serializer: Optional[FrameSerializer] = None
 
 
 class WebsocketClientCallbacks(BaseModel):
+    """Callback functions for WebSocket client events.
+
+    Parameters:
+        on_connected: Called when WebSocket connection is established.
+        on_disconnected: Called when WebSocket connection is closed.
+        on_message: Called when a message is received from the WebSocket.
+    """
+
     on_connected: Callable[[websockets.WebSocketClientProtocol], Awaitable[None]]
     on_disconnected: Callable[[websockets.WebSocketClientProtocol], Awaitable[None]]
     on_message: Callable[[websockets.WebSocketClientProtocol, websockets.Data], Awaitable[None]]
 
 
 class WebsocketClientSession:
+    """Manages a WebSocket client connection session.
+
+    Handles connection lifecycle, message sending/receiving, and provides
+    callback mechanisms for connection events.
+    """
+
     def __init__(
         self,
         uri: str,
@@ -52,6 +81,14 @@ class WebsocketClientSession:
         callbacks: WebsocketClientCallbacks,
         transport_name: str,
     ):
+        """Initialize the WebSocket client session.
+
+        Args:
+            uri: The WebSocket URI to connect to.
+            params: Configuration parameters for the session.
+            callbacks: Callback functions for session events.
+            transport_name: Name of the parent transport for logging.
+        """
         self._uri = uri
         self._params = params
         self._callbacks = callbacks
@@ -63,6 +100,14 @@ class WebsocketClientSession:
 
     @property
     def task_manager(self) -> BaseTaskManager:
+        """Get the task manager for this session.
+
+        Returns:
+            The task manager instance.
+
+        Raises:
+            Exception: If task manager is not initialized.
+        """
         if not self._task_manager:
             raise Exception(
                 f"{self._transport_name}::WebsocketClientSession: TaskManager not initialized (pipeline not started?)"
@@ -70,16 +115,22 @@ class WebsocketClientSession:
         return self._task_manager
 
     async def setup(self, task_manager: BaseTaskManager):
+        """Set up the session with a task manager.
+
+        Args:
+            task_manager: The task manager to use for session tasks.
+        """
         self._leave_counter += 1
         if not self._task_manager:
             self._task_manager = task_manager
 
     async def connect(self):
+        """Connect to the WebSocket server."""
         if self._websocket:
             return
 
         try:
-            self._websocket = await websockets.connect(uri=self._uri, open_timeout=10)
+            self._websocket = await websocket_connect(uri=self._uri, open_timeout=10)
             self._client_task = self.task_manager.create_task(
                 self._client_task_handler(),
                 f"{self._transport_name}::WebsocketClientSession::_client_task_handler",
@@ -89,6 +140,7 @@ class WebsocketClientSession:
             logger.error(f"Timeout connecting to {self._uri}")
 
     async def disconnect(self):
+        """Disconnect from the WebSocket server."""
         self._leave_counter -= 1
         if not self._websocket or self._leave_counter > 0:
             return
@@ -99,6 +151,11 @@ class WebsocketClientSession:
         self._websocket = None
 
     async def send(self, message: websockets.Data):
+        """Send a message through the WebSocket connection.
+
+        Args:
+            message: The message data to send.
+        """
         try:
             if self._websocket:
                 await self._websocket.send(message)
@@ -106,6 +163,7 @@ class WebsocketClientSession:
             logger.error(f"{self} exception sending data: {e.__class__.__name__} ({e})")
 
     async def _client_task_handler(self):
+        """Handle incoming messages from the WebSocket connection."""
         try:
             # Handle incoming messages
             async for message in self._websocket:
@@ -116,16 +174,30 @@ class WebsocketClientSession:
         await self._callbacks.on_disconnected(self._websocket)
 
     def __str__(self):
+        """String representation of the WebSocket client session."""
         return f"{self._transport_name}::WebsocketClientSession"
 
 
 class WebsocketClientInputTransport(BaseInputTransport):
+    """WebSocket client input transport for receiving frames.
+
+    Handles incoming WebSocket messages, deserializes them to frames,
+    and pushes them downstream in the processing pipeline.
+    """
+
     def __init__(
         self,
         transport: BaseTransport,
         session: WebsocketClientSession,
         params: WebsocketClientParams,
     ):
+        """Initialize the WebSocket client input transport.
+
+        Args:
+            transport: The parent transport instance.
+            session: The WebSocket session to use for communication.
+            params: Configuration parameters for the transport.
+        """
         super().__init__(params)
 
         self._transport = transport
@@ -136,10 +208,20 @@ class WebsocketClientInputTransport(BaseInputTransport):
         self._initialized = False
 
     async def setup(self, setup: FrameProcessorSetup):
+        """Set up the input transport with the frame processor setup.
+
+        Args:
+            setup: The frame processor setup configuration.
+        """
         await super().setup(setup)
         await self._session.setup(setup.task_manager)
 
     async def start(self, frame: StartFrame):
+        """Start the input transport and initialize the WebSocket connection.
+
+        Args:
+            frame: The start frame containing initialization parameters.
+        """
         await super().start(frame)
 
         if self._initialized:
@@ -153,18 +235,35 @@ class WebsocketClientInputTransport(BaseInputTransport):
         await self.set_transport_ready(frame)
 
     async def stop(self, frame: EndFrame):
+        """Stop the input transport and disconnect from WebSocket.
+
+        Args:
+            frame: The end frame signaling transport shutdown.
+        """
         await super().stop(frame)
         await self._session.disconnect()
 
     async def cancel(self, frame: CancelFrame):
+        """Cancel the input transport and disconnect from WebSocket.
+
+        Args:
+            frame: The cancel frame signaling immediate cancellation.
+        """
         await super().cancel(frame)
         await self._session.disconnect()
 
     async def cleanup(self):
+        """Clean up the input transport resources."""
         await super().cleanup()
         await self._transport.cleanup()
 
     async def on_message(self, websocket, message):
+        """Handle incoming WebSocket messages.
+
+        Args:
+            websocket: The WebSocket connection that received the message.
+            message: The received message data.
+        """
         if not self._params.serializer:
             return
         frame = await self._params.serializer.deserialize(message)
@@ -177,12 +276,25 @@ class WebsocketClientInputTransport(BaseInputTransport):
 
 
 class WebsocketClientOutputTransport(BaseOutputTransport):
+    """WebSocket client output transport for sending frames.
+
+    Handles outgoing frames, serializes them for WebSocket transmission,
+    and manages audio streaming with proper timing simulation.
+    """
+
     def __init__(
         self,
         transport: BaseTransport,
         session: WebsocketClientSession,
         params: WebsocketClientParams,
     ):
+        """Initialize the WebSocket client output transport.
+
+        Args:
+            transport: The parent transport instance.
+            session: The WebSocket session to use for communication.
+            params: Configuration parameters for the transport.
+        """
         super().__init__(params)
 
         self._transport = transport
@@ -201,10 +313,20 @@ class WebsocketClientOutputTransport(BaseOutputTransport):
         self._initialized = False
 
     async def setup(self, setup: FrameProcessorSetup):
+        """Set up the output transport with the frame processor setup.
+
+        Args:
+            setup: The frame processor setup configuration.
+        """
         await super().setup(setup)
         await self._session.setup(setup.task_manager)
 
     async def start(self, frame: StartFrame):
+        """Start the output transport and initialize the WebSocket connection.
+
+        Args:
+            frame: The start frame containing initialization parameters.
+        """
         await super().start(frame)
 
         if self._initialized:
@@ -219,21 +341,42 @@ class WebsocketClientOutputTransport(BaseOutputTransport):
         await self.set_transport_ready(frame)
 
     async def stop(self, frame: EndFrame):
+        """Stop the output transport and disconnect from WebSocket.
+
+        Args:
+            frame: The end frame signaling transport shutdown.
+        """
         await super().stop(frame)
         await self._session.disconnect()
 
     async def cancel(self, frame: CancelFrame):
+        """Cancel the output transport and disconnect from WebSocket.
+
+        Args:
+            frame: The cancel frame signaling immediate cancellation.
+        """
         await super().cancel(frame)
         await self._session.disconnect()
 
     async def cleanup(self):
+        """Clean up the output transport resources."""
         await super().cleanup()
         await self._transport.cleanup()
 
     async def send_message(self, frame: TransportMessageFrame | TransportMessageUrgentFrame):
+        """Send a transport message through the WebSocket.
+
+        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 with optional WAV header.
+
+        Args:
+            frame: The output audio frame to write.
+        """
         frame = OutputAudioRawFrame(
             audio=frame.audio,
             sample_rate=self.sample_rate,
@@ -260,6 +403,7 @@ class WebsocketClientOutputTransport(BaseOutputTransport):
         await self._write_audio_sleep()
 
     async def _write_frame(self, frame: Frame):
+        """Write a frame to the WebSocket after serialization."""
         if not self._params.serializer:
             return
         payload = await self._params.serializer.serialize(frame)
@@ -267,6 +411,7 @@ class WebsocketClientOutputTransport(BaseOutputTransport):
             await self._session.send(payload)
 
     async def _write_audio_sleep(self):
+        """Simulate audio playback timing with sleep delays."""
         # Simulate a clock.
         current_time = time.monotonic()
         sleep_duration = max(0, self._next_send_time - current_time)
@@ -278,11 +423,23 @@ class WebsocketClientOutputTransport(BaseOutputTransport):
 
 
 class WebsocketClientTransport(BaseTransport):
+    """WebSocket client transport for bidirectional communication.
+
+    Provides a complete WebSocket client transport implementation with
+    input and output capabilities, connection management, and event handling.
+    """
+
     def __init__(
         self,
         uri: str,
         params: Optional[WebsocketClientParams] = None,
     ):
+        """Initialize the WebSocket client transport.
+
+        Args:
+            uri: The WebSocket URI to connect to.
+            params: Optional configuration parameters for the transport.
+        """
         super().__init__()
 
         self._params = params or WebsocketClientParams()
@@ -304,21 +461,34 @@ class WebsocketClientTransport(BaseTransport):
         self._register_event_handler("on_disconnected")
 
     def input(self) -> WebsocketClientInputTransport:
+        """Get the input transport for receiving frames.
+
+        Returns:
+            The WebSocket client input transport instance.
+        """
         if not self._input:
             self._input = WebsocketClientInputTransport(self, self._session, self._params)
         return self._input
 
     def output(self) -> WebsocketClientOutputTransport:
+        """Get the output transport for sending frames.
+
+        Returns:
+            The WebSocket client output transport instance.
+        """
         if not self._output:
             self._output = WebsocketClientOutputTransport(self, self._session, self._params)
         return self._output
 
     async def _on_connected(self, websocket):
+        """Handle WebSocket connection established event."""
         await self._call_event_handler("on_connected", websocket)
 
     async def _on_disconnected(self, websocket):
+        """Handle WebSocket connection closed event."""
         await self._call_event_handler("on_disconnected", websocket)
 
     async def _on_message(self, websocket, message):
+        """Handle incoming WebSocket message."""
         if self._input:
             await self._input.on_message(websocket, message)