fix(streaming): accept bytearray and memoryview audio in stream()

assemblyai/__version__.py

@@ -1 +1 @@
-__version__ = "1.6.3"
+__version__ = "1.6.4"

assemblyai/streaming/v3/_base.py

@@ -214,6 +214,17 @@ def _resolve_options(
     )
 
 
+def _as_bytes_chunk(chunk: Any) -> Any:
+    """Copies a ``bytearray`` or ``memoryview`` chunk to ``bytes``.
+
+    The write loops only send ``bytes`` frames, so other bytes-like audio
+    buffers are converted when they are enqueued. Anything else passes through.
+    """
+    if isinstance(chunk, (bytearray, memoryview)):
+        return bytes(chunk)
+    return chunk
+
+
 class _BaseStreamingClient:
     """Sync/async-agnostic core for streaming clients.
 

assemblyai/streaming/v3/async_client.py

@@ -24,6 +24,7 @@
     _WS_HEADER_KW = "extra_headers"
 
 from ._base import (
+    _as_bytes_chunk,
     _BaseStreamingClient,
     _build_headers,
     _build_uri,
@@ -272,7 +273,9 @@ async def _close_websocket(self) -> None:
 
     async def stream(
         self,
-        data: Union[bytes, AsyncIterable[bytes], Iterable[bytes]],
+        data: Union[
+            bytes, bytearray, memoryview, AsyncIterable[bytes], Iterable[bytes]
+        ],
     ) -> None:
         # Loud on misuse (pre-connect), quiet on natural close (post-stop).
         # The first guards against silent data loss; the second keeps cleanup
@@ -281,21 +284,21 @@ async def stream(
         if stop_event.is_set():
             return
 
-        if isinstance(data, bytes):
-            await write_queue.put(data)
+        if isinstance(data, (bytes, bytearray, memoryview)):
+            await write_queue.put(bytes(data))
             return
 
         if isinstance(data, collections.abc.AsyncIterable):
             async for chunk in data:
                 if stop_event.is_set():
                     return
-                await write_queue.put(chunk)
+                await write_queue.put(_as_bytes_chunk(chunk))
             return
 
         for chunk in data:
             if stop_event.is_set():
                 return
-            await write_queue.put(chunk)
+            await write_queue.put(_as_bytes_chunk(chunk))
 
     async def set_params(self, params: RealTimeSessionParameters) -> None:
         write_queue, stop_event = self._ensure_connected("set_params")

assemblyai/streaming/v3/client.py

@@ -14,6 +14,7 @@
 from websockets.sync.client import connect as _ws_sync_connect
 
 from ._base import (
+    _as_bytes_chunk,
     _BaseStreamingClient,
     _build_headers,
     _build_uri,
@@ -208,26 +209,30 @@ def _close_websocket(self) -> None:
             logger.debug("Error closing websocket: %s", exc)
 
     def stream(
-        self, data: Union[bytes, Generator[bytes, None, None], Iterable[bytes]]
+        self,
+        data: Union[
+            bytes, bytearray, memoryview, Generator[bytes, None, None], Iterable[bytes]
+        ],
     ) -> None:
         """Send audio bytes to the server.
 
-        Accepts a raw ``bytes`` buffer or any (sync) iterable of ``bytes``.
+        Accepts a raw ``bytes``, ``bytearray`` or ``memoryview`` buffer, or any
+        (sync) iterable of those.
         Returns once all chunks are enqueued — the write thread does the
         actual sending. After ``disconnect()`` (or a connection drop) this
         becomes a silent no-op.
         """
         if self._stop_event.is_set():
             return
 
-        if isinstance(data, bytes):
-            self._write_queue.put(data)
+        if isinstance(data, (bytes, bytearray, memoryview)):
+            self._write_queue.put(bytes(data))
             return
 
         for chunk in data:
             if self._stop_event.is_set():
                 return
-            self._write_queue.put(chunk)
+            self._write_queue.put(_as_bytes_chunk(chunk))
 
     def set_params(self, params: RealTimeSessionParameters):
         message_dict = _normalize_min_turn_silence(_dump_model(params))

tests/unit/test_streaming.py

@@ -2533,3 +2533,36 @@ def test_options_zero_values_are_accepted():
     assert opts.max_connection_retries == 0
     assert opts.connection_retry_delay == 0
     assert opts.terminate_timeout == 0
+
+
+@pytest.mark.parametrize("wrap", [bytearray, memoryview])
+def test_client_stream_accepts_bytes_like_buffer(mocker: MockFixture, wrap):
+    _disable_rw_threads(mocker)
+    mocker.patch("assemblyai.streaming.v3.client.websocket_connect")
+
+    client = StreamingClient(
+        StreamingClientOptions(api_key="test", api_host="api.example.com")
+    )
+    client.connect(StreamingParameters(sample_rate=16000))
+    client.stream(wrap(b"\x01\x02\x03"))
+
+    queued = client._write_queue.get(timeout=1)
+    assert type(queued) is bytes
+    assert queued == b"\x01\x02\x03"
+    assert client._write_queue.empty()
+
+
+@pytest.mark.parametrize("wrap", [bytearray, memoryview])
+def test_client_stream_accepts_iterable_of_bytes_like_chunks(mocker: MockFixture, wrap):
+    _disable_rw_threads(mocker)
+    mocker.patch("assemblyai.streaming.v3.client.websocket_connect")
+
+    client = StreamingClient(
+        StreamingClientOptions(api_key="test", api_host="api.example.com")
+    )
+    client.connect(StreamingParameters(sample_rate=16000))
+    client.stream(iter([wrap(b"ab"), wrap(b"cd")]))
+
+    queued = [client._write_queue.get(timeout=1) for _ in range(2)]
+    assert queued == [b"ab", b"cd"]
+    assert all(type(chunk) is bytes for chunk in queued)

tests/unit/test_streaming_async.py

@@ -1661,3 +1661,26 @@ async def failing_connect(*_args, **_kwargs):
     assert attempts["n"] == 1
     assert len(errors) == 1
     assert errors[0].code == 401
+
+
+@pytest.mark.parametrize("wrap", [bytearray, memoryview])
+async def test_stream_accepts_bytes_like_buffer(mocker: MockFixture, wrap):
+    fake_ws = _FakeAsyncWebSocket()
+    _patch_connect(mocker, fake_ws)
+
+    client = AsyncStreamingClient(
+        StreamingClientOptions(api_key="test", api_host="api.example.com")
+    )
+    await client.connect(_default_params())
+
+    await client.stream(wrap(b"\x01\x02\x03"))
+    await client.stream(iter([wrap(b"ab"), wrap(b"cd")]))
+
+    for _ in range(50):
+        if len(fake_ws.sent) >= 3:
+            break
+        await asyncio.sleep(0.01)
+
+    assert fake_ws.sent == [b"\x01\x02\x03", b"ab", b"cd"]
+
+    await client.disconnect()