diff --git a/assemblyai/__version__.py b/assemblyai/__version__.py index 0a0a43a..38cf6db 100644 --- a/assemblyai/__version__.py +++ b/assemblyai/__version__.py @@ -1 +1 @@ -__version__ = "1.9.0" +__version__ = "1.9.1" diff --git a/assemblyai/tts/_base_ws.py b/assemblyai/tts/_base_ws.py index 83e34fb..1a0ffba 100644 --- a/assemblyai/tts/_base_ws.py +++ b/assemblyai/tts/_base_ws.py @@ -19,6 +19,8 @@ # Characters a single `Generate` frame may carry (Unicode code points) when # `Begin` does not report the limit in force. MAX_GENERATE_CHARS = 2000 +# The unspoken characters the API holds by default before `Begin` reports one. +DEFAULT_TEXT_MAX_BUFFER_CHARS = 8000 # Audio frames queued for pulling before the oldest is dropped: about 160 s # at the API's ~80 ms per frame. Bounds memory when audio is not read, or is # read more slowly than it arrives. @@ -62,6 +64,21 @@ def append(self, item: PullItem) -> bool: return True return False + def purge_before(self, cancel_id: int) -> None: + """Removes queued audio frames belonging to cancelled requests.""" + if not self._items: + return + new_items: Deque[PullItem] = deque() + audio_count = 0 + for item in self._items: + if item[0] == "audio" and item[1] is not None and item[1] <= cancel_id: + continue + if item[0] == "audio": + audio_count += 1 + new_items.append(item) + self._items = new_items + self._audio = audio_count + def popleft(self) -> Optional[PullItem]: """Removes and returns the oldest item; `None` when there is none.""" @@ -88,6 +105,21 @@ def max_generate_chars(begin: Any) -> int: return MAX_GENERATE_CHARS +def max_buffer_chars(begin: Any) -> int: + """ + The unspoken characters the API holds: the limit the session's + `Begin` reports (`limits.text_max_buffer_chars`, where `0` means + none), else `DEFAULT_TEXT_MAX_BUFFER_CHARS`. + """ + + configuration = getattr(begin, "configuration", None) or {} + limits = configuration.get("limits") + limit = limits.get("text_max_buffer_chars") if isinstance(limits, dict) else None + if isinstance(limit, int) and not isinstance(limit, bool) and limit >= 0: + return limit + return DEFAULT_TEXT_MAX_BUFFER_CHARS + + def check_text_frame(text: object, limit: int = MAX_GENERATE_CHARS) -> None: """Refuses what one `Generate` cannot carry; a `limit` of `0` is none.""" diff --git a/assemblyai/tts/v1/async_session.py b/assemblyai/tts/v1/async_session.py index d0b399f..e783b62 100644 --- a/assemblyai/tts/v1/async_session.py +++ b/assemblyai/tts/v1/async_session.py @@ -1,6 +1,7 @@ import asyncio import json import logging +import time from typing import Any, AsyncIterator, Callable, Dict, List, Optional, Union from urllib.parse import urlsplit @@ -14,6 +15,7 @@ check_session_config, check_text_frame, handshake_error, + max_buffer_chars, max_generate_chars, resolve_settings, ) @@ -114,6 +116,8 @@ def __init__( self._next_flush_id = 0 self._cancel_ids: List[int] = [] self._cancelled_before = -1 + self._done_through = -1 + self._last_heard_at = 0.0 self._drop_warned = False self.begin: Optional[TTSBeginEvent] = None self.termination: Optional[TTSTerminationEvent] = None @@ -152,6 +156,8 @@ async def connect(self) -> TTSBeginEvent: self._next_flush_id = 0 self._cancel_ids = [] self._cancelled_before = -1 + self._done_through = -1 + self._last_heard_at = 0.0 self._drop_warned = False self.begin = None self.termination = None @@ -203,12 +209,18 @@ async def terminate(self) -> Optional[TTSTerminationEvent]: await self._send({"type": "Terminate"}) except Exception: pass - try: - await asyncio.wait_for( - self._stopped.wait(), timeout=self._terminate_timeout - ) - except asyncio.TimeoutError: - pass + timeout = self._terminate_timeout + while not self._stopped.is_set(): + quiet = time.monotonic() - self._last_heard_at + remaining = timeout - quiet + if remaining <= 0: + break + try: + await asyncio.wait_for( + self._stopped.wait(), timeout=min(remaining, 0.25) + ) + except asyncio.TimeoutError: + pass finally: await self.close() return self.termination @@ -272,6 +284,8 @@ async def cancel(self) -> None: self._next_flush_id += 1 self._cancel_ids.append(cancel_id) self._cancelled_before = max(self._cancelled_before, cancel_id) + if self._queue is not None: + self._queue.purge_before(self._cancelled_before) await self._send({"type": "Cancel"}) async def keep_alive(self) -> None: @@ -287,6 +301,13 @@ async def speak(self, text: str) -> AsyncIterator[bytes]: self._check_not_in_handler("speak()") check_text_frame(text, 0) + backlog = max_buffer_chars(self.begin) + if backlog and len(text) > backlog: + raise ValueError( + f"text is {len(text)} characters, over the {backlog} unspoken characters " + "the API holds before it ends the session; split it and speak() each part " + "once the previous part's audio is read, or pace it with generate()" + ) limit = max_generate_chars(self.begin) or max(len(text), 1) for start in range(0, len(text), limit): await self.generate(text[start : start + limit]) @@ -294,7 +315,7 @@ async def speak(self, text: str) -> AsyncIterator[bytes]: async for chunk in self.audio(flush_id): yield chunk - async def audio(self, flush_id: Optional[int] = None) -> AsyncIterator[bytes]: + def audio(self, flush_id: Optional[int] = None) -> AsyncIterator[bytes]: """ Yields the audio of request `flush_id`; see `TTSSession.audio`. @@ -309,12 +330,19 @@ async def audio(self, flush_id: Optional[int] = None) -> AsyncIterator[bytes]: flush_id = self._next_flush_id # the request being written now elif flush_id > self._next_flush_id: raise ValueError(f"flush_id {flush_id} has not been flushed yet") + if flush_id <= self._done_through: + raise ValueError( + f"flush_id {flush_id} was already read: read each request once, in order" + ) if TTSEvents.Audio.value in self._handlers: raise RuntimeError( "audio goes to the registered Audio handler; pull with audio() " "or speak() only without one" ) connection = self._queue # this connection's; a reconnect makes a new one + return self._iter_audio(flush_id, connection) + + async def _iter_audio(self, flush_id: int, connection: Any) -> AsyncIterator[bytes]: while True: if self._queue is not connection: return @@ -334,6 +362,7 @@ async def audio(self, flush_id: Optional[int] = None) -> AsyncIterator[bytes]: if ident == flush_id: yield payload elif kind == "done": + self._done_through = max(self._done_through, ident) if ident == flush_id: return elif kind == "cancelled": @@ -478,6 +507,7 @@ async def _read_frames(self) -> None: return async def _dispatch(self, data: Dict[str, Any]) -> None: + self._last_heard_at = time.monotonic() kind = data.get("type") if kind == "Error": await self._fail(error_from_frame(data)) @@ -489,6 +519,9 @@ async def _dispatch(self, data: Dict[str, Any]) -> None: self.begin = event self._begun.set() elif kind == "Audio": + # Discard in-flight audio frames belonging to cancelled requests + if event.flush_id <= self._cancelled_before: + return if TTSEvents.Audio.value not in self._handlers: self._put(("audio", event.flush_id, event.audio)) elif kind == "FlushDone": diff --git a/assemblyai/tts/v1/session.py b/assemblyai/tts/v1/session.py index 03adb54..0f30bf1 100644 --- a/assemblyai/tts/v1/session.py +++ b/assemblyai/tts/v1/session.py @@ -1,6 +1,7 @@ import json import logging import threading +import time from typing import Any, Callable, Dict, Iterator, List, Optional, Union from urllib.parse import urlsplit @@ -14,6 +15,7 @@ check_session_config, check_text_frame, handshake_error, + max_buffer_chars, max_generate_chars, resolve_settings, ) @@ -173,6 +175,8 @@ def _reset(self) -> None: self._cancel_ids: List[int] = [] # Requests up to this flush_id were dropped by a Cancel (-1: none). self._cancelled_before = -1 + self._done_through = -1 + self._last_heard_at = 0.0 self._drop_warned = False self.begin = None self.termination = None @@ -270,7 +274,14 @@ def terminate(self) -> Optional[TTSTerminationEvent]: self._send({"type": "Terminate"}) except Exception: # already closed pass - self._stopped.wait(self._terminate_timeout) + timeout = self._terminate_timeout + while not self._stopped.is_set(): + quiet = time.monotonic() - self._last_heard_at + remaining = timeout - quiet + if remaining <= 0: + break + if self._stopped.wait(timeout=min(remaining, 0.25)): + break self.close() return self.termination @@ -347,6 +358,8 @@ def cancel(self) -> None: self._next_flush_id += 1 self._cancel_ids.append(cancel_id) self._cancelled_before = max(self._cancelled_before, cancel_id) + with self._queue_ready: + self._queue.purge_before(self._cancelled_before) self._send({"type": "Cancel"}) def keep_alive(self) -> None: @@ -369,11 +382,19 @@ def speak(self, text: str) -> Iterator[bytes]: is cancelled. Raises: + ValueError: for text over the unspoken backlog the API holds. TTSError: from the iteration, if the session fails. RuntimeError: from a session handler. """ self._check_not_in_handler("speak()") check_text_frame(text, 0) + backlog = max_buffer_chars(self.begin) + if backlog and len(text) > backlog: + raise ValueError( + f"text is {len(text)} characters, over the {backlog} unspoken characters " + "the API holds before it ends the session; split it and speak() each part " + "once the previous part's audio is read, or pace it with generate()" + ) limit = max_generate_chars(self.begin) or max(len(text), 1) for start in range(0, len(text), limit): self.generate(text[start : start + limit]) @@ -391,7 +412,7 @@ def audio(self, flush_id: Optional[int] = None) -> Iterator[bytes]: Raises: TTSError: if the session fails before the request completes. - ValueError: for a `flush_id` no `flush()` has returned. + ValueError: for a `flush_id` no `flush()` has returned or one already read. RuntimeError: when an `Audio` handler receives the audio instead, before the session first connects, or from a session handler. """ @@ -402,12 +423,19 @@ def audio(self, flush_id: Optional[int] = None) -> Iterator[bytes]: flush_id = self._next_flush_id # the request being written now elif flush_id > self._next_flush_id: raise ValueError(f"flush_id {flush_id} has not been flushed yet") + if flush_id <= self._done_through: + raise ValueError( + f"flush_id {flush_id} was already read: read each request once, in order" + ) if TTSEvents.Audio.value in self._handlers: raise RuntimeError( "audio goes to the registered Audio handler; pull with audio() " "or speak() only without one" ) connection = self._queue # this connection's; a reconnect makes a new one + return self._iter_audio(flush_id, connection) + + def _iter_audio(self, flush_id: int, connection: Any) -> Iterator[bytes]: while True: if self._queue is not connection: return @@ -427,6 +455,7 @@ def audio(self, flush_id: Optional[int] = None) -> Iterator[bytes]: if ident == flush_id: yield payload elif kind == "done": + self._done_through = max(self._done_through, ident) if ident == flush_id: return elif kind == "cancelled": @@ -574,6 +603,7 @@ def _fail(self, error: TTSError) -> None: self._emit(TTSEvents.Error, error) def _dispatch(self, data: Dict[str, Any]) -> None: + self._last_heard_at = time.monotonic() kind = data.get("type") if kind == "Error": self._fail(error_from_frame(data)) @@ -585,6 +615,9 @@ def _dispatch(self, data: Dict[str, Any]) -> None: self.begin = event self._begun.set() elif kind == "Audio": + # Discard in-flight audio frames belonging to cancelled requests + if event.flush_id <= self._cancelled_before: + return if TTSEvents.Audio.value not in self._handlers: self._put(("audio", event.flush_id, event.audio)) elif kind == "FlushDone": diff --git a/tests/unit/test_tts_session.py b/tests/unit/test_tts_session.py index bbb804a..b820a99 100644 --- a/tests/unit/test_tts_session.py +++ b/tests/unit/test_tts_session.py @@ -1073,3 +1073,52 @@ def test_a_session_over_ws_passes_no_tls_context(mocker, monkeypatch): [(uri, kwargs)] = calls assert uri.startswith("ws://localhost:8080/v1/ws?") assert not {"ssl", "ssl_context"} & kwargs.keys() + + +def test_audio_binds_request_id_eagerly_at_call_time(server): + with aai.TTSSession() as session: + reader = session.audio() + session.generate("First request.") + flush_0 = session.flush() + assert flush_0 == 0 + + session.generate("Second request.") + flush_1 = session.flush() + assert flush_1 == 1 + + assert b"".join(reader) == b"First request." + assert b"".join(session.audio(flush_1)) == b"Second request." + + +def test_rereading_a_finished_flush_id_raises(server): + with aai.TTSSession() as session: + assert b"".join(session.speak("First.")) == b"First." + with pytest.raises(ValueError, match="already read"): + session.audio(0) + + +def test_speak_refuses_text_over_backlog_limit(server): + with aai.TTSSession() as session: + with pytest.raises(ValueError, match="over the 8000 unspoken characters"): + session.speak("x" * 8001) + + +def test_speak_splits_text_over_generate_frame_limit(server): + with aai.TTSSession() as session: + long_text = "a" * 1500 + "b" * 1000 + assert b"".join(session.speak(long_text)) == long_text.encode() + generates = [f for f in server.sent if f["type"] == "Generate"] + assert len(generates) == 2 + assert len(generates[0]["text"]) == 2000 + assert len(generates[1]["text"]) == 500 + + +def test_cancel_purges_queued_audio_and_discards_in_flight_audio(server): + with aai.TTSSession() as session: + session.generate("Discard this.") + session.flush() + time.sleep(0.05) + session.cancel() + # Verify in-flight audio for cancelled request is discarded by _dispatch + session._dispatch({"type": "Audio", "audio": "AAAA", "flush_id": 0}) + assert b"".join(session.speak("Clean.")) == b"Clean." diff --git a/tests/unit/test_tts_session_async.py b/tests/unit/test_tts_session_async.py index e2da409..5efa54d 100644 --- a/tests/unit/test_tts_session_async.py +++ b/tests/unit/test_tts_session_async.py @@ -700,6 +700,50 @@ async def test_cancel_ends_a_reader_following_the_request_being_written(server): assert b"".join([c async for c in session.speak("Next.")]) == b"Next." +async def test_audio_binds_request_id_eagerly_at_call_time(server): + async with aai.AsyncTTSSession() as session: + reader = session.audio() + await session.generate("First async request.") + flush_0 = await session.flush() + assert flush_0 == 0 + + await session.generate("Second async request.") + flush_1 = await session.flush() + assert flush_1 == 1 + + assert b"".join([c async for c in reader]) == b"First async request." + assert ( + b"".join([c async for c in session.audio(flush_1)]) + == b"Second async request." + ) + + +async def test_rereading_a_finished_flush_id_raises(server): + async with aai.AsyncTTSSession() as session: + assert b"".join([c async for c in session.speak("First.")]) == b"First." + with pytest.raises(ValueError, match="already read"): + session.audio(0) + + +async def test_speak_refuses_text_over_backlog_limit(server): + async with aai.AsyncTTSSession() as session: + with pytest.raises(ValueError, match="over the 8000 unspoken characters"): + async for _ in session.speak("x" * 8001): + pass + + +async def test_speak_splits_text_over_generate_frame_limit(server): + async with aai.AsyncTTSSession() as session: + long_text = "a" * 1500 + "b" * 1000 + assert ( + b"".join([c async for c in session.speak(long_text)]) == long_text.encode() + ) + generates = [f for f in server.sent if f["type"] == "Generate"] + assert len(generates) == 2 + assert len(generates[0]["text"]) == 2000 + assert len(generates[1]["text"]) == 500 + + async def _collect(iterator): return [chunk async for chunk in iterator]