diff --git a/assemblyai/__version__.py b/assemblyai/__version__.py index 31e744e..dc79f8f 100644 --- a/assemblyai/__version__.py +++ b/assemblyai/__version__.py @@ -1 +1 @@ -__version__ = "1.6.3" +__version__ = "1.6.4" diff --git a/assemblyai/streaming/v3/_base.py b/assemblyai/streaming/v3/_base.py index 98db6dc..02b1f29 100644 --- a/assemblyai/streaming/v3/_base.py +++ b/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. diff --git a/assemblyai/streaming/v3/async_client.py b/assemblyai/streaming/v3/async_client.py index fe34bed..8a2b9cc 100644 --- a/assemblyai/streaming/v3/async_client.py +++ b/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") diff --git a/assemblyai/streaming/v3/client.py b/assemblyai/streaming/v3/client.py index d0c510f..fa78e05 100644 --- a/assemblyai/streaming/v3/client.py +++ b/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,11 +209,15 @@ 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. @@ -220,14 +225,14 @@ def stream( 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)) diff --git a/tests/unit/test_streaming.py b/tests/unit/test_streaming.py index fe5de94..85bd1ee 100644 --- a/tests/unit/test_streaming.py +++ b/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) diff --git a/tests/unit/test_streaming_async.py b/tests/unit/test_streaming_async.py index 77a455c..00174fd 100644 --- a/tests/unit/test_streaming_async.py +++ b/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()