dv-pipecat-ai 0.0.82.dev815__py3-none-any.whl → 0.0.82.dev857__py3-none-any.whl

pipecat/transports/network/fastapi_websocket.py

@@ -11,549 +11,15 @@ using FastAPI and WebSocket connections. Supports binary and text serialization
 with configurable session timeouts and WAV header generation.
 """
 
-import asyncio
-import io
-import time
-import traceback
-import typing
-import wave
-from typing import Awaitable, Callable, Optional
-
-from loguru import logger
-from pydantic import BaseModel
-
-from pipecat.frames.frames import (
-    CancelFrame,
-    EndFrame,
-    Frame,
-    InputAudioRawFrame,
-    OutputAudioRawFrame,
-    StartFrame,
-    StartInterruptionFrame,
-    TransportMessageFrame,
-    TransportMessageUrgentFrame,
-)
-from pipecat.processors.frame_processor import FrameDirection
-from pipecat.serializers.base_serializer import FrameSerializer, FrameSerializerType
-from pipecat.transports.base_input import BaseInputTransport
-from pipecat.transports.base_output import BaseOutputTransport
-from pipecat.transports.base_transport import BaseTransport, TransportParams
-
-try:
-    from fastapi import WebSocket
-    from starlette.websockets import WebSocketDisconnect, WebSocketState
-except ModuleNotFoundError as e:
-    logger.error(f"Exception: {e}")
-    logger.error(
-        "In order to use FastAPI websockets, you need to `pip install pipecat-ai[websocket]`."
-    )
-    raise Exception(f"Missing module: {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:
-    """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
-        self._callbacks = callbacks
-        self._leave_counter = 0
-        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:
-                    await self._websocket.send_bytes(data)
-                else:
-                    await self._websocket.send_text(data)
-        except Exception as e:
-            if isinstance(e, WebSocketDisconnect):
-                logger.warning(
-                    f"{self} WebSocket disconnected during send: {e}, application_state: {self._websocket.application_state}",
-                    call_id=self._conversation_id,
-                )
-            else:
-                logger.error(
-                    f"{self} exception sending data: {e.__class__.__name__} ({e}), application_state: {self._websocket.application_state}",
-                    call_id=self._conversation_id,
-                )
-            # For some reason the websocket is disconnected, and we are not able to send data
-            # So let's properly handle it and disconnect the transport if it is not already disconnecting
-            if (
-                self._websocket.application_state == WebSocketState.DISCONNECTED
-                and not self.is_closing
-            ):
-                logger.warning(
-                    "Closing already disconnected websocket!", call_id=self._conversation_id
-                )
-                self._closing = True
-                await self.trigger_client_disconnected()
-
-    async def disconnect(self):
-        """Disconnect the WebSocket client."""
-        self._leave_counter -= 1
-        if self._leave_counter > 0:
-            return
-        if self.is_connected and not self.is_closing:
-            self._closing = True
-            try:
-                await self._websocket.close()
-            except Exception as e:
-                logger.error(f"{self} exception while closing the websocket: {e}")
-            finally:
-                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,
-        client: FastAPIWebsocketClient,
-        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
-        self._params = params
-        self._receive_task = None
-        self._monitor_websocket_task = None
-
-        # Whether we have seen a StartFrame already.
-        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:
-            return
-
-        self._initialized = True
-
-        await self._client.setup(frame)
-        if self._params.serializer:
-            await self._params.serializer.setup(frame)
-        if not self._monitor_websocket_task and self._params.session_timeout:
-            self._monitor_websocket_task = self.create_task(self._monitor_websocket())
-        await self._client.trigger_client_connected()
-        if not self._receive_task:
-            self._receive_task = self.create_task(self._receive_messages())
-        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
-        if self._receive_task:
-            await self.cancel_task(self._receive_task)
-            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 self._client.receive():
-                if not self._params.serializer:
-                    continue
-
-                frame = await self._params.serializer.deserialize(message)
-
-                if not frame:
-                    continue
-
-                if isinstance(frame, InputAudioRawFrame):
-                    await self.push_audio_frame(frame)
-                else:
-                    await self.push_frame(frame)
-        except Exception as e:
-            logger.error(f"{self} exception receiving data: {e.__class__.__name__} ({e})")
-
-        await self._client.trigger_client_disconnected()
-
-    async def _monitor_websocket(self):
-        """Wait for self._params.session_timeout seconds, if the websocket is still open, trigger timeout event."""
-        await asyncio.sleep(self._params.session_timeout)
-        await self._client.trigger_client_timeout()
-
-
-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,
-        client: FastAPIWebsocketClient,
-        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
-        self._client = client
-        self._params = params
-
-        # write_audio_frame() is called quickly, as soon as we get audio
-        # (e.g. from the TTS), and since this is just a network connection we
-        # would be sending it to quickly. Instead, we want to block to emulate
-        # an audio device, this is what the send interval is. It will be
-        # computed on StartFrame.
-        self._send_interval = 0
-        self._next_send_time = 0
-
-        # Whether we have seen a StartFrame already.
-        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:
-            return
-
-        self._initialized = True
-
-        await self._client.setup(frame)
-        if self._params.serializer:
-            await self._params.serializer.setup(frame)
-        self._send_interval = (self.audio_chunk_size / self.sample_rate) / 2
-        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):
-            await self._write_frame(frame)
-            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):
-        """Write an audio frame to the WebSocket with timing simulation.
-
-        Args:
-            frame: The output audio frame to write.
-        """
-        if self._client.is_closing or not self._client.is_connected:
-            return
-
-        frame = OutputAudioRawFrame(
-            audio=frame.audio,
-            sample_rate=self.sample_rate,
-            num_channels=self._params.audio_out_channels,
-        )
-
-        if self._params.add_wav_header:
-            with io.BytesIO() as buffer:
-                with wave.open(buffer, "wb") as wf:
-                    wf.setsampwidth(2)
-                    wf.setnchannels(frame.num_channels)
-                    wf.setframerate(frame.sample_rate)
-                    wf.writeframes(frame.audio)
-                wav_frame = OutputAudioRawFrame(
-                    buffer.getvalue(),
-                    sample_rate=frame.sample_rate,
-                    num_channels=frame.num_channels,
-                )
-                frame = wav_frame
-
-        await self._write_frame(frame)
-
-        # Simulate audio playback with a sleep.
-        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
-
-        try:
-            payload = await self._params.serializer.serialize(frame)
-            if payload:
-                await self._client.send(payload)
-        except Exception as e:
-            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)
-        await asyncio.sleep(sleep_duration)
-        if sleep_duration == 0:
-            self._next_send_time = time.monotonic() + self._send_interval
-        else:
-            self._next_send_time += self._send_interval
-
-
-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,
-        params: FastAPIWebsocketParams,
-        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
-
-        self._callbacks = FastAPIWebsocketCallbacks(
-            on_client_connected=self._on_client_connected,
-            on_client_disconnected=self._on_client_disconnected,
-            on_session_timeout=self._on_session_timeout,
-        )
-
-        is_binary = False
-        if self._params.serializer:
-            is_binary = self._params.serializer.type == FrameSerializerType.BINARY
-        self._client = FastAPIWebsocketClient(websocket, is_binary, self._callbacks)
-
-        self._input = FastAPIWebsocketInputTransport(
-            self, self._client, self._params, name=self._input_name
-        )
-        self._output = FastAPIWebsocketOutputTransport(
-            self, self._client, self._params, name=self._output_name
-        )
-
-        # Register supported handlers. The user will only be able to register
-        # these handlers.
-        self._register_event_handler("on_client_connected")
-        self._register_event_handler("on_client_disconnected")
-        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)
+import warnings
+
+from pipecat.transports.websocket.fastapi import *
+
+with warnings.catch_warnings():
+    warnings.simplefilter("always")
+    warnings.warn(
+        "Module `pipecat.transports.network.fastapi_websocket` is deprecated, "
+        "use `pipecat.transports.websocket.fastapi` instead.",
+        DeprecationWarning,
+        stacklevel=2,
+    )