diff --git a/CLAUDE.md b/CLAUDE.md index a5c99eb..ede6f2a 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -54,9 +54,11 @@ aai.settings.api_key = "your-key" - `aai.DictationResponse` — Dictation result: `.text`, `.words` (`DictationWord` with `text` + `confidence`), `.confidence`, `.llm_response`, `.llm_error`, `.audio_duration_ms`, `.session_id`, `.request_time_ms`, `.sync_time_ms`, and `.final_text` (the LLM rewrite, falling back to `.text`) - `aai.TextToSpeech` — Text-to-Speech HTTP API: one text per request, audio out. Methods: `synthesize()` (complete audio as a `TTSResponse`), `stream()` (a `TTSStream` context manager yielding bytes as the API generates them), `close()` / `with` for a client built from `api_key=`. Host from `aai.settings.tts_base_url` (`https://tts.us.assemblyai.com`; set `https://tts.eu.assemblyai.com` for the EU zone) - `aai.AsyncTextToSpeech` — Asyncio counterpart of `TextToSpeech`; `synthesize()` is a coroutine and `stream()` returns an `AsyncTTSStream`. Owns an HTTP pool: use `async with` or `await aclose()` +- `aai.TTSSession` — Text-to-Speech WebSocket session: one connection, text in as it is written, audio out. `with` connects and waits for `Begin` (`session.begin`); `speak(text)` yields one request's audio; `generate(text)` / `flush()` -> flush_id / `audio(flush_id)` for streamed-in text and pipelining (`audio()` with no id follows the request being written, from another thread); `cancel()` for barge-in (ends the iteration at once); `keep_alive()`; `on(TTSEvents.X, handler(event))` for `Begin`, `Audio` (switches audio from pull to push), `FlushDone`, `WordBoundaries`, `Cancelled`, `Termination`, `Error`. Exit sends `Terminate` and sets `session.termination` (`None` after a failure). Refusals and failures raise `TTSError` (`status_code` None, `request_id` = session id). Raw PCM/G.711 only; 2,000 chars per generate unless `Begin` reports another `limits.max_generate_text_length`; ~100 chars/s sustained. `token=` for temporary tokens +- `aai.AsyncTTSSession` — Asyncio counterpart: `async with`, coroutine sends, `async for chunk in session.speak(text)` - `aai.TTSConfig` — Text-to-Speech options: `voice_id` (a library voice such as `aai-jane`, or a cloned voice by UUID or slug), `model_id` (`assemblyai-tts-1`), `language`, `sample_rate`, `encoding` (`aai.TTSEncoding`: `pcm_s16le` default | `pcm_mulaw` | `pcm_alaw` | `mp3`), `container` (`aai.TTSContainer`: `raw` default | `wav`), `bit_rate`, `extend_biometric_consent` (owner-only consent renewal of a cloned voice). Rejects unknown fields - `aai.TTSResponse` — Text-to-Speech result: `.audio` (bytes) plus the response headers: `.request_id`, `.region`, `.model`, `.voice_id`, `.language`, `.encoding`, `.sample_rate`, `.container`, `.bit_rate`, `.content_type`, `.voice_expires_at`, `.billed_characters`, `.audio_duration_ms`. `TTSStream.info` carries the same fields without `audio`; `audio_duration_ms` is `None` on a stream with audio (`0` when the text had nothing to speak) -- `aai.TTSError` — Raised on any Text-to-Speech refusal: `.status_code`, `.error_code` (stable API code; `biometric_terms_not_accepted` 403 until the account accepts the biometric data terms in the dashboard, `auth_unavailable` 503 retryable), `.param`, `.request_id`, `.retryable`, `.retry_after`, `.errors`, `.body` (the whole envelope) +- `aai.TTSError` — Raised on any Text-to-Speech refusal: `.status_code`, `.error_code` (stable API code; `biometric_terms_not_accepted` 403 until the account accepts the biometric data terms in the dashboard, `auth_unavailable` 503 retryable; on a session these arrive as close 3006 and 1011, with `status_code` None, the close code in `.body["close_code"]` and `.retryable` from the API's Error frame), `.param`, `.request_id`, `.retryable`, `.retry_after`, `.errors`, `.body` (the whole envelope) - `assemblyai.streaming.v3.RealTimeTranscriber` — Real-time streaming with event-based API (threaded) - `assemblyai.streaming.v3.AsyncRealTimeTranscriber` — Asyncio-native counterpart; same options/events - `aai.LLMGateway` — LLM Gateway client. Resources: `models.list()`, `chat.completions.create()` (incl. `stream=True`), `understanding.create()`/`.validate()` @@ -479,6 +481,13 @@ with tts.stream("Hello.", telephony) as speech: # A cloned or library voice by id tts.synthesize("Hello.", aai.TTSConfig(voice_id="aai-alba")) + +# A WebSocket session for a voice agent: no handshake per sentence, barge-in with cancel() +with aai.TTSSession(config=aai.TTSConfig(voice_id="aai-jane")) as session: + for sentence in llm_reply_sentences(): + for chunk in session.speak(sentence): + play(chunk) +# session.termination is the summary (None if the session failed); raw PCM/G.711 only, no mp3, wav or bit_rate ``` ## LLM Gateway diff --git a/README.md b/README.md index f6168a3..c90f9cf 100644 --- a/README.md +++ b/README.md @@ -1013,6 +1013,68 @@ asyncio.run(main()) +
+ Stream text in and audio out over a WebSocket session (voice agents) + +`aai.TTSSession` keeps one connection open and speaks text as it is written, so an LLM's reply can be voiced sentence by sentence with no per-request handshake. `speak()` sends one text and yields its audio; `generate()` / `flush()` do the same in two steps, so tokens can be sent as they arrive; `cancel()` is the barge-in. Leaving the `with` block ends the session cleanly and `session.termination` holds the summary (`None` if the session failed instead). + +A session streams raw `pcm_s16le` (the default), `pcm_mulaw` or `pcm_alaw` at the config's `sample_rate`; `mp3`, a `wav` container and `bit_rate` are HTTP-only and are refused. Chunk sizes vary and are not aligned to frames, so re-frame for a fixed-size sink (20 ms of 8 kHz mu-law is 160 bytes). `word_boundaries=True` and `inactivity_timeout=` are arguments of the session, not of `TTSConfig`; word boundaries are off unless asked for. A session that fails raises `TTSError` wherever the app meets it first: the `audio()` iteration, the next send, or, for an app that pushes audio to a handler and never reads, the end of the `with` block. + +```python +import assemblyai as aai + +aai.settings.api_key = "" + +with aai.TTSSession(config=aai.TTSConfig(voice_id="aai-jane")) as session: + for sentence in llm_reply_sentences(): + for chunk in session.speak(sentence): # raw 24 kHz 16-bit mono PCM + play(chunk) + +print("spoken:", session.termination) # the session's summary +``` + +Streaming the text in as it is generated, with a barge-in: + +```python +with aai.TTSSession(config=aai.TTSConfig(voice_id="aai-jane")) as session: + for token in llm_tokens(): + session.generate(token) # joined exactly as sent: include the spaces + request = session.flush() # the request's flush_id; audio may start before this + + for chunk in session.audio(request): + if user_started_talking(): + session.cancel() # ends this iteration at once; the server drops the rest + break + play(chunk) +``` + +Audio starts before `flush()` once whole sentences have arrived. To play it while the text is still being sent, read `session.audio()` with no id from another thread (or task): it follows the request being written, and ends at that request's `FlushDone`. Requests can be pipelined: call `flush()` again before reading the previous request's audio, then read each `audio(flush_id)` in order. Every `flush()` and every `cancel()` takes the next `flush_id`, so ids skip one after a cancel. + +Limits: a `Generate` frame takes at most 2,000 characters by default (`session.begin.configuration["limits"]["max_generate_text_length"]` has the limit in force; split longer text across calls), and the API accepts text at about 100 characters per second after an initial 1,500, holding at most 8,000 unspoken characters; sending far ahead of the speech ends the session with `input_rate_exceeded`, so pace text to the audio. The session ends after an hour, or after `inactivity_timeout` seconds (5 to 120, default 120) with no frame either way; send `keep_alive()` between sparse requests. Audio the SDK has already received does not hold the session open: the API counts idle time from the last frame it sent or received, so playing buffered audio, or waiting for the user, counts as idle. Send `keep_alive()` at least once per idle window during both. `session.begin` holds the `Begin` frame with the session id and the server's configuration and limits. + +Register handlers for the other frames with `session.on(event, handler)`, where `handler(event)` runs on the session's reader thread: `TTSEvents.Begin`, `TTSEvents.FlushDone`, `TTSEvents.WordBoundaries`, `TTSEvents.Cancelled`, `TTSEvents.Termination` and `TTSEvents.Error` (each delivers the matching `TTS...Event`, or the `TTSError`). Registering a `TTSEvents.Audio` handler switches audio from the pull iterators to the handler; `speak()` and `audio()` then raise. A handler may send (`generate()`, `flush()`, `cancel()`, `keep_alive()`, `close()`) but must not wait on the session: the reader that runs the handler is what `speak()`, `audio()` and `terminate()` wait for, so they raise `RuntimeError` there. Hand the work to another thread (or task) instead. Word timings (`word_boundaries=True`) arrive once per request after its `FlushDone`, usually after its `audio()` iteration has ended; their `start`/`end` count from the start of the session's audio, so subtract the event's `audio_start_ms` for a time within the request. + +A refusal of the session raises `TTSError` from the `with`; a failure mid-session raises it from the next `audio()` iteration or send, and reaches the `Error` handler. On a session `status_code` is `None` (there is no HTTP status), `request_id` is the session id and `body["close_code"]` is the WebSocket close code (`biometric_terms_not_accepted` is a `3006`, `auth_unavailable` a retryable `1011`); `retryable=True` means a new session can succeed, such as after an `inactivity_timeout`. + +Once the API has ended a session (an `inactivity_timeout`, say), the same `TTSSession` object opens a new one with `with session:` or `session.connect()`. + +For a session opened on a user's device, pass a temporary token instead of the key: `aai.TTSSession(token=token)` (mint it server-side on `https://token.assemblyai.com/v1/token?product=tts&expires_in_seconds=60`, which answers `{"token": ..., "expires_in_seconds": 60}`). The API spends the token when the session begins, so a token opens one session: to reconnect, mint a new token and build a new `TTSSession` with it. + +The asyncio twin is `aai.AsyncTTSSession`: + +```python +async with aai.AsyncTTSSession(config=aai.TTSConfig(voice_id="aai-jane")) as session: + async for chunk in session.speak("Hello there."): + await play(chunk) + await session.generate("Streamed ") + await session.generate("in.") + request = await session.flush() + async for chunk in session.audio(request): + await play(chunk) +``` + +
+
Renew a cloned voice's consent @@ -2102,7 +2164,7 @@ All custom SDK exceptions inherit from `aai.AssemblyAIError`, except the streami | `aai.AssemblyAIError` | Base exception for all AssemblyAI errors. | `status_code: Optional[int]` | | `aai.TranscriptError` | Raised when a prerecorded transcript API call fails at the HTTP level or polling times out. | `status_code: Optional[int]` | | `aai.SyncTranscriptError` | Raised when synchronous transcription fails (`aai.SyncTranscriber`). | `status_code: Optional[int]`, `error_code: Optional[str]`, `retry_after: Optional[int]` | -| `aai.TTSError` | Raised when a Text-to-Speech request fails (`aai.TextToSpeech`). | `status_code: Optional[int]`, `error_code: Optional[str]`, `param: Optional[str]`, `request_id: Optional[str]`, `retryable: Optional[bool]`, `retry_after: Optional[int]`, `errors: Optional[List[Dict[str, Optional[str]]]]`, `body: Optional[dict]` | +| `aai.TTSError` | Raised when a Text-to-Speech request fails (`aai.TextToSpeech`) or a session is refused or fails (`aai.TTSSession`, where `status_code` is `None`, `request_id` is the session id and `body["close_code"]` is the WebSocket close code). | `status_code: Optional[int]`, `error_code: Optional[str]`, `param: Optional[str]`, `request_id: Optional[str]`, `retryable: Optional[bool]`, `retry_after: Optional[int]`, `errors: Optional[List[Dict[str, Optional[str]]]]`, `body: Optional[dict]` | | `assemblyai.types.RedactedAudioIncompleteError` | Raised when requesting redacted audio before processing completes. | `status_code: Optional[int]` | | `assemblyai.types.RedactedAudioExpiredError` | Raised when requesting redacted audio that has expired and is no longer available. | `status_code: Optional[int]` | | `assemblyai.types.RedactedAudioUnavailableError` | Raised when requesting redacted audio that is unavailable at the given URL. | `status_code: Optional[int]` | diff --git a/assemblyai/__init__.py b/assemblyai/__init__.py index f625c97..a765353 100644 --- a/assemblyai/__init__.py +++ b/assemblyai/__init__.py @@ -31,15 +31,25 @@ from .transcriber import Transcriber, Transcript, TranscriptGroup from .tts.v1 import ( AsyncTextToSpeech, + AsyncTTSSession, AsyncTTSStream, TextToSpeech, + TTSAudioEvent, + TTSBeginEvent, + TTSCancelledEvent, TTSConfig, TTSContainer, TTSEncoding, TTSError, + TTSEvents, + TTSFlushDoneEvent, TTSResponse, TTSResponseInfo, + TTSSession, TTSStream, + TTSTerminationEvent, + TTSWord, + TTSWordBoundariesEvent, ) from .types import ( AssemblyAIError, @@ -140,6 +150,7 @@ def __getattr__(name: str): "AsyncDictationTranscriber", "AsyncLLMGateway", "AsyncSyncTranscriber", + "AsyncTTSSession", "AsyncTTSStream", "AsyncTextToSpeech", "AsyncTranscriber", @@ -211,13 +222,22 @@ def __getattr__(name: str): "SyncTranscriptResponse", "SyncWord", "Timestamp", + "TTSAudioEvent", + "TTSBeginEvent", + "TTSCancelledEvent", "TTSConfig", "TTSContainer", "TTSEncoding", "TTSError", + "TTSEvents", + "TTSFlushDoneEvent", "TTSResponse", "TTSResponseInfo", + "TTSSession", "TTSStream", + "TTSTerminationEvent", + "TTSWord", + "TTSWordBoundariesEvent", "TextToSpeech", "Transcriber", "TranscriptionConfig", diff --git a/assemblyai/__version__.py b/assemblyai/__version__.py index 14d9d2f..29654ee 100644 --- a/assemblyai/__version__.py +++ b/assemblyai/__version__.py @@ -1 +1 @@ -__version__ = "1.7.0" +__version__ = "1.8.0" diff --git a/assemblyai/_audio.py b/assemblyai/_audio.py index e0bab83..3728c11 100644 --- a/assemblyai/_audio.py +++ b/assemblyai/_audio.py @@ -164,7 +164,10 @@ def _config_to_json(config: Any, *, exclude: Tuple[str, ...] = ()) -> Optional[d Returns: the JSON-ready dict, or `None` when nothing is set. """ - data = config.dict(exclude_none=True) + if hasattr(config, "model_dump"): + data = config.model_dump(exclude_none=True) + else: + data = config.dict(exclude_none=True) for field in exclude: data.pop(field, None) return data or None diff --git a/assemblyai/tts/__init__.py b/assemblyai/tts/__init__.py index 4666a76..7c8d64b 100644 --- a/assemblyai/tts/__init__.py +++ b/assemblyai/tts/__init__.py @@ -7,21 +7,41 @@ from .v1 import ( AsyncTextToSpeech, + AsyncTTSSession, AsyncTTSStream, TextToSpeech, + TTSAudioEvent, + TTSBeginEvent, + TTSCancelledEvent, TTSConfig, TTSContainer, TTSEncoding, TTSError, + TTSEvents, + TTSFlushDoneEvent, TTSResponse, TTSResponseInfo, + TTSSession, TTSStream, + TTSTerminationEvent, + TTSWord, + TTSWordBoundariesEvent, ) __all__ = [ + "AsyncTTSSession", "AsyncTTSStream", "AsyncTextToSpeech", + "TTSAudioEvent", + "TTSBeginEvent", + "TTSCancelledEvent", "TTSConfig", + "TTSEvents", + "TTSFlushDoneEvent", + "TTSSession", + "TTSTerminationEvent", + "TTSWord", + "TTSWordBoundariesEvent", "TTSContainer", "TTSEncoding", "TTSError", diff --git a/assemblyai/tts/_base_ws.py b/assemblyai/tts/_base_ws.py new file mode 100644 index 0000000..83e34fb --- /dev/null +++ b/assemblyai/tts/_base_ws.py @@ -0,0 +1,203 @@ +"""Helpers shared by the sync and async Text-to-Speech sessions.""" + +import asyncio +import platform +import ssl +from collections import deque +from typing import Any, Deque, Dict, Optional, Tuple +from urllib.parse import urlencode, urlsplit + +from .. import _audio +from ..__version__ import __version__ +from ..types import Settings +from .v1.api import _RETRYABLE_STATUSES +from .v1.models import TTSConfig, TTSContainer, TTSError + +# The route of the WebSocket API, after the host and any path prefix of +# `settings.tts_base_url`. +WS_PATH = "/v1/ws" +# Characters a single `Generate` frame may carry (Unicode code points) when +# `Begin` does not report the limit in force. +MAX_GENERATE_CHARS = 2000 +# 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. +PULL_QUEUE_FRAMES = 2000 + +# An item `audio()` reads: (kind, flush_id or cancel id, payload). `audio` +# items carry a request's audio; the markers `done` and `cancelled` end a +# request, and `error` and `closed` end the session. +PullItem = Tuple[str, Any, Any] + + +class PullBuffer: + """ + The items `audio()` reads, in arrival order. + + Markers are always kept, since an `audio()` iteration ends only at one. + Past `max_audio` audio frames, the oldest audio frame is dropped. + """ + + def __init__(self, max_audio: int) -> None: + self.max_audio = max_audio + self._items: Deque[PullItem] = deque() + self._audio = 0 + + def __len__(self) -> int: + return len(self._items) + + def append(self, item: PullItem) -> bool: + """Adds `item`; returns whether an older audio frame was dropped.""" + + self._items.append(item) + if item[0] != "audio": + return False + self._audio += 1 + if self._audio <= self.max_audio: + return False + for index, queued in enumerate(self._items): + if queued[0] == "audio": + del self._items[index] + self._audio -= 1 + return True + return False + + def popleft(self) -> Optional[PullItem]: + """Removes and returns the oldest item; `None` when there is none.""" + + if not self._items: + return None + item = self._items.popleft() + if item[0] == "audio": + self._audio -= 1 + return item + + +def max_generate_chars(begin: Any) -> int: + """ + The characters one `Generate` may carry: the limit the session's + `Begin` reports (`limits.max_generate_text_length`, where `0` means + none), else `MAX_GENERATE_CHARS`. + """ + + configuration = getattr(begin, "configuration", None) or {} + limits = configuration.get("limits") + limit = limits.get("max_generate_text_length") if isinstance(limits, dict) else None + if isinstance(limit, int) and not isinstance(limit, bool) and limit >= 0: + return limit + return MAX_GENERATE_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.""" + + if not isinstance(text, str): + raise TypeError(f"text must be a str, not {type(text).__name__}") + if limit and len(text) > limit: + raise ValueError( + f"text is {len(text)} characters, over the {limit}-character " + "limit of one Generate frame; split it across several generate() calls" + ) + + +def resolve_settings( + owner: str, base: Optional[Settings], api_key: Optional[str], token: Optional[str] +) -> Settings: + """ + The settings a session connects with: a copy, so nothing shared changes. + + `base` is a client's settings, else the global `assemblyai.settings`; + `api_key` replaces the copy's key. Without a key or a `token`, raises + the SDK's usual message. No HTTP client is built: a session needs none. + """ + + if base is None: + from .. import settings as _global + + base = _global + settings = base.model_copy() if hasattr(base, "model_copy") else base.copy() + if api_key is not None: + settings.api_key = api_key + if token is None and not settings.api_key: + raise ValueError( + "Please provide an API key: set the ASSEMBLYAI_API_KEY environment " + f"variable, set aai.settings.api_key, or pass api_key= or token= to {owner}." + ) + return settings + + +def handshake_error(exc: BaseException) -> TTSError: + """ + The `TTSError` for a WebSocket handshake that failed with `exc`. + + An HTTP answer is retryable by its status, as on the HTTP API; with no + answer, a refused or dropped connection and a timeout are retryable. + """ + + # The legacy client's InvalidStatusCode carries `status_code`; the + # current clients' InvalidStatus carries the `response`. + status = getattr(exc, "status_code", None) + if status is None: + status = getattr(getattr(exc, "response", None), "status_code", None) + if isinstance(status, int): + retryable = status in _RETRYABLE_STATUSES + else: + from websockets.exceptions import InvalidMessage # a dropped handshake + + status = None + retryable = isinstance( + exc, (OSError, asyncio.TimeoutError, InvalidMessage) + ) and not isinstance(exc, ssl.CertificateError) + return TTSError( + f"WebSocket handshake failed: {exc}", status_code=status, retryable=retryable + ) + + +def _user_agent() -> str: + return ( + f"AssemblyAI/1.0 (sdk=Python/{__version__} " + f"runtime_env=Python/{platform.python_version()})" + ) + + +def check_session_config(config: TTSConfig) -> None: + """Refuses the options a session cannot honour: it streams raw audio.""" + + if config.container not in (None, TTSContainer.raw) or config.bit_rate is not None: + raise ValueError( + "a wav container and bit_rate do not apply to a session, which streams " + "raw pcm_s16le, pcm_mulaw or pcm_alaw; leave them unset in its TTSConfig" + ) + + +def build_uri( + base_url: str, + config: TTSConfig, + *, + token: Optional[str], + api_key: Optional[str], + word_boundaries: bool, + inactivity_timeout: Optional[int], +) -> Tuple[str, Dict[str, Optional[str]]]: + """The session URI and handshake headers from the settings and config.""" + + check_session_config(config) + parts = urlsplit(base_url if "://" in base_url else f"https://{base_url}") + scheme = "ws" if parts.scheme in ("http", "ws") else "wss" + params: Dict[str, Any] = ( + _audio._config_to_json(config, exclude=("container", "bit_rate")) or {} + ) + params["word_boundaries"] = word_boundaries + if inactivity_timeout is not None: + params["inactivity_timeout"] = inactivity_timeout + if token is not None: + params["token"] = token + # The API reads booleans case-insensitively; send the canonical form. + query = urlencode( + {k: (str(v).lower() if isinstance(v, bool) else v) for k, v in params.items()} + ) + headers: Dict[str, Optional[str]] = {"User-Agent": _user_agent()} + if token is None and api_key: + headers["Authorization"] = api_key + prefix = parts.path.rstrip("/") + return f"{scheme}://{parts.netloc}{prefix}{WS_PATH}?{query}", headers diff --git a/assemblyai/tts/v1/__init__.py b/assemblyai/tts/v1/__init__.py index fd8224e..ea39654 100644 --- a/assemblyai/tts/v1/__init__.py +++ b/assemblyai/tts/v1/__init__.py @@ -1,6 +1,7 @@ from .api import TTSStream from .async_api import AsyncTTSStream from .async_client import AsyncTextToSpeech +from .async_session import AsyncTTSSession from .client import TextToSpeech from .models import ( TTSConfig, @@ -10,11 +11,32 @@ TTSResponse, TTSResponseInfo, ) +from .session import TTSSession +from .ws_models import ( + TTSAudioEvent, + TTSBeginEvent, + TTSCancelledEvent, + TTSEvents, + TTSFlushDoneEvent, + TTSTerminationEvent, + TTSWord, + TTSWordBoundariesEvent, +) __all__ = [ + "AsyncTTSSession", "AsyncTTSStream", "AsyncTextToSpeech", + "TTSAudioEvent", + "TTSBeginEvent", + "TTSCancelledEvent", "TTSConfig", + "TTSEvents", + "TTSFlushDoneEvent", + "TTSSession", + "TTSTerminationEvent", + "TTSWord", + "TTSWordBoundariesEvent", "TTSContainer", "TTSEncoding", "TTSError", diff --git a/assemblyai/tts/v1/async_session.py b/assemblyai/tts/v1/async_session.py new file mode 100644 index 0000000..e241f57 --- /dev/null +++ b/assemblyai/tts/v1/async_session.py @@ -0,0 +1,515 @@ +import asyncio +import json +import logging +from typing import Any, AsyncIterator, Callable, Dict, List, Optional, Union + +from ... import _audio +from ... import async_client as _async_client +from .._base_ws import ( + PULL_QUEUE_FRAMES, + PullBuffer, + PullItem, + build_uri, + check_session_config, + check_text_frame, + handshake_error, + max_generate_chars, + resolve_settings, +) +from .models import TTSConfig, TTSError +from .ws_models import ( + TTSBeginEvent, + TTSEvents, + TTSTerminationEvent, + close_code_and_reason, + error_from_close, + error_from_frame, + parse_event, +) + +__all__ = ["AsyncTTSSession"] + +logger = logging.getLogger(__name__) + +_Handler = Callable[[Any], Any] + + +def websocket_connect_async( + uri: str, additional_headers: Dict[str, Optional[str]] +) -> Any: + """Opens an asyncio websocket; the streaming client's seam, imported + lazily so `import assemblyai` does not load websockets. Tests patch this.""" + + from ...streaming.v3.async_client import websocket_connect_async as _connect + + return _connect(uri, additional_headers=additional_headers) + + +class AsyncTTSSession: + """ + A Text-to-Speech WebSocket session for asyncio: the counterpart of + `TTSSession`. + + Open it with `async with`, which connects and waits for `Begin`. + `generate()`, `flush()`, `cancel()` and `keep_alive()` are coroutines; + `speak()` and `audio()` are async iterators of audio bytes. Handlers + registered with `on()` may be plain functions or coroutines; they run on + the session's reader task, so, as with `TTSSession`, they may send but + not wait on the session (`speak()`, `audio()`, `terminate()`). + + Example: + ```python + async with aai.AsyncTTSSession(config=aai.TTSConfig(voice_id="aai-jane")) as session: + async for chunk in session.speak("Hello there."): + await play(chunk) + ``` + """ + + def __init__( + self, + config: Optional[TTSConfig] = None, + *, + client: Optional[_async_client.AsyncClient] = None, + api_key: Optional[str] = None, + token: Optional[str] = None, + word_boundaries: bool = False, + inactivity_timeout: Optional[int] = None, + connect_timeout: float = 10.0, + terminate_timeout: float = 10.0, + ) -> None: + """Creates a session; see `TTSSession.__init__` for the arguments.""" + + _audio.check_config(type(self).__name__, config, TTSConfig) + self.config = config or TTSConfig() + check_session_config(self.config) + self._settings = resolve_settings( + type(self).__name__, client.settings if client else None, api_key, token + ) + self._token = token + self._word_boundaries = word_boundaries + self._inactivity_timeout = inactivity_timeout + self._connect_timeout = connect_timeout + self._terminate_timeout = terminate_timeout + + self._handlers: Dict[str, List[_Handler]] = {} + self._ws: Any = None + # Created in `connect()`, so they bind to the running loop. + self._reader: Optional["asyncio.Task[None]"] = None + self._begun: Any = None + self._stopped: Any = None + self._queue: Any = None + self._queue_ready: Any = None + self._next_flush_id = 0 + self._cancel_ids: List[int] = [] + self._cancelled_before = 0 + self._drop_warned = False + self.begin: Optional[TTSBeginEvent] = None + self.termination: Optional[TTSTerminationEvent] = None + self.error: Optional[TTSError] = None + + def on(self, event: Union[TTSEvents, str], handler: _Handler) -> None: + """Registers `handler(event)`, a function or coroutine function.""" + + name = str(event) + if name not in TTSEvents.__members__: + raise ValueError( + f"unknown event {name!r}; one of {', '.join(TTSEvents.__members__)}" + ) + self._handlers.setdefault(name, []).append(handler) + + async def connect(self) -> TTSBeginEvent: + """Opens the socket and waits for `Begin`; see `TTSSession.connect`.""" + + self._check_not_in_handler("connect()") + if self._ws is not None: + if not self._stopped.is_set(): + raise RuntimeError("the session is already connected") + await self.close() # the previous connection has ended + reader = self._reader + if reader is not None and not reader.done(): + # Its reader would write into the new connection's state. + await asyncio.wait({reader}, timeout=self._connect_timeout) + if not reader.done(): + raise RuntimeError( + "a handler of the previous connection is still running" + ) + self._begun = asyncio.Event() + self._stopped = asyncio.Event() + self._queue = PullBuffer(PULL_QUEUE_FRAMES) + self._queue_ready = asyncio.Event() + self._next_flush_id = 0 + self._cancel_ids = [] + self._cancelled_before = 0 + self._drop_warned = False + self.begin = None + self.termination = None + self.error = None + self._error_raised = False + + uri, headers = build_uri( + self._settings.tts_base_url, + self.config, + token=self._token, + api_key=self._settings.api_key, + word_boundaries=self._word_boundaries, + inactivity_timeout=self._inactivity_timeout, + ) + try: + self._ws = await asyncio.wait_for( + websocket_connect_async(uri, headers), timeout=self._connect_timeout + ) + except Exception as exc: # the handshake failed + await self.close() # ends any audio() waiting on the fresh state + raise handshake_error(exc) from exc + + self._reader = asyncio.create_task(self._read_loop()) + try: + await asyncio.wait_for(self._begun.wait(), timeout=self._connect_timeout) + except asyncio.TimeoutError: + await self.close() + raise TTSError("The session did not begin in time") from None + except BaseException: # cancelled: leave no socket or reader behind + await self.close() + raise + if self.error is not None: + await self.close() + raise self._surfaced() + if self.begin is None: + await self.close() + raise TTSError("The session closed before it began") + return self.begin + + async def terminate(self) -> Optional[TTSTerminationEvent]: + """Ends the session cleanly; see `TTSSession.terminate`.""" + + self._check_not_in_handler("terminate()") + if self._ws is None: + return None + try: + if not self._stopped.is_set(): + try: + await self._send({"type": "Terminate"}) + except Exception: + pass + try: + await asyncio.wait_for( + self._stopped.wait(), timeout=self._terminate_timeout + ) + except asyncio.TimeoutError: + pass + finally: + await self.close() + return self.termination + + async def close(self) -> None: + """Closes the socket without `Terminate`.""" + + ws, self._ws = self._ws, None + if self._stopped is not None: + self._stopped.set() + if self._begun is not None: + self._begun.set() + if self._queue is not None: + self._put(("closed", None, None)) + reader = self._reader + if reader is asyncio.current_task(): + reader = None + try: + if ws is not None: + try: + await ws.close() + except Exception: + pass + if reader is not None and not reader.done(): + await asyncio.wait({reader}, timeout=self._terminate_timeout) + finally: + # Also when the caller is cancelled, which then propagates. + if reader is not None and not reader.done(): + reader.cancel() + + async def __aenter__(self) -> "AsyncTTSSession": + await self.connect() + return self + + async def __aexit__(self, exc_type: Any, *exc: object) -> None: + if exc_type is not None: + await self.close() + return + await self.terminate() + if self.error is not None and not self._error_raised: + raise self._surfaced() + + async def generate(self, text: str) -> None: + """Appends `text` to the current request; see `TTSSession.generate`.""" + + check_text_frame(text, max_generate_chars(self.begin)) + await self._send({"type": "Generate", "text": text}) + + async def flush(self) -> int: + """Concludes the current request and returns its `flush_id`.""" + + flush_id = self._next_flush_id + self._next_flush_id += 1 + await self._send({"type": "Flush"}) + return flush_id + + async def cancel(self) -> None: + """Barge-in; see `TTSSession.cancel`.""" + + cancel_id = self._next_flush_id + self._next_flush_id += 1 + self._cancel_ids.append(cancel_id) + self._cancelled_before = max(self._cancelled_before, cancel_id) + await self._send({"type": "Cancel"}) + + async def keep_alive(self) -> None: + await self._send({"type": "KeepAlive"}) + + async def speak(self, text: str) -> AsyncIterator[bytes]: + """ + Speaks one text and yields its audio; see `TTSSession.speak`. + + Unlike `TTSSession.speak()`, nothing is sent until iteration starts, + so pipeline requests with `generate()`, `flush()` and `audio()`. + """ + + self._check_not_in_handler("speak()") + check_text_frame(text, 0) + 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]: + """ + Yields the audio of request `flush_id`; see `TTSSession.audio`. + + Without `flush_id`, the request being written now, so a task can + play it while another still sends text. + """ + + self._check_not_in_handler("audio()") + if self._queue is None: + raise RuntimeError("the session is not open") + if flush_id is None: + 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 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 + while True: + if self._queue is not connection: + return + if flush_id < self._cancelled_before: # dropped by a Cancel + return + item = await self._get(timeout=0.5) + if item is None: + if self._stopped.is_set(): + if self.error is not None: + raise self._surfaced() + return + continue + kind, ident, payload = item + if kind == "audio": + if ident == flush_id: + yield payload + elif kind == "done": + if ident == flush_id: + return + elif kind == "cancelled": + if ident > flush_id: + return + elif kind == "error": + raise self._surfaced() + elif kind == "closed": + if self.error is not None: + raise self._surfaced() + return + + def _surfaced(self) -> TTSError: + """The session's error, marked as raised to the caller.""" + + if self.error is None: + raise RuntimeError("the session has no error") + self._error_raised = True + return self.error + + def _check_not_in_handler(self, method: str) -> None: + """Refuses a method that waits on the reader when the reader, running + a handler, is the caller: it would wait for itself.""" + + if self._reader is not None and asyncio.current_task() is self._reader: + raise RuntimeError( + f"{method} cannot be called from a session handler: it waits " + "for the session's reader, which runs the handlers" + ) + + def _put(self, item: PullItem) -> None: + """Queues an item for `audio()`; past the cap the oldest audio drops.""" + + dropped = self._queue.append(item) + self._queue_ready.set() + if dropped and not self._drop_warned: + self._drop_warned = True + logger.warning( + "AsyncTTSSession: over %d audio frames are waiting to be read, so the " + "oldest are being dropped; read audio() as fast as the audio " + "arrives, or register an Audio handler", + self._queue.max_audio, + ) + + async def _get(self, timeout: float) -> Optional[PullItem]: + """The next item for `audio()`; `None` when none came in `timeout` s.""" + + if not self._queue: + self._queue_ready.clear() + try: + await asyncio.wait_for(self._queue_ready.wait(), timeout) + except asyncio.TimeoutError: + return None + return self._queue.popleft() + + async def _fail(self, error: TTSError) -> None: + first = self.error is None + if first: + if error.request_id is None and self.begin is not None: + error.request_id = self.begin.id + self.error = error + # Waiters learn of the failure before any handler runs. + self._put(("error", None, self.error)) + self._stopped.set() + self._begun.set() + if first: + await self._emit(TTSEvents.Error, error) + + async def _send(self, frame: Dict[str, Any]) -> None: + from websockets.exceptions import ConnectionClosed + + ws = self._ws + if ws is None or self._stopped.is_set(): + if self.error is not None: + raise self._surfaced() + raise RuntimeError( + "the session is not open: use `with` / `async with`, or connect()" + ) + try: + await ws.send(json.dumps(frame)) + except ConnectionClosed as exc: + raise await self._closed_error(ws, exc) from exc + + async def _closed_error(self, ws: Any, exc: Exception) -> Exception: + """ + What a send on a socket that closed raises: the session's error, once + the reader has taken in the frames that arrived before the close. + """ + + if asyncio.current_task() is not self._reader: + try: + await asyncio.wait_for(self._stopped.wait(), timeout=1) + except asyncio.TimeoutError: + pass + if self.error is not None: + return self.error + code, reason = close_code_and_reason(exc) + if self._ws is not ws or code == 1000: # closed by close() or cleanly + return RuntimeError("the session is not open") + return error_from_close(code, reason) + + async def _read_loop(self) -> None: + try: + await self._read_frames() + except Exception as exc: # a failure in the reader itself + await self._fail(TTSError(f"The session's reader failed: {exc!r}")) + finally: + # However the reader ended, nothing may wait for it any longer. + if not self._stopped.is_set(): + self._stopped.set() + self._begun.set() + self._put(("closed", None, None)) + + async def _read_frames(self) -> None: + from websockets.exceptions import ConnectionClosed + + ws = self._ws + while not self._stopped.is_set() and ws is not None: + try: + raw = await ws.recv() + except ConnectionClosed as exc: + await self._on_closed(exc) + return + except asyncio.CancelledError: + raise + except Exception as exc: + if not self._stopped.is_set(): + await self._on_closed(exc) + return + if isinstance(raw, bytes): + continue + try: + data = json.loads(raw) + except ValueError: + logger.warning("ignoring a non-JSON frame") + continue + if isinstance(data, dict): + try: + await self._dispatch(data) + except Exception as exc: # a frame this SDK cannot read + await self._fail(TTSError(f"Unreadable frame from the API: {exc}")) + return + + async def _dispatch(self, data: Dict[str, Any]) -> None: + kind = data.get("type") + if kind == "Error": + await self._fail(error_from_frame(data)) + return + event = parse_event(data) + if event is None: + return + if kind == "Begin": + self.begin = event + self._begun.set() + elif kind == "Audio": + if TTSEvents.Audio.value not in self._handlers: + self._put(("audio", event.flush_id, event.audio)) + elif kind == "FlushDone": + self._put(("done", event.flush_id, None)) + elif kind == "Cancelled": + cancel_id = self._cancel_ids.pop(0) if self._cancel_ids else -1 + self._put(("cancelled", cancel_id, None)) + elif kind == "Termination": + self.termination = event + await self._emit(str(kind), event) + if kind == "Termination": + self._stopped.set() + self._put(("closed", None, None)) + + async def _on_closed(self, exc: Exception) -> None: + if self._stopped.is_set(): + return + code, reason = close_code_and_reason(exc) + error = None + if code != 1000 and self.error is None: + error = error_from_close(code, reason) + if self.begin is not None: + error.request_id = self.begin.id + self.error = error + self._stopped.set() + self._begun.set() + self._put(("closed", None, None)) + if error is not None: + await self._emit(TTSEvents.Error, error) + + async def _emit(self, event: Union[TTSEvents, str], payload: Any) -> None: + for handler in self._handlers.get(str(event), ()): + try: + result = handler(payload) + if asyncio.iscoroutine(result): + await result + except Exception: + logger.exception("AsyncTTSSession handler for %s raised", event) diff --git a/assemblyai/tts/v1/models.py b/assemblyai/tts/v1/models.py index bdebdbd..d56afa5 100644 --- a/assemblyai/tts/v1/models.py +++ b/assemblyai/tts/v1/models.py @@ -50,10 +50,17 @@ class TTSError(AssemblyAIError): body over its limit adds `limit`). `biometric_terms_not_accepted` (`403`, not retryable) refuses one of the - account's own cloned voices, or `extend_biometric_consent`, until the - account accepts the biometric data terms in the AssemblyAI dashboard. - `auth_unavailable` (`503`, retryable) means the credential, or those - terms, could not be checked just now. + account's own cloned voices, an inline voice clip or + `extend_biometric_consent`, until the account accepts the biometric data + terms in the AssemblyAI dashboard. `auth_unavailable` (`503`, retryable) + means the credential, or those terms, could not be checked just now. + + On a WebSocket session (`TTSSession`) there is no HTTP status: + `status_code` is `None`, `request_id` is the session id, `retryable` is + the API's (from its `Error` frame, else from the close code), and + `body["close_code"]` is the WebSocket close code. There + `biometric_terms_not_accepted` arrives as close `3006` and + `auth_unavailable` as close `1011`. A streamed response cut after its first audio raises a `TTSError` with `status_code` `200`, no `error_code`, `retryable` `True` and the diff --git a/assemblyai/tts/v1/session.py b/assemblyai/tts/v1/session.py new file mode 100644 index 0000000..1eabaa6 --- /dev/null +++ b/assemblyai/tts/v1/session.py @@ -0,0 +1,612 @@ +import json +import logging +import threading +from typing import Any, Callable, Dict, Iterator, List, Optional, Union + +from ... import _audio +from ... import client as _client +from .._base_ws import ( + PULL_QUEUE_FRAMES, + PullBuffer, + PullItem, + build_uri, + check_session_config, + check_text_frame, + handshake_error, + max_generate_chars, + resolve_settings, +) +from .models import TTSConfig, TTSError +from .ws_models import ( + TTSBeginEvent, + TTSEvents, + TTSTerminationEvent, + close_code_and_reason, + error_from_close, + error_from_frame, + parse_event, +) + +__all__ = ["TTSSession"] + +logger = logging.getLogger(__name__) + +_Handler = Callable[[Any], None] + + +def websocket_connect(uri: str, **kwargs: Any) -> Any: + """Opens a sync websocket; the streaming client's seam, imported lazily so + `import assemblyai` does not load websockets. Tests patch this name.""" + + import inspect + + from websockets.sync.client import connect as _ws_connect + + from ...streaming.v3.client import websocket_connect as _connect + + # websockets 17 deprecates a bare connect() unless `legacy=True` (older + # releases have no such parameter). + if "legacy" in inspect.signature(_ws_connect).parameters: + kwargs["legacy"] = True + return _connect(uri, **kwargs) + + +class TTSSession: + """ + A Text-to-Speech WebSocket session: text in as it is written, audio out. + + Open it with `with`, which connects and waits for `Begin`. Send text with + `generate()` as it is produced, conclude each request with `flush()`, and + read its audio with `audio()`; `speak()` does all three for one text. + `cancel()` is the barge-in: it stops the request in flight and drops + whatever is queued. Leaving the `with` block sends `Terminate` and waits + for the session's `Termination` summary. + + Audio is delivered one of two ways: pulled with `speak()` / `audio()`, or + pushed to a handler registered with `on(TTSEvents.Audio, ...)`. Register + handlers for the other events (`FlushDone`, `WordBoundaries`, + `Cancelled`, `Termination`, `Error`) as needed; a handler takes the event + and runs on the session's reader thread. A handler may send + (`generate()`, `flush()`, `cancel()`, `keep_alive()`, `close()`) but not + wait on the session: `speak()`, `audio()` and `terminate()` raise + `RuntimeError` there. + + Example: + ```python + import assemblyai as aai + + aai.settings.api_key = "your-key" + + with aai.TTSSession(config=aai.TTSConfig(voice_id="aai-jane")) as session: + for sentence in reply_sentences(): + for chunk in session.speak(sentence): + play(chunk) # 24 kHz s16le mono + ``` + """ + + def __init__( + self, + config: Optional[TTSConfig] = None, + *, + client: Optional[_client.Client] = None, + api_key: Optional[str] = None, + token: Optional[str] = None, + word_boundaries: bool = False, + inactivity_timeout: Optional[int] = None, + connect_timeout: float = 10.0, + terminate_timeout: float = 10.0, + ) -> None: + """ + Creates a session; nothing connects until `with` or `connect()`. + + Args: + config: The voice and audio options (`voice_id`, `model_id`, + `language`, `sample_rate`, `encoding`). A session streams raw + PCM or G.711: `container` may only be `raw`, and `bit_rate` + does not apply. + client: A `Client` whose settings (API key, `tts_base_url`) to + use. Defaults to the shared default client. + api_key: The API key to authenticate with; takes precedence over + `client` and the global settings. + token: A temporary token instead of an API key, for a session + opened on a user's device. The API spends the token when the + session begins, so such a session connects once: to + reconnect, mint a new token and build a new session. + word_boundaries: Whether the server sends `WordBoundaries` frames. + Off by default here (the API's own default is on); pass + `True` to receive them. + inactivity_timeout: Seconds (5 to 120) with no frame either way + after which the API closes the session; `None` keeps the + API's 120 s. + connect_timeout: Seconds to wait for the handshake and `Begin`. + terminate_timeout: Seconds to wait for `Termination` on exit. + + Raises: + TypeError: if `config` is not a `TTSConfig`. + ValueError: if `config` sets a `wav` container or `bit_rate`. + """ + _audio.check_config(type(self).__name__, config, TTSConfig) + self.config = config or TTSConfig() + check_session_config(self.config) + self._settings = resolve_settings( + type(self).__name__, client.settings if client else None, api_key, token + ) + self._token = token + self._word_boundaries = word_boundaries + self._inactivity_timeout = inactivity_timeout + self._connect_timeout = connect_timeout + self._terminate_timeout = terminate_timeout + + self._handlers: Dict[str, List[_Handler]] = {} + self._ws: Any = None + self._reader: Optional[threading.Thread] = None + # Held while a frame is written. flush() and cancel() hold it while + # they also take the frame's flush_id, so the ids follow wire order. + self._send_lock = threading.RLock() + self.begin: Optional[TTSBeginEvent] = None + "The `Begin` frame, once connected." + self.termination: Optional[TTSTerminationEvent] = None + """The `Termination` summary once the session ended cleanly; `None` + before that, after a failure, or when the wait for it timed out.""" + self.error: Optional[TTSError] = None + "The error that ended the session, if one did." + self._reset() + + def _reset(self) -> None: + """Fresh per-connection state, so a session object can be reused.""" + + self._begun = threading.Event() + self._stopped = threading.Event() + self._queue = PullBuffer(PULL_QUEUE_FRAMES) + self._queue_ready = threading.Condition() + self._next_flush_id = 0 + self._cancel_ids: List[int] = [] + # Requests with a lower flush_id than this were dropped by a Cancel. + self._cancelled_before = 0 + self._drop_warned = False + self.begin = None + self.termination = None + self.error = None + self._error_raised = False + + # -- lifecycle ----------------------------------------------------------- + + def on(self, event: Union[TTSEvents, str], handler: _Handler) -> None: + """Registers `handler(event)` for a server frame type.""" + + name = str(event) + if name not in TTSEvents.__members__: + raise ValueError( + f"unknown event {name!r}; one of {', '.join(TTSEvents.__members__)}" + ) + self._handlers.setdefault(name, []).append(handler) + + def connect(self) -> TTSBeginEvent: + """ + Opens the socket and waits for `Begin`. + + A session whose previous connection has ended (cleanly, by `close()` + or by the API, such as after an `inactivity_timeout`) connects again; + one built with `token=` cannot, since its token is already spent. + + Raises: + TTSError: if the API refuses the session (its `Error` frame), the + handshake is rejected, or `Begin` does not arrive in time. + RuntimeError: if the session is still connected, from a session + handler, or while a handler of the previous connection still + runs after `connect_timeout`. + """ + self._check_not_in_handler("connect()") + if self._ws is not None: + if not self._stopped.is_set(): + raise RuntimeError("the session is already connected") + self.close() # the previous connection has ended + if self._reader is not None and self._reader.is_alive(): + # Its reader would write into the new connection's state. + self._reader.join(self._connect_timeout) + if self._reader.is_alive(): + raise RuntimeError( + "a handler of the previous connection is still running" + ) + self._reset() + + uri, headers = build_uri( + self._settings.tts_base_url, + self.config, + token=self._token, + api_key=self._settings.api_key, + word_boundaries=self._word_boundaries, + inactivity_timeout=self._inactivity_timeout, + ) + try: + self._ws = websocket_connect( + uri, additional_headers=headers, open_timeout=self._connect_timeout + ) + except Exception as exc: # the handshake failed + self.close() # ends any audio() waiting on the fresh state + raise handshake_error(exc) from exc + + self._reader = threading.Thread( + target=self._read_loop, name="assemblyai-tts-session", daemon=True + ) + self._reader.start() + + if not self._begun.wait(self._connect_timeout): + self.close() + raise TTSError("The session did not begin in time") + if self.error is not None: + self.close() + raise self._surfaced() + if self.begin is None: + self.close() + raise TTSError("The session closed before it began") + return self.begin + + def terminate(self) -> Optional[TTSTerminationEvent]: + """ + Ends the session cleanly: `Terminate`, then wait for `Termination`. + + Text sent and not flushed is spoken first. Returns the summary, or + `None` if the session had already ended or the wait timed out. + + Raises: + RuntimeError: from a session handler. + """ + self._check_not_in_handler("terminate()") + if self._ws is None: + return None + if not self._stopped.is_set(): + try: + self._send({"type": "Terminate"}) + except Exception: # already closed + pass + self._stopped.wait(self._terminate_timeout) + self.close() + return self.termination + + def close(self) -> None: + """Closes the socket without `Terminate`; in-flight audio is dropped.""" + + ws, self._ws = self._ws, None + self._stopped.set() + self._begun.set() + if ws is not None: + try: + ws.close() + except Exception: # already closed + pass + self._put(("closed", None, None)) + if self._reader is not None and self._reader is not threading.current_thread(): + self._reader.join(timeout=self._terminate_timeout) + + def __enter__(self) -> "TTSSession": + self.connect() + return self + + def __exit__(self, exc_type: Any, *exc: object) -> None: + if exc_type is not None: + self.close() + return + self.terminate() + # A failure nothing in the block observed (push mode, or text sent + # and never read) must not pass silently. + if self.error is not None and not self._error_raised: + raise self._surfaced() + + # -- sending --------------------------------------------------------------- + + def generate(self, text: str) -> None: + """ + Appends `text` to the current request. + + Consecutive texts are joined exactly as sent, so include the spaces. + Audio may start before `flush()` once enough complete sentences have + arrived; `flush()` is what concludes the request. + + Raises: + ValueError: for a text over the limit of one `Generate` frame + (split it): the limit `Begin` reports, else 2,000 characters. + """ + check_text_frame(text, max_generate_chars(self.begin)) + self._send({"type": "Generate", "text": text}) + + def flush(self) -> int: + """ + Concludes the current request and returns its `flush_id`. + + Everything sent so far is spoken, then `FlushDone` arrives for this + id. Flushes can be pipelined; the audio comes in order. + """ + with self._send_lock: + flush_id = self._next_flush_id + self._next_flush_id += 1 + self._send({"type": "Flush"}) + return flush_id + + def cancel(self) -> None: + """ + Barge-in: stops the request in flight and drops queued requests. + + An `audio()` iteration of a cancelled request ends at once; audio + of it still in transit is discarded. The server answers `Cancelled`. + Text sent and not flushed is discarded. Like `flush()`, a cancel + takes a `flush_id`, so the next request's id skips one. + """ + with self._send_lock: + cancel_id = self._next_flush_id + self._next_flush_id += 1 + self._cancel_ids.append(cancel_id) + self._cancelled_before = max(self._cancelled_before, cancel_id) + self._send({"type": "Cancel"}) + + def keep_alive(self) -> None: + """ + Restarts the API's idle timer. Send one at least once per + `inactivity_timeout` while the session waits, which includes the + time already received audio takes to play. + """ + + self._send({"type": "KeepAlive"}) + + def speak(self, text: str) -> Iterator[bytes]: + """ + Speaks one text: `generate()`, `flush()`, then its audio. + + Text longer than one `Generate` frame is split across frames; very + long text (several thousand characters) should be paced instead, as + the API limits unspoken text. Returns an iterator of the request's + audio bytes that ends at its `FlushDone`, or early if the request + is cancelled. + + Raises: + TTSError: from the iteration, if the session fails. + RuntimeError: from a session handler. + """ + self._check_not_in_handler("speak()") + check_text_frame(text, 0) + 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]) + return self.audio(self.flush()) + + def audio(self, flush_id: Optional[int] = None) -> Iterator[bytes]: + """ + Yields the audio of request `flush_id` until its `FlushDone`. + + Without `flush_id`, the request being written now: a thread can read + it while `generate()` is still sending text, since audio starts + before `flush()`. Requests are served in order, so read them in + order. Audio is queued for pulling only while no `TTSEvents.Audio` + handler is registered. + + Raises: + TTSError: if the session fails before the request completes. + ValueError: for a `flush_id` no `flush()` has returned. + RuntimeError: when an `Audio` handler receives the audio instead, + before the session first connects, or from a session handler. + """ + self._check_not_in_handler("audio()") + if self._reader is None: + raise RuntimeError("the session is not open") + if flush_id is None: + 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 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 + while True: + if self._queue is not connection: + return + if flush_id < self._cancelled_before: # dropped by a Cancel + return + item = self._get(timeout=0.5) + if item is None: + if self._stopped.is_set(): + if self.error is not None: + raise self._surfaced() + return + continue + kind, ident, payload = item + if kind == "audio": + if ident == flush_id: + yield payload + elif kind == "done": + if ident == flush_id: + return + elif kind == "cancelled": + if ident > flush_id: # a Cancel sent after this request + return + elif kind == "error": + raise self._surfaced() + elif kind == "closed": + if self.error is not None: + raise self._surfaced() + return + + # -- internals --------------------------------------------------------------- + + def _surfaced(self) -> TTSError: + """The session's error, marked as raised to the caller.""" + + if self.error is None: + raise RuntimeError("the session has no error") + self._error_raised = True + return self.error + + def _check_not_in_handler(self, method: str) -> None: + """Refuses a method that waits on the reader when the reader, running + a handler, is the caller: it would wait for itself.""" + + if self._reader is not None and threading.current_thread() is self._reader: + raise RuntimeError( + f"{method} cannot be called from a session handler: it waits " + "for the session's reader, which runs the handlers" + ) + + def _put(self, item: PullItem) -> None: + """Queues an item for `audio()`; past the cap the oldest audio drops.""" + + with self._queue_ready: + dropped = self._queue.append(item) + self._queue_ready.notify_all() + if dropped and not self._drop_warned: + self._drop_warned = True + logger.warning( + "TTSSession: over %d audio frames are waiting to be read, so the " + "oldest are being dropped; read audio() as fast as the audio " + "arrives, or register an Audio handler", + self._queue.max_audio, + ) + + def _get(self, timeout: float) -> Optional[PullItem]: + """The next item for `audio()`; `None` when none came in `timeout` s.""" + + with self._queue_ready: + if not self._queue: + self._queue_ready.wait(timeout) + return self._queue.popleft() + + def _send(self, frame: Dict[str, Any]) -> None: + from websockets.exceptions import ConnectionClosed + + with self._send_lock: + ws = self._ws + if ws is None or self._stopped.is_set(): + if self.error is not None: + raise self._surfaced() + raise RuntimeError( + "the session is not open: use `with` / `async with`, or connect()" + ) + try: + ws.send(json.dumps(frame)) + return + except ConnectionClosed as exc: + closed = exc + # Outside the lock: a handler that sends must not block the reader. + raise self._closed_error(ws, closed) from closed + + def _closed_error(self, ws: Any, exc: Exception) -> Exception: + """ + What a send on a socket that closed raises: the session's error, once + the reader has taken in the frames that arrived before the close. + """ + + if threading.current_thread() is not self._reader: + self._stopped.wait(1) + if self.error is not None: + return self.error + code, reason = close_code_and_reason(exc) + if self._ws is not ws or code == 1000: # closed by close() or cleanly + return RuntimeError("the session is not open") + return error_from_close(code, reason) + + def _read_loop(self) -> None: + try: + self._read_frames() + except Exception as exc: # a failure in the reader itself + self._fail(TTSError(f"The session's reader failed: {exc!r}")) + finally: + # However the reader ended, nothing may wait for it any longer. + if not self._stopped.is_set(): + self._stopped.set() + self._begun.set() + self._put(("closed", None, None)) + + def _read_frames(self) -> None: + from websockets.exceptions import ConnectionClosed + + ws = self._ws + while not self._stopped.is_set() and ws is not None: + try: + raw = ws.recv(timeout=1) + except TimeoutError: + continue + except ConnectionClosed as exc: + self._on_closed(exc) + return + except Exception as exc: # the socket was closed under us + if not self._stopped.is_set(): + self._on_closed(exc) + return + if isinstance(raw, bytes): + continue + try: + data = json.loads(raw) + except ValueError: + logger.warning("ignoring a non-JSON frame") + continue + if isinstance(data, dict): + try: + self._dispatch(data) + except Exception as exc: # a frame this SDK cannot read + self._fail(TTSError(f"Unreadable frame from the API: {exc}")) + return + + def _fail(self, error: TTSError) -> None: + """Ends the session on an error: handlers, waiters, state.""" + + first = self.error is None + if first: + if error.request_id is None and self.begin is not None: + error.request_id = self.begin.id + self.error = error + # Waiters learn of the failure before any handler runs. + self._put(("error", None, self.error)) + self._stopped.set() + self._begun.set() + if first: + self._emit(TTSEvents.Error, error) + + def _dispatch(self, data: Dict[str, Any]) -> None: + kind = data.get("type") + if kind == "Error": + self._fail(error_from_frame(data)) + return + event = parse_event(data) + if event is None: + return + if kind == "Begin": + self.begin = event + self._begun.set() + elif kind == "Audio": + if TTSEvents.Audio.value not in self._handlers: + self._put(("audio", event.flush_id, event.audio)) + elif kind == "FlushDone": + self._put(("done", event.flush_id, None)) + elif kind == "Cancelled": + # cancel() already raised the watermark for this id. + cancel_id = self._cancel_ids.pop(0) if self._cancel_ids else -1 + self._put(("cancelled", cancel_id, None)) + elif kind == "Termination": + self.termination = event + self._emit(str(kind), event) + if kind == "Termination": + self._stopped.set() + self._put(("closed", None, None)) + + def _on_closed(self, exc: Exception) -> None: + if self._stopped.is_set(): + return + code, reason = close_code_and_reason(exc) + error = None + if code != 1000 and self.error is None: + error = error_from_close(code, reason) + if self.begin is not None: + error.request_id = self.begin.id + self.error = error + self._stopped.set() + self._begun.set() + self._put(("closed", None, None)) + if error is not None: + self._emit(TTSEvents.Error, error) + + def _emit(self, event: Union[TTSEvents, str], payload: Any) -> None: + for handler in self._handlers.get(str(event), ()): + try: + handler(payload) + except Exception: + logger.exception("TTSSession handler for %s raised", event) diff --git a/assemblyai/tts/v1/ws_models.py b/assemblyai/tts/v1/ws_models.py new file mode 100644 index 0000000..27c3ebc --- /dev/null +++ b/assemblyai/tts/v1/ws_models.py @@ -0,0 +1,218 @@ +"""Models of the Text-to-Speech WebSocket API: the frames a session receives. + +The client frames (`Generate`, `Flush`, `Cancel`, `KeepAlive`, `Terminate`) are +plain dicts; the server frames are the events below, dispatched to handlers +registered with `TTSSession.on()`. +""" + +import base64 +from enum import Enum +from typing import Any, Dict, List, Optional, Tuple, Type, TypeVar + +try: + # pydantic v2 import (`model_validator` exists only in v2) + from pydantic import BaseModel, Field, model_validator # noqa: F401 + + pydantic_v2 = True +except ImportError: + # pydantic v1 import (fallback for Python < 3.14) + from pydantic import BaseModel, Field + + pydantic_v2 = False + +from .models import TTSError + +__all__ = [ + "TTSAudioEvent", + "TTSBeginEvent", + "TTSCancelledEvent", + "TTSEvents", + "TTSFlushDoneEvent", + "TTSTerminationEvent", + "TTSWord", + "TTSWordBoundariesEvent", +] + +# Close codes after which the same session parameters can succeed on a new +# connection, for a close that came without an `Error` frame: a lost +# connection, a transient check failure, a server restart, a server-side +# failure, an expired session and the account's session limit. +_RETRYABLE_CLOSE_CODES = frozenset({1006, 1011, 1012, 3005, 3008, 3009}) +# The code of a connection that ended without a close frame. +_ABNORMAL_CLOSURE = 1006 + +_M = TypeVar("_M", bound=BaseModel) + + +def _parse(model: Type[_M], data: Dict[str, Any]) -> _M: + if hasattr(model, "model_validate"): + return model.model_validate(data) + return model.parse_obj(data) + + +class TTSEvents(str, Enum): + """The server frames of a session, as `TTSSession.on()` names them.""" + + Begin = "Begin" + "Handler receives a `TTSBeginEvent`." + Audio = "Audio" + "Handler receives a `TTSAudioEvent`; registering one turns off pulling." + FlushDone = "FlushDone" + "Handler receives a `TTSFlushDoneEvent`." + WordBoundaries = "WordBoundaries" + "Handler receives a `TTSWordBoundariesEvent`, after the request's `FlushDone`." + Cancelled = "Cancelled" + "Handler receives a `TTSCancelledEvent`." + Termination = "Termination" + "Handler receives a `TTSTerminationEvent`." + Error = "Error" + "Handler receives the `TTSError` that ended the session." + + def __str__(self): + return self.value + + +class TTSBeginEvent(BaseModel): + """The first frame of a session: it is open and configured.""" + + id: str + "The session id. Quote it in support requests." + + expires_at: int + "Unix time in seconds at which the session's maximum duration ends." + + configuration: Dict[str, Any] = Field(default_factory=dict) + """The session as created: the effective connect parameters, `model`, + `voice_id`, `session_ttl_sec`, `limits` and more. An open set; empty + when the API sends none.""" + + +class TTSAudioEvent(BaseModel): + """A piece of a request's audio.""" + + flush_id: int + "The request the audio belongs to." + + audio: bytes = Field(..., repr=False) + "Raw codec bytes (no header), decoded from the frame's base64." + + +class TTSFlushDoneEvent(BaseModel): + """A request's audio has all been delivered.""" + + flush_id: int + audio_duration_ms: int + "The playback duration of the request's audio; `0` for nothing to speak." + + +class TTSWord(BaseModel): + word: str + "One whitespace-delimited token of the request's text, verbatim." + start: int + """Milliseconds on the session's whole audio timeline, not the request's: + subtract the event's `audio_start_ms` for a time within the request.""" + end: int + confidence: float + + +class TTSWordBoundariesEvent(BaseModel): + """ + Word timings of a request: best effort, at most one per request. + + Arrives after the request's `FlushDone`, so usually after its `audio()` + iteration has ended, and possibly after the next request's first audio; + match it by `flush_id`. `start` and `end` count from the start of the + session's audio; `start - audio_start_ms` is the time within the request. + """ + + flush_id: int + audio_start_ms: int + "Where the request's audio starts on the session's timeline." + words: List[TTSWord] + + +class TTSCancelledEvent(BaseModel): + """The answer to a `Cancel`.""" + + audio_duration_seconds: float + "Audio streamed for the cancelled request before it stopped." + + +class TTSTerminationEvent(BaseModel): + """The session's summary, after `Terminate`.""" + + session_duration_seconds: int + total_input_char_length: int + "The billed characters." + audio_duration_seconds: float + voice_expires_at: Optional[int] = None + + +def parse_event(data: Dict[str, Any]) -> Any: + """The event model for a server frame, or `None` for an unknown type.""" + + kind = data.get("type") + if kind == "Audio": + return TTSAudioEvent( + flush_id=data["flush_id"], audio=base64.b64decode(data["audio"]) + ) + if kind == "Begin": + return _parse( + TTSBeginEvent, {**data, "configuration": data.get("configuration") or {}} + ) + if kind == "FlushDone": + return _parse(TTSFlushDoneEvent, data) + if kind == "WordBoundaries": + return _parse(TTSWordBoundariesEvent, data) + if kind == "Cancelled": + return _parse(TTSCancelledEvent, data) + if kind == "Termination": + return _parse(TTSTerminationEvent, data) + return None + + +def error_from_frame(data: Dict[str, Any]) -> TTSError: + """ + The `TTSError` for an `Error` frame. `body` keeps the frame, plus + `close_code`: the code of the close that follows it (the frame's + `error_code`), under the same key as for a close without a frame. + """ + + retry_after = data.get("retry_after_seconds") + body = dict(data) + body.setdefault("close_code", data.get("error_code")) + return TTSError( + data.get("error") or "text-to-speech session failed", + error_code=data.get("error_type"), + param=data.get("param"), + retryable=data.get("retryable"), + retry_after=int(retry_after) if isinstance(retry_after, (int, float)) else None, + body=body, + ) + + +def close_code_and_reason(exc: BaseException) -> Tuple[int, str]: + """ + The code and reason of the close frame the API sent, from the exception + that ended the connection; `1006` and no reason when none arrived. + """ + + rcvd = getattr(exc, "rcvd", None) + if rcvd is None: + return _ABNORMAL_CLOSURE, "" + return int(rcvd.code), rcvd.reason or "" + + +def error_from_close(code: Optional[int], reason: str) -> TTSError: + """The `TTSError` for a connection closed without a clean `Termination`.""" + + if code is None: + code = _ABNORMAL_CLOSURE + text = f"Connection closed (code={code})" + if reason: + text = f"{text}: {reason}" + return TTSError( + text, + retryable=code in _RETRYABLE_CLOSE_CODES, + body={"close_code": code, "reason": reason}, + ) diff --git a/tests/unit/test_tts_session.py b/tests/unit/test_tts_session.py new file mode 100644 index 0000000..b96608a --- /dev/null +++ b/tests/unit/test_tts_session.py @@ -0,0 +1,997 @@ +"""Tests for the threaded Text-to-Speech WebSocket session. + +`_FakeServer` plays the API: it answers `Generate`/`Flush` with the text's +bytes as one `Audio` frame and a `FlushDone`, `Cancel` with `Cancelled`, and +`Terminate` with `Termination` and a clean close. +""" + +import base64 +import json +import queue +import threading +import time +from types import SimpleNamespace +from typing import Any, Iterable, List, Optional +from urllib.parse import parse_qs, urlsplit + +import pytest +from websockets.exceptions import ConnectionClosedError, ConnectionClosedOK +from websockets.frames import Close + +import assemblyai as aai + +aai.settings.api_key = "test" + +_BEGIN = { + "type": "Begin", + "id": "sess-1", + "expires_at": 1789556741, + "configuration": { + "voice_id": "aai-jane", + "model": "assemblyai-tts-1", + "sample_rate": 24000, + }, +} +_TERMINATION = { + "type": "Termination", + "session_duration_seconds": 3, + "total_input_char_length": 20, + "audio_duration_seconds": 1.5, +} + + +_INACTIVITY_TIMEOUT = { + "type": "Error", + "error_code": 3006, + "error": "No message received for 120 seconds; the idle session was closed.", + "error_type": "inactivity_timeout", + "retryable": True, + "param": None, + "retry_after_seconds": None, +} +_INPUT_RATE_EXCEEDED = { + "type": "Error", + "error_code": 3010, + "error": "Text arrived faster than it can be spoken.", + "error_type": "input_rate_exceeded", + "retryable": False, + "param": None, + "retry_after_seconds": None, +} + + +class _Rejected(Exception): + """A handshake answered with an HTTP status, as the current clients raise.""" + + def __init__(self, status: int) -> None: + super().__init__(f"server rejected WebSocket connection: HTTP {status}") + self.response = SimpleNamespace(status_code=status) + + +class _LegacyRejected(Exception): + """The same, as the legacy client of websockets 11 and 12 raises it.""" + + def __init__(self, status: int) -> None: + super().__init__(f"server rejected WebSocket connection: HTTP {status}") + self.status_code = status + + +_HANDSHAKE_FAILURES = [ + pytest.param(_Rejected(503), 503, True, id="503"), + pytest.param(_Rejected(502), 502, True, id="502"), + pytest.param(_Rejected(401), 401, False, id="401"), + pytest.param(_LegacyRejected(429), 429, True, id="legacy-429"), + pytest.param(_LegacyRejected(403), 403, False, id="legacy-403"), + pytest.param( + ConnectionRefusedError(111, "Connection refused"), None, True, id="refused" + ), + pytest.param(TimeoutError("timed out"), None, True, id="timeout"), + pytest.param(ValueError("not a WebSocket URI"), None, False, id="bad-uri"), +] + + +def _closed_ok() -> ConnectionClosedOK: + return ConnectionClosedOK(Close(1000, "Session ended"), None) + + +class _FakeServer: + def __init__(self, begin: Optional[dict] = _BEGIN, answer_flush: bool = True): + self.sent: List[dict] = [] + self.inbox: "queue.Queue[Any]" = queue.Queue() + self.closed = False + self.text = "" + self.flush_id = 0 + self.answer_flush = answer_flush + self.generate_delay = 0.0 + self.closing: Optional[List[Any]] = None + if begin is not None: + self.push(begin) + + def close_on_next_send(self, code: int, error: Optional[dict] = None) -> None: + """The API closes the session: the next send raises, and the reader + has yet to take in the frames that came first.""" + self.closing = ([error] if error else []) + [ + ConnectionClosedError(Close(code, "See Error message for details"), None) + ] + + def push(self, item: Any) -> None: + self.inbox.put(json.dumps(item) if isinstance(item, dict) else item) + + def send(self, raw: str) -> None: + if self.closing is not None: + closing, self.closing = self.closing, None + for item in closing: + self.push(item) + raise ConnectionClosedError(closing[-1].rcvd, None) + frame = json.loads(raw) + self.sent.append(frame) + kind = frame["type"] + if kind == "Generate": + time.sleep(self.generate_delay) + self.text += frame["text"] + elif kind == "Flush": + if self.answer_flush: + audio = base64.b64encode(self.text.encode()).decode() + self.push({"type": "Audio", "audio": audio, "flush_id": self.flush_id}) + self.push( + { + "type": "FlushDone", + "flush_id": self.flush_id, + "audio_duration_ms": 80, + } + ) + self.flush_id += 1 + self.text = "" + elif kind == "Cancel": + self.flush_id += 1 + self.text = "" + self.push({"type": "Cancelled", "audio_duration_seconds": 0.1}) + elif kind == "Terminate": + self.push(_TERMINATION) + self.push(_closed_ok()) + + def recv(self, timeout: Optional[float] = None) -> str: + try: + item = self.inbox.get(timeout=0.02) + except queue.Empty: + raise TimeoutError() + if isinstance(item, Exception): + raise item + return item + + def close(self) -> None: + self.closed = True + + +@pytest.fixture +def server(mocker): + fake = _FakeServer() + connect = mocker.patch( + "assemblyai.tts.v1.session.websocket_connect", return_value=fake + ) + fake.connect = connect + return fake + + +def _drain(items: Iterable[Any], timeout: float = 5.0) -> List[Any]: + """Collects `items` on a thread, failing the test if it does not end.""" + + got: List[Any] = [] + raised: List[BaseException] = [] + + def collect() -> None: + try: + got.extend(items) + except BaseException as exc: + raised.append(exc) + + thread = threading.Thread(target=collect, daemon=True) + thread.start() + thread.join(timeout) + if thread.is_alive(): + pytest.fail(f"the iteration did not end within {timeout} s") + if raised: + raise raised[0] + return got + + +def _frames(server: _FakeServer, kind: str) -> List[dict]: + return [f for f in server.sent if f["type"] == kind] + + +def test_with_connects_waits_for_begin_and_terminates_on_exit(server): + with aai.TTSSession( + config=aai.TTSConfig(voice_id="aai-alba", sample_rate=8000) + ) as session: + assert session.begin is not None + assert session.begin.id == "sess-1" + assert session.begin.configuration["model"] == "assemblyai-tts-1" + + uri = server.connect.call_args.args[0] + headers = server.connect.call_args.kwargs["additional_headers"] + parts = urlsplit(uri) + assert parts.scheme == "wss" and parts.netloc == "tts.us.assemblyai.com" + assert parts.path == "/v1/ws" + assert parse_qs(parts.query) == { + "voice_id": ["aai-alba"], + "sample_rate": ["8000"], + "word_boundaries": ["false"], + } + assert headers["Authorization"] == "test" + assert "AssemblyAI" in headers["User-Agent"] + assert [f["type"] for f in server.sent] == ["Terminate"] + assert session.termination is not None + assert session.termination.total_input_char_length == 20 + assert server.closed + + +def test_speak_sends_generate_and_flush_and_yields_the_audio(server): + with aai.TTSSession() as session: + chunks = _drain(session.speak("Hello there.")) + + assert b"".join(chunks) == b"Hello there." + assert server.sent[:2] == [ + {"type": "Generate", "text": "Hello there."}, + {"type": "Flush"}, + ] + + +def test_requests_are_pipelined_and_read_in_flush_id_order(server): + with aai.TTSSession() as session: + session.generate("One. ") + first = session.flush() + session.generate("Two.") + second = session.flush() + assert (first, second) == (0, 1) + assert b"".join(_drain(session.audio(first))) == b"One. " + assert b"".join(_drain(session.audio(second))) == b"Two." + + +def test_cancel_ends_the_current_request_but_not_a_later_one(server): + server.answer_flush = False # the request stays in flight + with aai.TTSSession() as session: + session.generate("A long reply that gets interrupted.") + flush_id = session.flush() + stream = session.audio(flush_id) + session.cancel() # barge-in + assert _drain(stream) == [] + assert session._next_flush_id == 2 + + server.answer_flush = True + assert b"".join(_drain(session.speak("Next."))) == b"Next." + assert [f["type"] for f in server.sent] == [ + "Generate", + "Flush", + "Cancel", + "Generate", + "Flush", + "Terminate", + ] + + +def test_a_cancel_from_another_thread_keeps_flush_ids_in_wire_order(server): + server.generate_delay = 0.02 # a slow write the cancel lands in + with aai.TTSSession() as session: + for _ in range(5): + flush_ids: List[int] = [] + + def write_reply() -> None: + session.generate("Hello there. ") + flush_ids.append(session.flush()) + + writer = threading.Thread(target=write_reply) + writer.start() + time.sleep(0.01) + session.cancel() # barge-in from a VAD callback, say + writer.join(5) + _drain(session.audio(flush_ids[0])) # ends, cancelled or spoken + + # every Flush and Cancel took the id the API gives it + assert session._next_flush_id == server.flush_id == 10 + + +def test_cancel_ends_the_iteration_before_the_server_answers(server): + server.answer_flush = False + with aai.TTSSession() as session: + flush_id = session.flush() + server.push({"type": "Audio", "audio": "AA==", "flush_id": 0}) + stream = session.audio(flush_id) + assert next(stream) == b"\x00" + session.cancel() + # in-transit audio of the cancelled request is discarded, not yielded + server.push({"type": "Audio", "audio": "AA==", "flush_id": 0}) + assert _drain(stream) == [] + + +def test_audio_for_an_unflushed_id_raises(server): + with aai.TTSSession() as session: + with pytest.raises(ValueError, match="not been flushed"): + next(session.audio(5)) + + +def test_audio_without_an_id_follows_the_request_being_written(server): + with aai.TTSSession() as session: + got = [] + reader = threading.Thread(target=lambda: got.extend(session.audio())) + reader.start() # reading before the flush, like a player thread + session.generate("Early ") + session.generate("audio.") + session.flush() + reader.join(timeout=5) + assert b"".join(got) == b"Early audio." + + +def test_a_misspelled_event_name_is_refused(server): + session = aai.TTSSession() + with pytest.raises(ValueError, match="FlushDone"): + session.on("FlushDon", print) + session.on("FlushDone", print) # a string spelling of a member is fine + + +def test_sending_before_connect_says_how_to_open(server): + with pytest.raises(RuntimeError, match="with"): + aai.TTSSession().generate("Hi.") + + +def test_audio_before_connect_says_the_session_is_not_open(server): + with pytest.raises(RuntimeError, match="not open"): + _drain(aai.TTSSession().audio()) + + +@pytest.mark.parametrize( + "config", + [aai.TTSConfig(container="wav"), aai.TTSConfig(bit_rate=64000)], + ids=["wav", "bit_rate"], +) +def test_http_only_options_are_refused_when_the_session_is_built(server, config): + with pytest.raises(ValueError, match="wav container and bit_rate"): + aai.TTSSession(config=config) + assert not server.connect.called + + +@pytest.mark.parametrize( + "container", ["raw", aai.TTSContainer.raw], ids=["str", "enum"] +) +def test_a_raw_container_is_what_a_session_streams(server, container): + with aai.TTSSession(config=aai.TTSConfig(container=container)): + pass + assert "container" not in parse_qs(urlsplit(server.connect.call_args.args[0]).query) + + +def test_a_cancel_before_a_request_does_not_end_it(server): + with aai.TTSSession() as session: + session.cancel() # nothing in flight; Cancelled arrives later + assert b"".join(_drain(session.speak("Hi."))) == b"Hi." + + +def test_an_error_before_begin_raises_from_with(mocker): + fake = _FakeServer(begin=None) + fake.push( + { + "type": "Error", + "error_code": 3012, + "error": "The requested voice was not found.", + "error_type": "voice_not_found", + "retryable": False, + "param": "voice_id", + "retry_after_seconds": None, + } + ) + fake.push(ConnectionClosedError(Close(3012, "See Error message for details"), None)) + mocker.patch("assemblyai.tts.v1.session.websocket_connect", return_value=fake) + + with pytest.raises(aai.TTSError) as exc_info: + with aai.TTSSession(config=aai.TTSConfig(voice_id="aai-nobody")): + pass + + err = exc_info.value + assert err.error_code == "voice_not_found" + assert err.param == "voice_id" + assert err.retryable is False + assert err.body["error_code"] == 3012 + assert err.body["close_code"] == 3012 # the key a bare close uses too + assert err.request_id is None # no Begin, so no session id yet + assert fake.closed + + +@pytest.mark.parametrize("failure, status, retryable", _HANDSHAKE_FAILURES) +def test_a_failed_handshake_keeps_its_status_and_says_if_retryable( + mocker, failure, status, retryable +): + mocker.patch("assemblyai.tts.v1.session.websocket_connect", side_effect=failure) + with pytest.raises(aai.TTSError, match="handshake failed") as exc_info: + with aai.TTSSession(): + pass + assert exc_info.value.status_code == status + assert exc_info.value.retryable is retryable + assert exc_info.value.__cause__ is failure + + +def test_an_error_mid_session_raises_from_the_audio_and_reaches_the_handler(server): + errors = [] + with pytest.raises(aai.TTSError) as exc_info: + with aai.TTSSession() as session: + session.on(aai.TTSEvents.Error, errors.append) + server.answer_flush = False + flush_id = session.flush() + server.push( + { + "type": "Error", + "error_code": 3005, + "error": "No room", + "error_type": "at_capacity", + "retryable": True, + "param": None, + "retry_after_seconds": 2, + } + ) + _drain(session.audio(flush_id)) + + assert exc_info.value.error_code == "at_capacity" + assert exc_info.value.retry_after == 2 + assert exc_info.value.request_id == "sess-1" # the session id + assert errors == [exc_info.value] + assert session.error is exc_info.value + + +def test_a_lost_connection_is_a_retryable_error(server): + errors = [] + with pytest.raises(aai.TTSError) as exc_info: + with aai.TTSSession() as session: + session.on(aai.TTSEvents.Error, errors.append) + server.answer_flush = False + flush_id = session.flush() + server.push(ConnectionClosedError(None, None)) # abnormal: no close frame + _drain(session.audio(flush_id)) + + assert exc_info.value.retryable is True + assert exc_info.value.body == {"close_code": 1006, "reason": ""} + assert len(errors) == 1 + + +@pytest.mark.filterwarnings("error::DeprecationWarning") +@pytest.mark.parametrize( + "code, retryable", + [ + (1011, True), + (1012, True), + (3005, True), + (3008, True), + (3009, True), + (3006, False), + (3010, False), + ], +) +def test_a_close_without_an_error_frame_is_retryable_by_its_code( + server, code, retryable +): + with pytest.raises(aai.TTSError) as exc_info: + with aai.TTSSession() as session: + server.answer_flush = False + flush_id = session.flush() + server.push(ConnectionClosedError(Close(code, "Upstream failed"), None)) + _drain(session.audio(flush_id)) + + assert exc_info.value.retryable is retryable + assert exc_info.value.body == {"close_code": code, "reason": "Upstream failed"} + + +def test_a_socket_error_is_a_retryable_lost_connection(server): + with pytest.raises(aai.TTSError) as exc_info: + with aai.TTSSession() as session: + server.answer_flush = False + flush_id = session.flush() + server.push(OSError("Connection reset by peer")) + _drain(session.audio(flush_id)) + + assert exc_info.value.retryable is True + assert exc_info.value.body == {"close_code": 1006, "reason": ""} + + +def test_a_failure_in_the_reader_still_ends_the_iteration(server, mocker): + mocker.patch( + "assemblyai.tts.v1.session.error_from_close", side_effect=RuntimeError("bug") + ) + with pytest.raises(aai.TTSError, match="reader failed"): + with aai.TTSSession() as session: + server.answer_flush = False + flush_id = session.flush() + server.push(ConnectionClosedError(None, None)) + _drain(session.audio(flush_id)) + + +def test_a_send_after_a_failure_raises_the_stored_error(server): + with pytest.raises(aai.TTSError) as exc_info: + with aai.TTSSession() as session: + server.push(ConnectionClosedError(None, None)) + session._stopped.wait(2) + session.generate("too late") + assert exc_info.value.body["close_code"] == 1006 + + +def test_a_send_the_api_closed_raises_the_sessions_error(server): + with pytest.raises(aai.TTSError) as exc_info: + with aai.TTSSession() as session: + server.close_on_next_send(3010, _INPUT_RATE_EXCEEDED) + session.generate("Too much, too fast.") + + assert exc_info.value.error_code == "input_rate_exceeded" + assert exc_info.value.retryable is False + assert exc_info.value is session.error + + +def test_a_send_on_a_bare_close_raises_a_tts_error(server): + with pytest.raises(aai.TTSError) as exc_info: + with aai.TTSSession() as session: + server.close_on_next_send(1012) + session.flush() + + assert exc_info.value.retryable is True + assert exc_info.value.body["close_code"] == 1012 + + +def test_a_cancel_drops_every_pipelined_request(server): + server.answer_flush = False + with aai.TTSSession() as session: + first = session.flush() + second = session.flush() + session.cancel() + assert _drain(session.audio(first)) == [] + assert _drain(session.audio(second)) == [] # also dropped, must not hang + + +def test_an_unreadable_frame_ends_the_session_with_an_error(server): + with pytest.raises(aai.TTSError, match="Unreadable frame"): + with aai.TTSSession() as session: + server.answer_flush = False + flush_id = session.flush() + server.push({"type": "Audio", "flush_id": "zero"}) # no audio, wrong type + _drain(session.audio(flush_id)) + + +def test_a_session_object_can_be_reused(mocker): + fakes = [_FakeServer(), _FakeServer()] + mocker.patch("assemblyai.tts.v1.session.websocket_connect", side_effect=fakes) + session = aai.TTSSession() + with session: + assert b"".join(_drain(session.speak("One."))) == b"One." + with session: + assert b"".join(_drain(session.speak("Two."))) == b"Two." + assert [f["type"] for f in fakes[1].sent] == ["Generate", "Flush", "Terminate"] + + +def test_connect_after_the_api_ended_the_session_opens_a_new_one(mocker): + first, second = _FakeServer(), _FakeServer() + mocker.patch( + "assemblyai.tts.v1.session.websocket_connect", side_effect=[first, second] + ) + session = aai.TTSSession() + session.connect() + first.push(_INACTIVITY_TIMEOUT) + first.push( + ConnectionClosedError(Close(3006, "See Error message for details"), None) + ) + assert session._stopped.wait(5) + assert session.error is not None and session.error.retryable is True + + session.connect() # no close() needed first + assert first.closed + assert session.error is None + assert b"".join(_drain(session.speak("Hello again."))) == b"Hello again." + session.terminate() + + +def test_connect_on_a_live_session_is_refused(server): + with aai.TTSSession() as session: + with pytest.raises(RuntimeError, match="already connected"): + session.connect() + + +def test_a_new_connection_waits_for_the_previous_connections_handler(mocker): + first, second = _FakeServer(), _FakeServer() + mocker.patch( + "assemblyai.tts.v1.session.websocket_connect", side_effect=[first, second] + ) + in_handler = threading.Event() + + def blocking_player(event: aai.TTSAudioEvent) -> None: + in_handler.set() + time.sleep(1) # a playback write that blocks + + session = aai.TTSSession(terminate_timeout=0.2) + session.on(aai.TTSEvents.Audio, blocking_player) + with session: + session.generate("Hi.") + session.flush() + assert in_handler.wait(5) + # the previous reader is still in the handler here + with session: + time.sleep(1.5) # the previous handler has returned by now + session.generate("Still open.") + assert not session._stopped.is_set() + assert session.termination is None + + +def test_unread_audio_is_bounded_and_dropped_oldest_first(server, caplog): + from assemblyai.tts._base_ws import PULL_QUEUE_FRAMES + + with aai.TTSSession() as session: + # audio of a request nobody reads floods the pull queue + for _ in range(PULL_QUEUE_FRAMES + 5): + server.push({"type": "Audio", "audio": "AA==", "flush_id": 7}) + deadline = time.monotonic() + 5 + while not server.inbox.empty() and time.monotonic() < deadline: + time.sleep(0.01) # let the reader queue the flood before anyone reads + assert len(session._queue) == PULL_QUEUE_FRAMES + session.generate("y") + flush_id = session.flush() + # still served after the flood + assert b"".join(_drain(session.audio(flush_id))) == b"y" + assert session._queue.max_audio == PULL_QUEUE_FRAMES + assert "oldest are being dropped" in caplog.text + + +def test_an_overflow_drops_audio_but_never_the_end_of_a_request(server, mocker, caplog): + mocker.patch("assemblyai.tts.v1.session.PULL_QUEUE_FRAMES", 6) + server.answer_flush = False + done: List[int] = [] + with aai.TTSSession() as session: + session.on(aai.TTSEvents.FlushDone, lambda event: done.append(event.flush_id)) + flush_ids = [session.flush() for _ in range(3)] + for flush_id in flush_ids: # 4 frames each, all before anything is read + for _ in range(4): + server.push({"type": "Audio", "audio": "AA==", "flush_id": flush_id}) + server.push( + {"type": "FlushDone", "flush_id": flush_id, "audio_duration_ms": 320} + ) + deadline = time.monotonic() + 5 + while len(done) < 3 and time.monotonic() < deadline: + time.sleep(0.01) + + got = [_drain(session.audio(flush_id)) for flush_id in flush_ids] + + # 12 frames for room for 6: the oldest 6 went, every request still ended + assert [len(chunks) for chunks in got] == [0, 2, 4] + assert "oldest are being dropped" in caplog.text + + +def test_an_audio_handler_receives_the_frames_instead_of_the_queue(server): + received = [] + done = [] + with aai.TTSSession() as session: + session.on(aai.TTSEvents.Audio, received.append) + session.on(aai.TTSEvents.FlushDone, done.append) + session.generate("Pushed.") + flush_id = session.flush() + with pytest.raises(RuntimeError, match="Audio handler"): + _drain(session.audio(flush_id)) + session._stopped.wait(0.2) # let the frames land + + assert [e.audio for e in received] == [b"Pushed."] + assert received[0].flush_id == 0 + assert done[0].audio_duration_ms == 80 + + +def test_the_error_reaches_waiters_before_the_error_handler(server): + app_saw_error = threading.Event() + handler_waited: List[bool] = [] + + def on_error(error: aai.TTSError) -> None: + handler_waited.append(app_saw_error.wait(5)) # waits on the app + + with pytest.raises(aai.TTSError): + with aai.TTSSession() as session: + session.on(aai.TTSEvents.Error, on_error) + server.answer_flush = False + flush_id = session.flush() + server.push(_INPUT_RATE_EXCEEDED) + try: + _drain(session.audio(flush_id), timeout=3) + finally: + app_saw_error.set() + + assert handler_waited == [True] + + +def test_waiting_from_a_handler_raises_instead_of_deadlocking(server): + refused: List[str] = [] + + def speak_next(event: aai.TTSFlushDoneEvent) -> None: + for wait in (lambda: list(session.speak("Next.")), session.terminate): + try: + wait() + except RuntimeError as exc: + refused.append(str(exc)) + + with aai.TTSSession(terminate_timeout=3) as session: + session.on(aai.TTSEvents.FlushDone, speak_next) + assert b"".join(_drain(session.speak("Hi."))) == b"Hi." + + assert len(refused) == 2 + assert all("from a session handler" in message for message in refused) + assert _frames(server, "Flush") == [{"type": "Flush"}] # nothing more was sent + + +def test_word_boundaries_and_cancelled_reach_their_handlers(server): + words = [] + cancelled = [] + with aai.TTSSession(word_boundaries=True) as session: + session.on(aai.TTSEvents.WordBoundaries, words.append) + session.on(aai.TTSEvents.Cancelled, cancelled.append) + _drain(session.speak("Hello there.")) + server.push( + { + "type": "WordBoundaries", + "flush_id": 0, + "audio_start_ms": 0, + "words": [ + {"word": "Hello", "start": 265, "end": 530, "confidence": 0.9} + ], + } + ) + session.cancel() + session.keep_alive() + _drain(session.speak("Again.")) + + assert parse_qs(urlsplit(server.connect.call_args.args[0]).query)[ + "word_boundaries" + ] == ["true"] + assert words[0].words[0].word == "Hello" + assert cancelled[0].audio_duration_seconds == 0.1 + assert [f["type"] for f in server.sent].count("KeepAlive") == 1 + + +def test_a_temporary_token_goes_in_the_query_and_not_the_header(server, monkeypatch): + monkeypatch.setattr(aai.settings, "api_key", None) # a token needs no key + with aai.TTSSession(token="tmp-token", inactivity_timeout=30): + pass + + uri = server.connect.call_args.args[0] + headers = server.connect.call_args.kwargs["additional_headers"] + query = parse_qs(urlsplit(uri).query) + assert query["token"] == ["tmp-token"] + assert query["inactivity_timeout"] == ["30"] + assert "Authorization" not in headers + + +def test_generate_refuses_text_over_one_frame(server): + with aai.TTSSession() as session: + with pytest.raises(ValueError, match="2000"): + session.generate("x" * 2001) + with pytest.raises(TypeError): + session.generate(b"bytes") # type: ignore[arg-type] + + +@pytest.mark.parametrize( + "limit, allowed, refused", [(4000, 4000, 4001), (500, 500, 501)] +) +def test_generate_follows_the_limit_begin_reports(mocker, limit, allowed, refused): + begin = dict(_BEGIN, configuration={"limits": {"max_generate_text_length": limit}}) + mocker.patch( + "assemblyai.tts.v1.session.websocket_connect", + return_value=_FakeServer(begin=begin), + ) + with aai.TTSSession() as session: + session.generate("x" * allowed) + with pytest.raises(ValueError, match=f"{limit}-character"): + session.generate("x" * refused) + + +def test_the_eu_host_comes_from_the_settings(server): + client = aai.Client(api_key="test") + client.settings.tts_base_url = "https://tts.eu.assemblyai.com" + with aai.TTSSession(client=client): + pass + assert urlsplit(server.connect.call_args.args[0]).netloc == "tts.eu.assemblyai.com" + + +def test_a_plain_http_host_opens_a_plain_websocket(server): + client = aai.Client(api_key="test") + client.settings.tts_base_url = "http://localhost:8080" + with aai.TTSSession(client=client): + pass + assert server.connect.call_args.args[0].startswith("ws://localhost:8080/v1/ws?") + + +@pytest.mark.parametrize( + "base_url, expected", + [ + ("https://proxy.example.com/tts", "wss://proxy.example.com/tts/v1/ws?"), + ("https://proxy.example.com/tts/", "wss://proxy.example.com/tts/v1/ws?"), + ("https://tts.us.assemblyai.com/", "wss://tts.us.assemblyai.com/v1/ws?"), + ], +) +def test_a_path_prefix_in_the_host_setting_is_kept(server, base_url, expected): + client = aai.Client(api_key="test") + client.settings.tts_base_url = base_url + with aai.TTSSession(client=client): + pass + assert server.connect.call_args.args[0].startswith(expected) + + +def test_api_key_argument_is_used_without_touching_the_shared_settings(server): + with aai.TTSSession( + api_key="other", config=aai.TTSConfig(extend_biometric_consent=True) + ): + pass + headers = server.connect.call_args.kwargs["additional_headers"] + query = parse_qs(urlsplit(server.connect.call_args.args[0]).query) + assert headers["Authorization"] == "other" + assert query["extend_biometric_consent"] == ["true"] + assert aai.settings.api_key == "test" + + +@pytest.mark.parametrize("configuration", [None, "missing"]) +def test_a_begin_without_configuration_opens_the_session(mocker, configuration): + begin = {k: v for k, v in _BEGIN.items() if k != "configuration"} + if configuration != "missing": + begin["configuration"] = configuration + mocker.patch( + "assemblyai.tts.v1.session.websocket_connect", + return_value=_FakeServer(begin=begin), + ) + with aai.TTSSession() as session: + assert session.begin is not None and session.begin.configuration == {} + assert b"".join(_drain(session.speak("Hi."))) == b"Hi." + + +def test_missing_key_is_a_clear_error(monkeypatch): + monkeypatch.delenv("ASSEMBLYAI_API_KEY", raising=False) + monkeypatch.setattr(aai.settings, "api_key", None) + with pytest.raises(ValueError, match="api_key= or token="): + aai.TTSSession() + + +def test_an_exception_inside_with_closes_without_terminate(server): + with pytest.raises(RuntimeError): + with aai.TTSSession(): + raise RuntimeError("boom") + assert server.sent == [] + assert server.closed + + +def test_wrong_config_type_is_a_type_error(): + with pytest.raises(TypeError): + aai.TTSSession(config=aai.DictationConfig()) # type: ignore[arg-type] + + +def test_a_failure_nobody_read_is_raised_when_the_block_ends(server): + # push mode: an Audio handler, text flushed, then the API fails + with pytest.raises(aai.TTSError) as exc_info: + with aai.TTSSession() as session: + session.on(aai.TTSEvents.Audio, lambda e: None) + server.answer_flush = False + session.generate("Never read.") + session.flush() + server.push( + { + "type": "Error", + "error_code": 3010, + "error": "Too fast", + "error_type": "input_rate_exceeded", + "retryable": False, + "param": None, + "retry_after_seconds": None, + } + ) + session._stopped.wait(2) + assert exc_info.value.error_code == "input_rate_exceeded" + + +def test_an_error_the_block_already_saw_is_not_raised_twice(server): + seen = [] + with aai.TTSSession() as session: + server.answer_flush = False + flush_id = session.flush() + server.push( + { + "type": "Error", + "error_code": 3005, + "error": "No room", + "error_type": "at_capacity", + "retryable": True, + "param": None, + "retry_after_seconds": 2, + } + ) + try: + list(session.audio(flush_id)) + except aai.TTSError as err: + seen.append(err) + assert len(seen) == 1 # the with block exited quietly + + +def test_speak_splits_long_text_across_generate_frames(server): + with aai.TTSSession() as session: + text = "a" * 2000 + "b" * 2000 + "c" * 1000 + assert b"".join(session.speak(text)) == text.encode() + sent = [f["text"] for f in server.sent if f["type"] == "Generate"] + assert [len(t) for t in sent] == [2000, 2000, 1000] + + +def test_a_stale_audio_iterator_ends_when_the_session_reconnects(mocker): + fakes = [_FakeServer(answer_flush=False), _FakeServer()] + mocker.patch("assemblyai.tts.v1.session.websocket_connect", side_effect=fakes) + session = aai.TTSSession() + session.connect() + stale = session.audio(session.flush()) # a player still holding request 0 + fakes[0].push(ConnectionClosedError(Close(3006, "idle"), None)) + session._stopped.wait(2) + with pytest.raises(aai.TTSError): + next(stale) # the old connection's failure, raised once + + session.connect() # same object, new connection, ids start again at 0 + assert list(stale) == [] # the old iterator does not read the new queue + assert b"".join(session.speak("NEW")) == b"NEW" + session.terminate() + + +def test_a_failed_reconnect_does_not_leave_audio_waiting(mocker): + import time + + first = _FakeServer() + mocker.patch( + "assemblyai.tts.v1.session.websocket_connect", + side_effect=[first, OSError("refused")], + ) + session = aai.TTSSession() + with session: + pass + with pytest.raises(aai.TTSError) as exc_info: + session.connect() + assert exc_info.value.retryable is True + started = time.monotonic() + assert list(session.audio(0)) == [] # ends at once instead of polling forever + assert time.monotonic() - started < 2 + + +def test_a_dropped_handshake_is_retryable(): + from websockets.exceptions import InvalidMessage + + from assemblyai.tts._base_ws import handshake_error + + assert handshake_error( + InvalidMessage("did not receive a valid HTTP response") + ).retryable + assert handshake_error(ConnectionRefusedError()).retryable + assert not handshake_error(ValueError("bad uri")).retryable + + +def test_a_send_that_meets_a_close_does_not_hold_the_lock_while_waiting(server): + import time + + server.answer_flush = False + blocked_for = {} + + def on_flush_done(event): + t0 = time.monotonic() + try: + session.keep_alive() + except Exception: + pass + blocked_for["s"] = time.monotonic() - t0 + + with pytest.raises(aai.TTSError) as exc_info: + with aai.TTSSession() as session: + session.on(aai.TTSEvents.FlushDone, on_flush_done) + session.flush() + # the API closes with an Error frame; the main thread's send races it + server.inbox.put( + json.dumps({"type": "FlushDone", "flush_id": 0, "audio_duration_ms": 1}) + ) + server.push( + { + "type": "Error", + "error_code": 3010, + "error": "Too fast", + "error_type": "input_rate_exceeded", + "retryable": False, + "param": None, + "retry_after_seconds": None, + } + ) + server.push(ConnectionClosedError(Close(3010, ""), None)) + session._stopped.wait(2) + session.generate("late") + assert exc_info.value.error_code == "input_rate_exceeded" + assert blocked_for.get("s", 0) < 0.5 diff --git a/tests/unit/test_tts_session_async.py b/tests/unit/test_tts_session_async.py new file mode 100644 index 0000000..4c4c5f3 --- /dev/null +++ b/tests/unit/test_tts_session_async.py @@ -0,0 +1,685 @@ +"""Tests for the asyncio Text-to-Speech WebSocket session.""" + +import asyncio +import base64 +import json +from types import SimpleNamespace +from typing import Any, AsyncIterable, List, Optional +from urllib.parse import parse_qs, urlsplit + +import pytest +import pytest_asyncio +from websockets.exceptions import ConnectionClosedError, ConnectionClosedOK +from websockets.frames import Close + +import assemblyai as aai + +pytestmark = pytest.mark.asyncio + +aai.settings.api_key = "test" + +_BEGIN = { + "type": "Begin", + "id": "sess-1", + "expires_at": 1789556741, + "configuration": {"model": "assemblyai-tts-1"}, +} +_TERMINATION = { + "type": "Termination", + "session_duration_seconds": 3, + "total_input_char_length": 20, + "audio_duration_seconds": 1.5, +} + + +class _Rejected(Exception): + """A handshake answered with an HTTP status, as the current clients raise.""" + + def __init__(self, status: int) -> None: + super().__init__(f"server rejected WebSocket connection: HTTP {status}") + self.response = SimpleNamespace(status_code=status) + + +class _LegacyRejected(Exception): + """The same, as the legacy client of websockets 11 and 12 raises it.""" + + def __init__(self, status: int) -> None: + super().__init__(f"server rejected WebSocket connection: HTTP {status}") + self.status_code = status + + +_HANDSHAKE_FAILURES = [ + pytest.param(_Rejected(503), 503, True, id="503"), + pytest.param(_Rejected(502), 502, True, id="502"), + pytest.param(_Rejected(401), 401, False, id="401"), + pytest.param(_LegacyRejected(429), 429, True, id="legacy-429"), + pytest.param(_LegacyRejected(403), 403, False, id="legacy-403"), + pytest.param( + ConnectionRefusedError(111, "Connection refused"), None, True, id="refused" + ), + pytest.param(TimeoutError("timed out"), None, True, id="timeout"), + pytest.param(asyncio.TimeoutError(), None, True, id="asyncio-timeout"), + pytest.param(ValueError("not a WebSocket URI"), None, False, id="bad-uri"), +] + + +_INACTIVITY_TIMEOUT = { + "type": "Error", + "error_code": 3006, + "error": "No message received for 120 seconds; the idle session was closed.", + "error_type": "inactivity_timeout", + "retryable": True, + "param": None, + "retry_after_seconds": None, +} +_INPUT_RATE_EXCEEDED = { + "type": "Error", + "error_code": 3010, + "error": "Text arrived faster than it can be spoken.", + "error_type": "input_rate_exceeded", + "retryable": False, + "param": None, + "retry_after_seconds": None, +} + + +class _FakeAsyncServer: + def __init__(self, begin: Optional[dict] = _BEGIN, answer_flush: bool = True): + self.sent: List[dict] = [] + self.inbox: "asyncio.Queue[Any]" = asyncio.Queue() + self.closed = False + self.text = "" + self.flush_id = 0 + self.answer_flush = answer_flush + self.answer_terminate = True + self.closing: Optional[List[Any]] = None + if begin is not None: + self.push(begin) + + def close_on_next_send(self, code: int, error: Optional[dict] = None) -> None: + """The API closes the session: the next send raises, and the reader + has yet to take in the frames that came first.""" + self.closing = ([error] if error else []) + [ + ConnectionClosedError(Close(code, "See Error message for details"), None) + ] + + def push(self, item: Any) -> None: + self.inbox.put_nowait(json.dumps(item) if isinstance(item, dict) else item) + + async def send(self, raw: str) -> None: + if self.closing is not None: + closing, self.closing = self.closing, None + for item in closing: + self.push(item) + raise ConnectionClosedError(closing[-1].rcvd, None) + frame = json.loads(raw) + self.sent.append(frame) + kind = frame["type"] + if kind == "Generate": + self.text += frame["text"] + elif kind == "Flush": + if self.answer_flush: + audio = base64.b64encode(self.text.encode()).decode() + self.push({"type": "Audio", "audio": audio, "flush_id": self.flush_id}) + self.push( + { + "type": "FlushDone", + "flush_id": self.flush_id, + "audio_duration_ms": 80, + } + ) + self.flush_id += 1 + self.text = "" + elif kind == "Cancel": + self.flush_id += 1 + self.text = "" + self.push({"type": "Cancelled", "audio_duration_seconds": 0.1}) + elif kind == "Terminate" and self.answer_terminate: + self.push(_TERMINATION) + self.push(ConnectionClosedOK(Close(1000, "Session ended"), None)) + + async def recv(self) -> str: + item = await self.inbox.get() + if isinstance(item, Exception): + raise item + return item + + async def close(self) -> None: + self.closed = True + self.push(ConnectionClosedOK(Close(1000, ""), None)) + + +async def _drain(items: AsyncIterable[Any], timeout: float = 5.0) -> List[Any]: + """Collects `items`, failing the test if it does not end in time.""" + + async def collect() -> List[Any]: + return [item async for item in items] + + return await asyncio.wait_for(collect(), timeout) + + +@pytest_asyncio.fixture +async def server(mocker): + # Built inside the running loop: on Python 3.9 an asyncio.Queue binds to + # the loop that exists when it is created. + fake = _FakeAsyncServer() + + async def connect(uri, additional_headers): + connect.uri = uri + connect.headers = additional_headers + return fake + + mocker.patch( + "assemblyai.tts.v1.async_session.websocket_connect_async", side_effect=connect + ) + fake.connect = connect + return fake + + +async def test_async_with_connects_speaks_and_terminates(server): + async with aai.AsyncTTSSession( + config=aai.TTSConfig(voice_id="aai-alba") + ) as session: + assert session.begin is not None and session.begin.id == "sess-1" + chunks = await _drain(session.speak("Hello there.")) + + assert b"".join(chunks) == b"Hello there." + assert parse_qs(urlsplit(server.connect.uri).query)["voice_id"] == ["aai-alba"] + assert server.connect.headers["Authorization"] == "test" + assert [f["type"] for f in server.sent] == ["Generate", "Flush", "Terminate"] + assert session.termination is not None + assert server.closed + + +async def test_pipelined_requests_and_cancel(server): + async with aai.AsyncTTSSession() as session: + await session.generate("One. ") + first = await session.flush() + second = await session.flush() + assert b"".join(await _drain(session.audio(first))) == b"One. " + assert b"".join(await _drain(session.audio(second))) == b"" + + server.answer_flush = False + third = await session.flush() + await session.cancel() + assert await _drain(session.audio(third)) == [] + + +async def test_an_error_before_begin_raises_from_async_with(mocker): + fake = _FakeAsyncServer(begin=None) + fake.push( + { + "type": "Error", + "error_code": 3009, + "error": "Too many sessions", + "error_type": "rate_limited", + "retryable": True, + "param": None, + "retry_after_seconds": 17, + } + ) + + async def connect(uri, additional_headers): + return fake + + mocker.patch( + "assemblyai.tts.v1.async_session.websocket_connect_async", side_effect=connect + ) + + with pytest.raises(aai.TTSError) as exc_info: + async with aai.AsyncTTSSession(): + pass + assert exc_info.value.error_code == "rate_limited" + assert exc_info.value.retry_after == 17 + assert exc_info.value.body["close_code"] == 3009 + assert fake.closed + + +@pytest.mark.parametrize("failure, status, retryable", _HANDSHAKE_FAILURES) +async def test_a_failed_handshake_keeps_its_status_and_says_if_retryable( + mocker, failure, status, retryable +): + async def connect(uri, additional_headers): + raise failure + + mocker.patch( + "assemblyai.tts.v1.async_session.websocket_connect_async", side_effect=connect + ) + with pytest.raises(aai.TTSError, match="handshake failed") as exc_info: + async with aai.AsyncTTSSession(): + pass + assert exc_info.value.status_code == status + assert exc_info.value.retryable is retryable + assert exc_info.value.__cause__ is failure + + +async def test_handlers_may_be_coroutines(server): + received = [] + + async def on_audio(event): + received.append(event.audio) + + async with aai.AsyncTTSSession() as session: + session.on(aai.TTSEvents.Audio, on_audio) + await session.generate("Pushed.") + flush_id = await session.flush() + with pytest.raises(RuntimeError, match="Audio handler"): + await session.audio(flush_id).__anext__() + await asyncio.sleep(0.05) + + assert received == [b"Pushed."] + + +async def test_the_error_reaches_waiters_before_the_error_handler(server): + app_saw_error = asyncio.Event() + handler_waited: List[bool] = [] + + async def on_error(error: aai.TTSError) -> None: + try: # waits on the app + await asyncio.wait_for(app_saw_error.wait(), 5) + handler_waited.append(True) + except asyncio.TimeoutError: + handler_waited.append(False) + + with pytest.raises(aai.TTSError): + async with aai.AsyncTTSSession() as session: + session.on(aai.TTSEvents.Error, on_error) + server.answer_flush = False + flush_id = await session.flush() + server.push(_INPUT_RATE_EXCEEDED) + try: + await _drain(session.audio(flush_id), timeout=3) + finally: + app_saw_error.set() + + assert handler_waited == [True] + + +async def test_waiting_from_a_handler_raises_instead_of_deadlocking(server): + refused: List[str] = [] + + async def speak_next(event: aai.TTSFlushDoneEvent) -> None: + try: + [chunk async for chunk in session.speak("Next.")] + except RuntimeError as exc: + refused.append(str(exc)) + try: + await session.terminate() + except RuntimeError as exc: + refused.append(str(exc)) + + async with aai.AsyncTTSSession(terminate_timeout=3) as session: + session.on(aai.TTSEvents.FlushDone, speak_next) + assert b"".join(await _drain(session.speak("Hi."))) == b"Hi." + + assert len(refused) == 2 + assert all("from a session handler" in message for message in refused) + assert [f for f in server.sent if f["type"] == "Flush"] == [{"type": "Flush"}] + + +async def test_a_lost_connection_raises_a_retryable_error(server): + with pytest.raises(aai.TTSError) as exc_info: + async with aai.AsyncTTSSession() as session: + server.answer_flush = False + flush_id = await session.flush() + server.push(ConnectionClosedError(None, None)) # abnormal: no close frame + await _drain(session.audio(flush_id)) + assert exc_info.value.retryable is True + assert exc_info.value.body["close_code"] == 1006 + + +@pytest.mark.filterwarnings("error::DeprecationWarning") +@pytest.mark.parametrize( + "code, retryable", + [ + (1011, True), + (1012, True), + (3005, True), + (3008, True), + (3009, True), + (3006, False), + (3010, False), + ], +) +async def test_a_close_without_an_error_frame_is_retryable_by_its_code( + server, code, retryable +): + with pytest.raises(aai.TTSError) as exc_info: + async with aai.AsyncTTSSession() as session: + server.answer_flush = False + flush_id = await session.flush() + server.push(ConnectionClosedError(Close(code, "Upstream failed"), None)) + await _drain(session.audio(flush_id)) + + assert exc_info.value.retryable is retryable + assert exc_info.value.body == {"close_code": code, "reason": "Upstream failed"} + + +async def test_a_socket_error_is_a_retryable_lost_connection(server): + with pytest.raises(aai.TTSError) as exc_info: + async with aai.AsyncTTSSession() as session: + server.answer_flush = False + flush_id = await session.flush() + server.push(OSError("Connection reset by peer")) + await _drain(session.audio(flush_id)) + + assert exc_info.value.retryable is True + assert exc_info.value.body == {"close_code": 1006, "reason": ""} + + +async def test_a_failure_in_the_reader_still_ends_the_iteration(server, mocker): + mocker.patch( + "assemblyai.tts.v1.async_session.error_from_close", + side_effect=RuntimeError("bug"), + ) + with pytest.raises(aai.TTSError, match="reader failed"): + async with aai.AsyncTTSSession() as session: + server.answer_flush = False + flush_id = await session.flush() + server.push(ConnectionClosedError(None, None)) + await _drain(session.audio(flush_id)) + + +async def test_a_send_the_api_closed_raises_the_sessions_error(server): + with pytest.raises(aai.TTSError) as exc_info: + async with aai.AsyncTTSSession() as session: + server.close_on_next_send(3010, _INPUT_RATE_EXCEEDED) + await session.generate("Too much, too fast.") + + assert exc_info.value.error_code == "input_rate_exceeded" + assert exc_info.value.retryable is False + assert exc_info.value is session.error + + +async def test_a_send_on_a_bare_close_raises_a_tts_error(server): + with pytest.raises(aai.TTSError) as exc_info: + async with aai.AsyncTTSSession() as session: + server.close_on_next_send(1012) + await session.flush() + + assert exc_info.value.retryable is True + assert exc_info.value.body["close_code"] == 1012 + + +async def test_a_cancel_drops_every_pipelined_request(server): + server.answer_flush = False + async with aai.AsyncTTSSession() as session: + first = await session.flush() + second = await session.flush() + await session.cancel() + assert await _drain(session.audio(first)) == [] + assert await _drain(session.audio(second)) == [] + + +async def test_an_overflow_drops_audio_but_never_the_end_of_a_request( + server, mocker, caplog +): + mocker.patch("assemblyai.tts.v1.async_session.PULL_QUEUE_FRAMES", 6) + server.answer_flush = False + done: List[int] = [] + async with aai.AsyncTTSSession() as session: + session.on(aai.TTSEvents.FlushDone, lambda event: done.append(event.flush_id)) + flush_ids = [await session.flush() for _ in range(3)] + for flush_id in flush_ids: # 4 frames each, all before anything is read + for _ in range(4): + server.push({"type": "Audio", "audio": "AA==", "flush_id": flush_id}) + server.push( + {"type": "FlushDone", "flush_id": flush_id, "audio_duration_ms": 320} + ) + for _ in range(500): + if len(done) == 3: + break + await asyncio.sleep(0.01) + + got = [await _drain(session.audio(flush_id)) for flush_id in flush_ids] + + # 12 frames for room for 6: the oldest 6 went, every request still ended + assert [len(chunks) for chunks in got] == [0, 2, 4] + assert "oldest are being dropped" in caplog.text + + +async def test_a_session_object_can_be_reused(mocker): + fakes = [_FakeAsyncServer(), _FakeAsyncServer()] + + async def connect(uri, additional_headers): + return fakes.pop(0) + + mocker.patch( + "assemblyai.tts.v1.async_session.websocket_connect_async", side_effect=connect + ) + session = aai.AsyncTTSSession() + async with session: + assert b"".join(await _drain(session.speak("One."))) == b"One." + async with session: + assert b"".join(await _drain(session.speak("Two."))) == b"Two." + + +async def test_a_cancelled_connect_leaves_nothing_open(mocker): + slow = _FakeAsyncServer(begin=None) # Begin never comes + fakes = [slow, _FakeAsyncServer()] + + async def connect(uri, additional_headers): + return fakes.pop(0) + + mocker.patch( + "assemblyai.tts.v1.async_session.websocket_connect_async", side_effect=connect + ) + session = aai.AsyncTTSSession() + with pytest.raises(asyncio.TimeoutError): + await asyncio.wait_for(session.connect(), 0.3) + + assert slow.closed + assert session._reader is not None and session._reader.done() + begin = await asyncio.wait_for(session.connect(), 5) # not "already connected" + assert begin.id == "sess-1" + await session.close() + + +async def test_a_cancelled_terminate_still_closes(server): + server.answer_terminate = False + session = aai.AsyncTTSSession(terminate_timeout=30) + await session.connect() + ending = asyncio.create_task(session.terminate()) + await asyncio.sleep(0.05) + ending.cancel() + await asyncio.wait({ending}, timeout=5) + + assert ending.cancelled() + assert server.closed + assert session._ws is None + + +async def test_a_cancelled_close_propagates_the_cancellation(server): + in_handler = asyncio.Event() + + async def busy_handler(event): + in_handler.set() + await asyncio.sleep(30) + + session = aai.AsyncTTSSession(terminate_timeout=30) + await session.connect() + session.on(aai.TTSEvents.FlushDone, busy_handler) + await session.flush() + await asyncio.wait_for(in_handler.wait(), 5) + + closing = asyncio.create_task(session.close()) + await asyncio.sleep(0.05) # close() now waits for the busy reader + closing.cancel() + await asyncio.wait({closing}, timeout=5) + + assert closing.cancelled() + await asyncio.sleep(0) + assert session._reader is not None and session._reader.done() + + +def _connect_to(mocker, *fakes: "_FakeAsyncServer") -> None: + queued = list(fakes) + + async def connect(uri, additional_headers): + return queued.pop(0) + + mocker.patch( + "assemblyai.tts.v1.async_session.websocket_connect_async", side_effect=connect + ) + + +async def test_connect_after_the_api_ended_the_session_opens_a_new_one(mocker): + first, second = _FakeAsyncServer(), _FakeAsyncServer() + _connect_to(mocker, first, second) + session = aai.AsyncTTSSession() + await session.connect() + first.push(_INACTIVITY_TIMEOUT) + first.push( + ConnectionClosedError(Close(3006, "See Error message for details"), None) + ) + await asyncio.wait_for(session._stopped.wait(), 5) + assert session.error is not None and session.error.retryable is True + + await session.connect() # no close() needed first + assert first.closed + assert session.error is None + assert b"".join(await _drain(session.speak("Hello again."))) == b"Hello again." + await session.terminate() + + +async def test_connect_on_a_live_session_is_refused(server): + async with aai.AsyncTTSSession() as session: + with pytest.raises(RuntimeError, match="already connected"): + await session.connect() + + +async def test_connect_right_after_close_gave_up_on_a_handler(mocker): + first, second = _FakeAsyncServer(), _FakeAsyncServer() + _connect_to(mocker, first, second) + in_handler = asyncio.Event() + + async def busy_handler(event): + in_handler.set() + await asyncio.sleep(30) + + session = aai.AsyncTTSSession(terminate_timeout=0.1) + await session.connect() + session.on(aai.TTSEvents.Cancelled, busy_handler) + await session.cancel() + await asyncio.wait_for(in_handler.wait(), 5) + await session.close() # gives up on the busy reader and cancels it + + await asyncio.wait_for(session.connect(), 5) # at once + await asyncio.sleep(0.05) + assert not session._stopped.is_set() + assert b"".join(await _drain(session.speak("Hi."))) == b"Hi." + await session.terminate() + + +@pytest.mark.parametrize( + "limit, allowed, refused", [(4000, 4000, 4001), (500, 500, 501)] +) +async def test_generate_follows_the_limit_begin_reports( + mocker, limit, allowed, refused +): + begin = dict(_BEGIN, configuration={"limits": {"max_generate_text_length": limit}}) + _connect_to(mocker, _FakeAsyncServer(begin=begin)) + async with aai.AsyncTTSSession() as session: + await session.generate("x" * allowed) + with pytest.raises(ValueError, match=f"{limit}-character"): + await session.generate("x" * refused) + + +@pytest.mark.parametrize("configuration", [None, "missing"]) +async def test_a_begin_without_configuration_opens_the_session(mocker, configuration): + begin = {k: v for k, v in _BEGIN.items() if k != "configuration"} + if configuration != "missing": + begin["configuration"] = configuration + _connect_to(mocker, _FakeAsyncServer(begin=begin)) + async with aai.AsyncTTSSession() as session: + assert session.begin is not None and session.begin.configuration == {} + assert b"".join(await _drain(session.speak("Hi."))) == b"Hi." + + +async def test_audio_before_connect_says_the_session_is_not_open(server): + with pytest.raises(RuntimeError, match="not open"): + await _drain(aai.AsyncTTSSession().audio()) + + +async def test_a_path_prefix_in_the_host_setting_is_kept(server, monkeypatch): + monkeypatch.setattr(aai.settings, "tts_base_url", "https://proxy.example.com/tts/") + async with aai.AsyncTTSSession(): + pass + assert server.connect.uri.startswith("wss://proxy.example.com/tts/v1/ws?") + + +@pytest.mark.parametrize( + "config", + [aai.TTSConfig(container="wav"), aai.TTSConfig(bit_rate=64000)], + ids=["wav", "bit_rate"], +) +async def test_http_only_options_are_refused_when_the_session_is_built(config): + with pytest.raises(ValueError, match="wav container and bit_rate"): + aai.AsyncTTSSession(config=config) + + +async def test_a_raw_container_is_what_a_session_streams(server): + async with aai.AsyncTTSSession(config=aai.TTSConfig(container="raw")): + pass + assert "container" not in parse_qs(urlsplit(server.connect.uri).query) + + +async def test_missing_key_is_a_clear_error(monkeypatch): + monkeypatch.delenv("ASSEMBLYAI_API_KEY", raising=False) + monkeypatch.setattr(aai.settings, "api_key", None) + with pytest.raises(ValueError, match="api_key= or token="): + aai.AsyncTTSSession() + + +async def test_a_failure_nobody_read_is_raised_when_the_block_ends(server): + with pytest.raises(aai.TTSError) as exc_info: + async with aai.AsyncTTSSession() as session: + session.on(aai.TTSEvents.Audio, lambda e: None) + server.answer_flush = False + await session.generate("Never read.") + await session.flush() + server.push( + { + "type": "Error", + "error_code": 3010, + "error": "Too fast", + "error_type": "input_rate_exceeded", + "retryable": False, + "param": None, + "retry_after_seconds": None, + } + ) + await asyncio.wait_for(session._stopped.wait(), 2) + assert exc_info.value.error_code == "input_rate_exceeded" + + +async def test_speak_splits_long_text_across_generate_frames(server): + async with aai.AsyncTTSSession() as session: + text = "a" * 2000 + "b" * 500 + assert b"".join([c async for c in session.speak(text)]) == text.encode() + sent = [f["text"] for f in server.sent if f["type"] == "Generate"] + assert [len(t) for t in sent] == [2000, 500] + + +async def test_a_stale_audio_iterator_ends_when_the_session_reconnects(mocker): + first, second = _FakeAsyncServer(answer_flush=False), _FakeAsyncServer() + servers = [first, second] + + async def connect(uri, additional_headers): + return servers.pop(0) + + mocker.patch( + "assemblyai.tts.v1.async_session.websocket_connect_async", side_effect=connect + ) + session = aai.AsyncTTSSession() + await session.connect() + stale = session.audio(await session.flush()) + first.push(ConnectionClosedError(Close(3006, "idle"), None)) + await asyncio.wait_for(session._stopped.wait(), 2) + with pytest.raises(aai.TTSError): + await stale.__anext__() + + await session.connect() + assert [c async for c in stale] == [] + assert b"".join([c async for c in session.speak("NEW")]) == b"NEW" + await session.terminate()