Fix FastAPIWebsocketTransport to handle both binary and text messages

This commit is contained in:
Luke Payyapilli
2026-01-08 11:25:18 -05:00
parent 973769b8bc
commit b52ae0e56b
2 changed files with 96 additions and 11 deletions

View File

@@ -33,7 +33,7 @@ from pipecat.frames.frames import (
StartFrame,
)
from pipecat.processors.frame_processor import FrameDirection
from pipecat.serializers.base_serializer import FrameSerializer, FrameSerializerType
from pipecat.serializers.base_serializer import FrameSerializer
from pipecat.transports.base_input import BaseInputTransport
from pipecat.transports.base_output import BaseOutputTransport
from pipecat.transports.base_transport import BaseTransport, TransportParams
@@ -77,6 +77,26 @@ class FastAPIWebsocketCallbacks(BaseModel):
on_session_timeout: Callable[[WebSocket], Awaitable[None]]
class _WebSocketMessageIterator:
"""Async iterator for WebSocket messages that yields both binary and text."""
def __init__(self, websocket: WebSocket):
self._websocket = websocket
def __aiter__(self):
return self
async def __anext__(self) -> bytes | str:
message = await self._websocket.receive()
if message["type"] == "websocket.disconnect":
raise StopAsyncIteration
if "bytes" in message and message["bytes"] is not None:
return message["bytes"]
if "text" in message and message["text"] is not None:
return message["text"]
raise StopAsyncIteration
class FastAPIWebsocketClient:
"""WebSocket client wrapper for handling connections and message passing.
@@ -84,17 +104,15 @@ class FastAPIWebsocketClient:
with support for both binary and text message types.
"""
def __init__(self, websocket: WebSocket, is_binary: bool, callbacks: FastAPIWebsocketCallbacks):
def __init__(self, websocket: WebSocket, 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
@@ -110,9 +128,9 @@ class FastAPIWebsocketClient:
"""Get an async iterator for receiving WebSocket messages.
Returns:
An async iterator yielding bytes or strings based on message type.
An async iterator yielding bytes or strings.
"""
return self._websocket.iter_bytes() if self._is_binary else self._websocket.iter_text()
return _WebSocketMessageIterator(self._websocket)
async def send(self, data: str | bytes):
"""Send data through the WebSocket connection.
@@ -122,7 +140,7 @@ class FastAPIWebsocketClient:
"""
try:
if self._can_send():
if self._is_binary:
if isinstance(data, bytes):
await self._websocket.send_bytes(data)
else:
await self._websocket.send_text(data)
@@ -510,10 +528,7 @@ class FastAPIWebsocketTransport(BaseTransport):
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._client = FastAPIWebsocketClient(websocket, self._callbacks)
self._input = FastAPIWebsocketInputTransport(
self, self._client, self._params, name=self._input_name

View File

@@ -0,0 +1,70 @@
#
# Copyright (c) 2024-2026, Daily
#
# SPDX-License-Identifier: BSD 2-Clause License
#
import unittest
from unittest.mock import AsyncMock
from pipecat.transports.websocket.fastapi import _WebSocketMessageIterator
class TestWebSocketMessageIterator(unittest.IsolatedAsyncioTestCase):
async def test_yields_binary_message(self):
mock_websocket = AsyncMock()
mock_websocket.receive.side_effect = [
{"type": "websocket.receive", "bytes": b"binary data", "text": None},
{"type": "websocket.disconnect"},
]
iterator = _WebSocketMessageIterator(mock_websocket)
messages = [msg async for msg in iterator]
self.assertEqual(len(messages), 1)
self.assertEqual(messages[0], b"binary data")
async def test_yields_text_message(self):
mock_websocket = AsyncMock()
mock_websocket.receive.side_effect = [
{"type": "websocket.receive", "bytes": None, "text": "text data"},
{"type": "websocket.disconnect"},
]
iterator = _WebSocketMessageIterator(mock_websocket)
messages = [msg async for msg in iterator]
self.assertEqual(len(messages), 1)
self.assertEqual(messages[0], "text data")
async def test_yields_mixed_messages(self):
mock_websocket = AsyncMock()
mock_websocket.receive.side_effect = [
{"type": "websocket.receive", "bytes": b"binary", "text": None},
{"type": "websocket.receive", "bytes": None, "text": "text"},
{"type": "websocket.receive", "bytes": b"more binary", "text": None},
{"type": "websocket.disconnect"},
]
iterator = _WebSocketMessageIterator(mock_websocket)
messages = [msg async for msg in iterator]
self.assertEqual(len(messages), 3)
self.assertEqual(messages[0], b"binary")
self.assertEqual(messages[1], "text")
self.assertEqual(messages[2], b"more binary")
async def test_stops_on_disconnect(self):
mock_websocket = AsyncMock()
mock_websocket.receive.side_effect = [
{"type": "websocket.disconnect"},
]
iterator = _WebSocketMessageIterator(mock_websocket)
messages = [msg async for msg in iterator]
self.assertEqual(len(messages), 0)
if __name__ == "__main__":
unittest.main()