fix(tts): eager audio binding, dynamic terminate wait, and cancel audio purge

assemblyai/__version__.py

@@ -1 +1 @@
-__version__ = "1.8.2"
+__version__ = "1.8.3"

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."""
 

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 ... import _audio
@@ -13,6 +14,7 @@
     check_session_config,
     check_text_frame,
     handshake_error,
+    max_buffer_chars,
     max_generate_chars,
     resolve_settings,
 )
@@ -102,6 +104,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
@@ -140,6 +144,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
@@ -191,12 +197,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
@@ -260,6 +272,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:
@@ -275,14 +289,21 @@ 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])
         flush_id = await self.flush()
         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`.
 
@@ -297,12 +318,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
@@ -322,6 +350,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":
@@ -466,6 +495,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))
@@ -477,6 +507,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":

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 ... import _audio
@@ -13,6 +14,7 @@
     check_session_config,
     check_text_frame,
     handshake_error,
+    max_buffer_chars,
     max_generate_chars,
     resolve_settings,
 )
@@ -163,6 +165,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
@@ -260,7 +264,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
 
@@ -337,6 +348,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:
@@ -359,11 +372,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])
@@ -381,7 +402,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.
         """
@@ -392,12 +413,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
@@ -417,6 +445,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":
@@ -564,6 +593,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))
@@ -575,6 +605,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":

tests/unit/test_tts_session.py

@@ -1017,3 +1017,52 @@ def test_cancel_ends_a_reader_following_the_request_being_written(server):
 def test_the_first_request_is_not_dropped_by_the_initial_watermark(server):
     with aai.TTSSession() as session:
         assert b"".join(session.speak("Zero.")) == b"Zero."
+
+
+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."