diff --git a/CLAUDE.md b/CLAUDE.md index 56dc4c3..3a0bf9f 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -38,6 +38,9 @@ aai.settings.api_key = "your-key" ## Key classes - `aai.Transcriber` — Transcribe files, URLs, or streams. Methods: `transcribe()`, `transcribe_async()`, `submit()`, `list_transcripts()` +- `aai.AsyncTranscriber` — Asyncio counterpart of `Transcriber`. Same options, every API call a coroutine. Methods: `transcribe()`, `submit()`, `transcribe_group()`, `get_by_id()`, `delete_by_id()`, `list_transcripts()`, `upload_file()` +- `aai.AsyncTranscript` — What `AsyncTranscriber` returns. Same fields as `Transcript`, coroutine methods +- `aai.AsyncClient` — Shared `httpx.AsyncClient` pool for one or more `AsyncTranscriber`s - `aai.TranscriptionConfig` — All transcription options: `speech_models`, `speaker_labels`, `sentiment_analysis`, `entity_detection`, `auto_chapters`, `content_safety`, `language_detection`, `summarization`, `word_boost`, `disfluencies` - `aai.Transcript` — Result object with `.text`, `.status`, `.utterances`, `.words`, `.chapters`, `.entities`, `.sentiment_analysis`. Methods: `get_sentences()`, `get_paragraphs()`, `export_subtitles_srt()`, `export_subtitles_vtt()` - `aai.SyncTranscriber` — Synchronous pre-recorded transcription: audio in, transcript out, one request (no polling). Methods: `transcribe()`, `transcribe_async()` @@ -80,6 +83,62 @@ config.set_redact_pii( transcript = aai.Transcript.get_by_id("transcript-id") ``` +## Asyncio transcription (`AsyncTranscriber`) + +`AsyncTranscriber` is `Transcriber` with coroutines instead of a thread pool. Use it in +asyncio code (FastAPI, aiohttp, voice agents). Do not use `transcribe_async()` there: its +`concurrent.futures.Future` is not awaitable, and `.result()` blocks the event loop. Same +`TranscriptionConfig`, same fields on the result. + +```python +import asyncio +import assemblyai as aai + +aai.settings.api_key = os.environ["ASSEMBLYAI_API_KEY"] + +async def main(): + async with aai.AsyncTranscriber() as transcriber: + transcript = await transcriber.transcribe("./audio.mp3") + if transcript.status == aai.TranscriptStatus.error: + raise RuntimeError(transcript.error) + print(transcript.text) + + sentences = await transcript.get_sentences() # follow-ups are coroutines too + srt = await transcript.export_subtitles_srt() + +asyncio.run(main()) +``` + +**Lifecycle**: the transcriber owns an HTTP connection pool. Use `async with`, or call +`await transcriber.aclose()`. To share one pool, pass a client, which stays yours to close: +```python +async with aai.AsyncClient(settings=aai.settings) as client: + transcriber = aai.AsyncTranscriber(client=client) +``` +There is **no** process-wide default async client, unlike sync `Client`. An +`httpx.AsyncClient` pool is bound to the event loop that first used it, so a global one +breaks under a second `asyncio.run(...)`. + +**Concurrency**: use `asyncio.gather(...)`, or the group helpers to cap in-flight work. +Results come back in input order: +```python +transcripts = await transcriber.transcribe_group(files, max_concurrency=8) +transcripts, errors = await transcriber.transcribe_group(files, return_failures=True) +``` +Unlike the sync `transcribe_group`, failures are never dropped. The first error raises +after the batch settles, or you collect them with `return_failures=True`. + +**Method mapping** (sync → async): `Transcript.get_by_id(id)` → `await transcriber.get_by_id(id)`; +`Transcript.delete_by_id(id)` → `await transcriber.delete_by_id(id)`. Both moved onto the +transcriber, which owns the connection pool. + +**Uploads**: local paths and file objects stream to the upload endpoint in chunks read on +a worker thread. The request sets `Content-Length` when the size is known. A large file +never blocks the loop and never loads fully into memory. + +**LeMUR** is sync-only. An `AsyncTranscript` works as a `LemurSource`, but the LeMUR call +blocks. Run it off the loop, for example with `asyncio.to_thread`. + ## Sync transcription (pre-recorded, single request) `SyncTranscriber` posts a whole audio file and returns the finished transcript in one @@ -323,7 +382,7 @@ async with AsyncStreamingClient(StreamingClientOptions(token=token_from_server)) - **PII redaction uses `set_redact_pii()`**, not a constructor parameter - **Streaming v3 lives in its own module**: `assemblyai.streaming.v3` (there is no other streaming API in this SDK). See the "Streaming (real-time)" section above. - **Microphone streaming needs extras**: `pip install "assemblyai[extras]"` for `pyaudio` -- **`transcribe_async()` returns a `concurrent.futures.Future`**, not an asyncio coroutine +- **`transcribe_async()` returns a `concurrent.futures.Future`**, not an asyncio coroutine. In asyncio code use `aai.AsyncTranscriber` (see "Asyncio transcription" above) - **Timestamps are in milliseconds** throughout the SDK - **Minimum Python**: 3.8+ diff --git a/README.md b/README.md index e9b9d8a..902c477 100644 --- a/README.md +++ b/README.md @@ -49,6 +49,7 @@ See [Coding agent prompts](https://www.assemblyai.com/docs/coding-agent-prompts) - [Installation](#installation) - [Examples](#examples) - [**Core Examples**](#core-examples) + - [**Asyncio Examples**](#asyncio-examples) - [**Speech Understanding Examples**](#speech-understanding-examples) - [**Streaming Examples**](#streaming-examples) - [**Change the default settings**](#change-the-default-settings) @@ -58,6 +59,7 @@ See [Coding agent prompts](https://www.assemblyai.com/docs/coding-agent-prompts) - [Defining Defaults](#defining-defaults) - [Overriding Defaults](#overriding-defaults) - [Synchronous vs Asynchronous](#synchronous-vs-asynchronous) + - [Asyncio](#asyncio) - [Getting the HTTP status code](#getting-the-http-status-code) - [Polling Intervals](#polling-intervals) - [Retrieving Existing Transcripts](#retrieving-existing-transcripts) @@ -518,6 +520,146 @@ except aai.SyncTranscriptError as error: --- +### **Asyncio Examples** + +`aai.AsyncTranscriber` is the asyncio counterpart of `aai.Transcriber`. Every method +that calls the API is a coroutine. Hundreds of transcriptions run concurrently on one +thread, with no thread pool. + +Use it in asyncio code (FastAPI, aiohttp, voice agents). Do not use +`transcribe_async` there: it returns a +[`concurrent.futures.Future`](https://docs.python.org/3/library/concurrent.futures.html), +which is not awaitable. + +
+ Transcribe a file with asyncio + +```python +import asyncio + +import assemblyai as aai + +aai.settings.api_key = "" + + +async def main(): + config = aai.TranscriptionConfig( + speech_models=["universal-3-5-pro", "universal-2"], + speaker_labels=True, + ) + + async with aai.AsyncTranscriber(config=config) as transcriber: + transcript = await transcriber.transcribe("./example.mp3") + + if transcript.status == aai.TranscriptStatus.error: + raise RuntimeError(f"Transcription failed: {transcript.error}") + + print(transcript.text) + + +asyncio.run(main()) +``` + +The transcriber owns an HTTP connection pool. Close it with an async context manager, +or call `await transcriber.aclose()`. + +
+ +
+ Transcribe many files concurrently + +```python +import asyncio + +import assemblyai as aai + +aai.settings.api_key = "" + + +async def main(): + async with aai.AsyncTranscriber() as transcriber: + # Plain asyncio - one coroutine per file, all in flight at once. + transcripts = await asyncio.gather( + transcriber.transcribe("./one.mp3"), + transcriber.transcribe("./two.mp3"), + ) + + # Or let the transcriber cap how many run at a time. Results come back + # in the order of the input. + transcripts = await transcriber.transcribe_group( + ["./one.mp3", "./two.mp3", "./three.mp3"], + max_concurrency=2, + ) + + for transcript in transcripts: + print(transcript.text) + + +asyncio.run(main()) +``` + +Pass `return_failures=True` to get `(transcripts, errors)` instead of a raised error. + +
+ +
+ Submit now, collect later + +```python +import asyncio + +import assemblyai as aai + +aai.settings.api_key = "" + + +async def main(): + async with aai.AsyncTranscriber() as transcriber: + # Returns as soon as the job is queued - no polling. + transcript = await transcriber.submit("https://example.org/audio.mp3") + print(transcript.id, transcript.status) + + # Later, in the same or another process: + transcript = await transcriber.get_by_id(transcript.id) + print(transcript.text) + + # Follow-up operations are coroutines too. + sentences = await transcript.get_sentences() + srt = await transcript.export_subtitles_srt() + + +asyncio.run(main()) +``` + +
+ +
+ Share one connection pool across transcribers + +```python +import assemblyai as aai + +aai.settings.api_key = "" + + +async def main(): + async with aai.AsyncClient(settings=aai.settings) as client: + verbatim = aai.AsyncTranscriber( + client=client, + config=aai.TranscriptionConfig(disfluencies=True), + ) + clean = aai.AsyncTranscriber(client=client) + + # A client that was passed in is not closed by the transcribers - the + # `async with` above owns it. + await verbatim.transcribe("./interview.mp3") + await clean.transcribe("./interview.mp3") +``` + +
+ +--- + ### **Speech Understanding Examples**
@@ -1176,6 +1318,45 @@ The asynchronous approach allows the application to continue running while the t You can identify those two approaches by the `_async` suffix in the `Transcriber`'s method name (e.g. `transcribe` vs `transcribe_async`). +A `concurrent.futures.Future` is not awaitable, and its `.result()` blocks the event +loop. In an asyncio application, use [`AsyncTranscriber`](#asyncio) instead. + +## Asyncio + +`aai.AsyncTranscriber` mirrors `aai.Transcriber` with coroutines instead of threads: + +| `Transcriber` (threads) | `AsyncTranscriber` (asyncio) | +| -------------------------------------- | ------------------------------------------------ | +| `transcribe(...)` | `await transcribe(...)` | +| `transcribe_async(...)` -> `Future` | `await transcribe(...)` (or `asyncio.gather`) | +| `submit(...)` | `await submit(...)` | +| `transcribe_group(...)` | `await transcribe_group(...)` | +| `Transcript.get_by_id(id)` | `await transcriber.get_by_id(id)` | +| `Transcript.delete_by_id(id)` | `await transcriber.delete_by_id(id)` | +| `transcript.get_sentences()` | `await transcript.get_sentences()` | +| `list_transcripts(...)` | `await list_transcripts(...)` | + +Notes: + +- `AsyncTranscriber` owns an HTTP connection pool. Close it with an async context + manager, or call `await transcriber.aclose()`. +- Pass `client=aai.AsyncClient(settings=aai.settings)` to share one pool between + transcribers. A client you pass in stays yours to close. +- There is no process-wide default async client. An `httpx.AsyncClient` pool belongs to + the event loop that first used it, so a global pool fails on a second `asyncio.run()`. +- `AsyncTranscript` carries the same fields as `Transcript`. Only the methods that call + the API became coroutines. +- Uploads stream local files and file objects through a thread, so a large upload never + blocks the loop. +- `transcribe_group` and `submit_group` return results in input order. Both cap in-flight + work at `max_concurrency`, which defaults to 8. +- Neither group method drops a failure. Either the first error is raised, or you pass + `return_failures=True` and get `(transcripts, errors)`. +- LeMUR is sync-only. An `AsyncTranscript` works as a `LemurSource`, but the LeMUR call + blocks. Run it in a thread, for example with `asyncio.to_thread`. + +For real-time streaming, use `assemblyai.streaming.v3.AsyncStreamingClient`. + ## Getting the HTTP status code There are two ways of accessing the HTTP status code: diff --git a/assemblyai/__init__.py b/assemblyai/__init__.py index 7d92a77..62439f0 100644 --- a/assemblyai/__init__.py +++ b/assemblyai/__init__.py @@ -1,5 +1,7 @@ from . import extras from .__version__ import __version__ +from .async_client import AsyncClient +from .async_transcriber import AsyncTranscriber, AsyncTranscript from .client import Client from .lemur import Lemur from .sync import SyncTranscriber @@ -91,6 +93,9 @@ __all__ = [ # types "AssemblyAIError", + "AsyncClient", + "AsyncTranscriber", + "AsyncTranscript", "AutohighlightResponse", "AutohighlightResult", "Chapter", diff --git a/assemblyai/async_api.py b/assemblyai/async_api.py new file mode 100644 index 0000000..a701044 --- /dev/null +++ b/assemblyai/async_api.py @@ -0,0 +1,215 @@ +""" +Asyncio counterparts of the request functions in `assemblyai.api`. + +Each function calls the same endpoint as its sync twin. Both share the helpers +in `api`, so both raise the same exception type and message. +""" + +from typing import AsyncIterable, Dict, List, Optional, Union + +import httpx + +from . import api, types +from .api import ENDPOINT_TRANSCRIPT, ENDPOINT_UPLOAD + +__all__ = [ + "create_transcript", + "delete_transcript", + "export_subtitles_srt", + "export_subtitles_vtt", + "get_paragraphs", + "get_redacted_audio", + "get_sentences", + "get_transcript", + "list_transcripts", + "upload_file", + "word_search", +] + + +async def create_transcript( + client: httpx.AsyncClient, + request: types.TranscriptRequest, +) -> types.TranscriptResponse: + response = await client.post( + ENDPOINT_TRANSCRIPT, + json=api._transcript_request_json(request), + ) + + api._raise_for_status(response, f"failed to transcribe url {request.audio_url}") + + return types.TranscriptResponse.parse_obj(response.json()) + + +async def get_transcript( + client: httpx.AsyncClient, + transcript_id: str, +) -> types.TranscriptResponse: + response = await client.get( + f"{ENDPOINT_TRANSCRIPT}/{transcript_id}", + ) + + api._raise_for_status(response, f"failed to retrieve transcript {transcript_id}") + + return types.TranscriptResponse.parse_obj(response.json()) + + +async def delete_transcript( + client: httpx.AsyncClient, + transcript_id: str, +) -> types.TranscriptResponse: + response = await client.delete( + f"{ENDPOINT_TRANSCRIPT}/{transcript_id}", + ) + + api._raise_for_status(response, f"failed to delete transcript {transcript_id}") + + return types.TranscriptResponse.parse_obj(response.json()) + + +async def upload_file( + client: httpx.AsyncClient, + audio_file: Union[bytes, AsyncIterable[bytes]], + headers: Optional[Dict[str, str]] = None, +) -> str: + """ + Uploads the given audio. + + Args: + `client`: the HTTP client + `audio_file`: the raw audio bytes, or an async iterable of bytes. Do + not pass a blocking file object. httpx would read it on the event + loop. See `AsyncTranscriber.upload_file`. + `headers`: extra headers for the body, such as `Content-Length`. Without + it, httpx uses chunked transfer encoding for an async iterable. + + Returns: The URL of the uploaded audio file. + """ + + response = await client.post( + ENDPOINT_UPLOAD, + content=audio_file, + headers=headers, + ) + + api._raise_for_status(response, "Failed to upload audio file") + + return response.json()["upload_url"] + + +async def export_subtitles_srt( + client: httpx.AsyncClient, + transcript_id: str, + chars_per_caption: Optional[int], +) -> str: + response = await client.get( + f"{ENDPOINT_TRANSCRIPT}/{transcript_id}/srt", + params=api._subtitles_params(chars_per_caption), + ) + + api._raise_for_status( + response, f"failed to export SRT for transcript {transcript_id}" + ) + + return response.text + + +async def export_subtitles_vtt( + client: httpx.AsyncClient, + transcript_id: str, + chars_per_caption: Optional[int], +) -> str: + response = await client.get( + f"{ENDPOINT_TRANSCRIPT}/{transcript_id}/vtt", + params=api._subtitles_params(chars_per_caption), + ) + + api._raise_for_status( + response, f"failed to export VTT for transcript {transcript_id}" + ) + + return response.text + + +async def word_search( + client: httpx.AsyncClient, + transcript_id: str, + words: List[str], +) -> types.WordSearchMatchResponse: + response = await client.get( + f"{ENDPOINT_TRANSCRIPT}/{transcript_id}/word-search", + params=api._word_search_params(words), + ) + + api._raise_for_status( + response, f"failed to search words in transcript {transcript_id}" + ) + + return types.WordSearchMatchResponse.parse_obj(response.json()) + + +async def get_redacted_audio( + client: httpx.AsyncClient, + transcript_id: str, +) -> types.RedactedAudioResponse: + """ + Retrieves the object containing the redacted audio URL for the given transcript. + + Raises: + RedactedAudioIncompleteError: If response indicates that the redacted audio is still processing + RedactedAudioExpiredError: If response indicates that the redacted audio is no longer available + TranscriptError: If we fail to get a valid response from the API at all + + Returns: + `RedactedAudioResponse`, which contains the URL of the redacted audio + """ + + response = await client.get(f"{ENDPOINT_TRANSCRIPT}/{transcript_id}/redacted-audio") + + return api._parse_redacted_audio_response(response, transcript_id) + + +async def get_sentences( + client: httpx.AsyncClient, + transcript_id: str, +) -> types.SentencesResponse: + response = await client.get( + f"{ENDPOINT_TRANSCRIPT}/{transcript_id}/sentences", + ) + + api._raise_for_status( + response, f"failed to retrieve sentences for transcript {transcript_id}" + ) + + return types.SentencesResponse.parse_obj(response.json()) + + +async def get_paragraphs( + client: httpx.AsyncClient, + transcript_id: str, +) -> types.ParagraphsResponse: + response = await client.get( + f"{ENDPOINT_TRANSCRIPT}/{transcript_id}/paragraphs", + ) + + api._raise_for_status( + response, f"failed to retrieve paragraphs for transcript {transcript_id}" + ) + + return types.ParagraphsResponse.parse_obj(response.json()) + + +async def list_transcripts( + client: httpx.AsyncClient, + params: Optional[types.ListTranscriptParameters], +) -> types.ListTranscriptResponse: + response = await client.get( + ENDPOINT_TRANSCRIPT, + params=api._list_transcripts_params(params), + ) + + api._raise_for_status( + response, "failed to retrieve transcripts", types.AssemblyAIError + ) + + return types.ListTranscriptResponse.parse_obj(response.json()) diff --git a/assemblyai/async_client.py b/assemblyai/async_client.py new file mode 100644 index 0000000..76382d4 --- /dev/null +++ b/assemblyai/async_client.py @@ -0,0 +1,113 @@ +from types import TracebackType +from typing import Optional, Type + +import httpx +from typing_extensions import Self + +from . import types +from .client import _build_headers, _build_limits + + +class AsyncClient: + """ + The asyncio counterpart of `Client`. Holds an `httpx.AsyncClient`. + + `AsyncClient` has no process-wide default instance. An `httpx.AsyncClient` + pool belongs to the event loop that first used it, so a global pool fails on + a second event loop. `AsyncTranscriber` creates one client per instance. + Pass an `AsyncClient` to share one pool between transcribers. + + Close the pool with `async with` or `aclose()`. + + Example: + ```python + import assemblyai as aai + + async with aai.AsyncClient(settings=aai.settings) as client: + transcriber = aai.AsyncTranscriber(client=client) + transcript = await transcriber.transcribe("./audio.mp3") + ``` + """ + + def __init__( + self, + *, + settings: types.Settings, + api_key_required: bool = True, + ) -> None: + """ + Creates the asyncio AssemblyAI client. + + Args: + settings: The settings to use for the client. + api_key_required: If an API key is required (either as environment variable or the global settings). + Can be set to `False` if a different authentication method is used, e.g., a temporary token. + """ + + self._settings = settings.copy() + + if api_key_required and not self._settings.api_key: + raise ValueError( + "Please provide an API key via the ASSEMBLYAI_API_KEY environment variable or the global settings." + ) + + self._last_response: Optional[httpx.Response] = None + + async def _store_response(response: httpx.Response) -> None: + self._last_response = response + + self._http_client = httpx.AsyncClient( + base_url=self._settings.base_url, + headers=_build_headers(self._settings), + timeout=self._settings.http_timeout, + limits=_build_limits(self._settings), + event_hooks={"response": [_store_response]}, + ) + + @property + def last_response(self) -> Optional[httpx.Response]: + """ + Get the last HTTP response, corresponding to the last request sent from this client. + + Returns: + The last HTTP response. + """ + return self._last_response + + @property + def settings(self) -> types.Settings: + """ + Get the current settings. + + Returns: + The current settings. + """ + + return self._settings + + @property + def http_client(self) -> httpx.AsyncClient: + """ + Get the current HTTP client. + + Returns: + The current HTTP client. + """ + + return self._http_client + + async def aclose(self) -> None: + """Closes the underlying HTTP connection pool.""" + + await self._http_client.aclose() + + async def __aenter__(self) -> Self: + return self + + async def __aexit__( + self, + exc_type: Optional[Type[BaseException]], + exc_value: Optional[BaseException], + traceback: Optional[TracebackType], + ) -> None: + await self.aclose() diff --git a/assemblyai/async_transcriber.py b/assemblyai/async_transcriber.py new file mode 100644 index 0000000..28b3aac --- /dev/null +++ b/assemblyai/async_transcriber.py @@ -0,0 +1,736 @@ +from __future__ import annotations + +import asyncio +import os +import stat +from types import TracebackType +from typing import ( + Any, + AsyncIterator, + Awaitable, + BinaryIO, + Callable, + Dict, + List, + Optional, + Tuple, + Type, + TypeVar, + Union, + cast, +) +from urllib.parse import urlparse + +import httpx +from typing_extensions import Self + +from . import async_api, types +from . import async_client as _async_client +from ._transcript_fields import TranscriptFields, config_from_response + +AudioSource = Union[str, bytes, "os.PathLike[str]", BinaryIO] +"""An audio URL, a local file path, raw `bytes`, or an opened binary file.""" + +# Read this much per thread hop when streaming a file off disk into an upload. +_UPLOAD_CHUNK_SIZE = 1024 * 1024 + +# Matches the thread-pool size the sync `Transcriber` uses for group submissions. +_DEFAULT_MAX_CONCURRENCY = 8 + +_T = TypeVar("_T") + + +async def _run_in_thread(func: Callable[..., _T], *args: Any) -> _T: + """Runs a blocking call on the default executor.""" + + loop = asyncio.get_event_loop() + + return await loop.run_in_executor(None, func, *args) + + +def _open_binary(path: str, mode: str) -> BinaryIO: + """Opens a file in binary `mode`. Call it via `_run_in_thread`.""" + + return cast(BinaryIO, open(path, mode)) + + +def _peek_length(stream: BinaryIO) -> Optional[int]: + """ + Returns the bytes left in `stream`, or `None` if the size is unknown. + + A known size lets the upload send `Content-Length`. Pipes and sockets + cannot report a size. + """ + + try: + offset = stream.tell() + except (AttributeError, OSError, ValueError): + offset = 0 + + try: + stat_result = os.fstat(stream.fileno()) + if not stat.S_ISREG(stat_result.st_mode): + # A pipe or terminal reports st_size 0, which is not its length. + raise OSError + length = stat_result.st_size + except (AttributeError, OSError, ValueError): + try: + length = stream.seek(0, os.SEEK_END) + stream.seek(offset) + except (AttributeError, OSError, ValueError): + return None + + return max(length - offset, 0) + + +async def _aiter_stream( + stream: BinaryIO, + chunk_size: int = _UPLOAD_CHUNK_SIZE, +) -> AsyncIterator[bytes]: + """Yields `stream` in chunks. Each read runs off the event loop.""" + + while True: + chunk = await _run_in_thread(stream.read, chunk_size) + if not chunk: + break + + yield chunk + + +async def _aiter_path( + path: str, + chunk_size: int = _UPLOAD_CHUNK_SIZE, +) -> AsyncIterator[bytes]: + """Yields the file at `path` in chunks. Opens and reads it off the loop.""" + + audio_file = await _run_in_thread(_open_binary, path, "rb") + try: + async for chunk in _aiter_stream(audio_file, chunk_size): + yield chunk + finally: + await _run_in_thread(audio_file.close) + + +def _upload_request( + data: AudioSource, +) -> Tuple[Union[bytes, AsyncIterator[bytes]], Dict[str, str]]: + """ + Turns audio into an upload body and the headers that describe it. + + Streams paths and file objects instead of reading them into memory. Sets + `Content-Length` when the size is known, so httpx skips chunked encoding. + + Returns: `(content, headers)`. + """ + + if isinstance(data, (bytes, bytearray)): + # httpx derives Content-Length from bytes on its own. + return bytes(data), {} + + if isinstance(data, (str, os.PathLike)): + path = os.fspath(data) + content: Union[bytes, AsyncIterator[bytes]] = _aiter_path(path) + try: + length: Optional[int] = os.path.getsize(path) + except OSError: + length = None + elif hasattr(data, "read"): + content = _aiter_stream(data) + length = _peek_length(data) + else: + raise TypeError(f"unsupported audio input type: {type(data).__name__}") + + headers = {} if length is None else {"Content-Length": str(length)} + + return content, headers + + +class AsyncTranscript(TranscriptFields, types.Sourcable): + """ + The asyncio counterpart of `Transcript`. + + Carries the same fields as `Transcript`, such as `text` and `utterances`. + Every method that calls the API is a coroutine. An `AsyncTranscriber` + creates these and owns the HTTP client they use. + """ + + def __init__( + self, + transcript_id: Optional[str], + client: _async_client.AsyncClient, + ) -> None: + """ + Creates an `AsyncTranscript` for an existing transcript id. + + Args: + transcript_id: The id of the transcript. + client: The `AsyncClient` whose connection pool to use. + """ + self._client = client + self._transcript_id = transcript_id + self._transcript: Optional[types.TranscriptResponse] = None + self._redacted_audio_url: Optional[str] = None + + @classmethod + def from_response( + cls, + *, + client: _async_client.AsyncClient, + response: types.TranscriptResponse, + ) -> Self: + self = cls(transcript_id=response.id, client=client) + self._transcript = response + + return self + + def _response(self) -> types.TranscriptResponse: + if not self._transcript: + raise ValueError("The internal Transcript object is None.") + + return self._transcript + + def _require_id(self, action: str) -> str: + if not self._transcript_id: + raise ValueError(f"Cannot {action}. The internal transcript ID is None.") + + return self._transcript_id + + @property + def id(self) -> Optional[str]: + "The unique identifier of your transcription" + + return self._transcript_id + + @property + def config(self) -> types.TranscriptionConfig: + "Return the corresponding configurations for the given transcript." + + if self._transcript is None: + raise ValueError( + "Cannot access the configuration. The internal Transcript object is None." + ) + + return config_from_response(self._transcript) + + async def wait_for_completion(self) -> Self: + """ + Polls the transcript until its status is `completed` or `error`. + + Sleeps `settings.polling_interval` seconds between polls. Other tasks + run during the sleep. + + Returns: this `AsyncTranscript`, with the finished response. + """ + + transcript_id = self._require_id("wait for completion") + + while True: + # No try-except - if there is an HTTP error then surface it to user + self._transcript = await async_api.get_transcript( + self._client.http_client, + transcript_id, + ) + + if self._transcript.status in ( + types.TranscriptStatus.completed, + types.TranscriptStatus.error, + ): + return self + + await asyncio.sleep(self._client.settings.polling_interval) + + async def export_subtitles_srt( + self, + chars_per_caption: Optional[int] = None, + ) -> str: + """ + You can export your complete transcripts in SRT format, + to be plugged into a video player for subtitles and closed captions. + + Args: + chars_per_caption: To control the maximum number of characters per caption + + Returns: A string containing the all subtitles in SRT format. + """ + + return await async_api.export_subtitles_srt( + client=self._client.http_client, + transcript_id=self._require_id("export subtitles"), + chars_per_caption=chars_per_caption, + ) + + async def export_subtitles_vtt( + self, + chars_per_caption: Optional[int] = None, + ) -> str: + """ + You can export your complete transcripts in VTT format, + to be plugged into a video player for subtitles and closed captions. + + Args: + chars_per_caption: To control the maximum number of characters per caption + + Returns: A string containing the all subtitles in VTT format. + """ + + return await async_api.export_subtitles_vtt( + client=self._client.http_client, + transcript_id=self._require_id("export subtitles"), + chars_per_caption=chars_per_caption, + ) + + async def word_search( + self, + words: List[str], + ) -> List[types.WordSearchMatch]: + """ + Once a transcript has been completed, you can search through the transcript for a specific set of keywords. + You can search for individual words, numbers, or phrases containing up to five words or numbers. + + Args: + words: A list of words, numbers, or phrases (containing up to five words or numbers) + + Returns: A list of matches + """ + + response = await async_api.word_search( + client=self._client.http_client, + transcript_id=self._require_id("perform word search"), + words=words, + ) + + return response.matches + + async def get_sentences(self) -> List[types.Sentence]: + """ + Semantically segment your transcript into sentences to create more reader-friendly transcripts. + + Returns: A list of sentence objects. + """ + + response = await async_api.get_sentences( + client=self._client.http_client, + transcript_id=self._require_id("get sentences"), + ) + + return response.sentences + + async def get_paragraphs(self) -> List[types.Paragraph]: + """ + Semantically segment your transcript into paragraphs to create more reader-friendly transcripts. + + Returns: A list of paragraph objects. + """ + + response = await async_api.get_paragraphs( + client=self._client.http_client, + transcript_id=self._require_id("get paragraphs"), + ) + + return response.paragraphs + + async def get_redacted_audio_url(self) -> str: + """ + Retrieve the URL for the PII-redacted audio file, if `redact_pii_audio` was enabled on the `TranscriptionConfig`. + Polls until the redacted audio is ready. Later calls return the cached + URL. + + Returns: The URL of the redacted audio file. + """ + + if self._redacted_audio_url is not None: + return self._redacted_audio_url + + if not self.config.redact_pii or not self.config.redact_pii_audio: + raise ValueError( + "Redacted audio is only available when `redact_pii` and `redact_pii_audio` are set to `True`." + ) + + transcript_id = self._require_id("get redacted audio url") + + while True: + try: + response = await async_api.get_redacted_audio( + client=self._client.http_client, + transcript_id=transcript_id, + ) + except types.RedactedAudioIncompleteError: + await asyncio.sleep(self._client.settings.polling_interval) + continue + + self._redacted_audio_url = response.redacted_audio_url + + return self._redacted_audio_url + + async def save_redacted_audio(self, filepath: str) -> None: + """ + Retrieve the PII-redacted audio file, if `redact_pii_audio` was enabled on the `TranscriptionConfig` + + Args: + filepath: The path to save the redacted audio file to. + """ + + url = await self.get_redacted_audio_url() + + # The redacted audio lives behind a pre-signed URL, so it is fetched + # with a bare client rather than the API-authenticated one. + async with httpx.AsyncClient() as client: + async with client.stream(method="GET", url=url) as response: + if response.status_code not in ( + httpx.codes.OK, + httpx.codes.NOT_MODIFIED, + ): + raise types.RedactedAudioUnavailableError( + f"Fetching redacted audio failed with status code {response.status_code}", + response.status_code, + ) + + audio_file = await _run_in_thread(_open_binary, filepath, "wb") + try: + async for chunk in response.aiter_bytes(): + await _run_in_thread(audio_file.write, chunk) + finally: + await _run_in_thread(audio_file.close) + + +class AsyncTranscriber: + """ + The asyncio counterpart of `Transcriber`. Transcribes URLs and local audio + files without blocking the event loop. + + Every method that calls the API is a coroutine. Many transcriptions run + concurrently on one thread, with no thread pool and no blocking + `concurrent.futures.Future`. + + The transcriber owns an HTTP connection pool. Close it with `aclose()`, or + use the transcriber as an async context manager. + + Example: + ```python + import asyncio + import assemblyai as aai + + aai.settings.api_key = "your-key" + + async def main(): + async with aai.AsyncTranscriber() as transcriber: + transcript = await transcriber.transcribe("./audio.mp3") + print(transcript.text) + + asyncio.run(main()) + ``` + + Transcribing several files concurrently is plain asyncio: + ```python + async with aai.AsyncTranscriber() as transcriber: + transcripts = await asyncio.gather( + transcriber.transcribe("./one.mp3"), + transcriber.transcribe("./two.mp3"), + ) + ``` + """ + + def __init__( + self, + *, + client: Optional[_async_client.AsyncClient] = None, + config: Optional[types.TranscriptionConfig] = None, + ) -> None: + """ + Initializes the `AsyncTranscriber` with the given parameters. + + Args: + client: The `AsyncClient` to use. If `None`, the transcriber + creates one from the global `settings` and closes it on + `aclose()`. Pass a client to share one pool between + transcribers. + config: The default configuration for the `AsyncTranscriber`. If + `None`, a default `TranscriptionConfig` is used. + """ + from . import settings as default_settings + + self._owns_client = client is None + self._client = client or _async_client.AsyncClient(settings=default_settings) + self.config = config or types.TranscriptionConfig() + + @property + def client(self) -> _async_client.AsyncClient: + """The `AsyncClient` this transcriber sends requests with.""" + + return self._client + + async def upload_file(self, data: AudioSource) -> str: + """ + Uploads an audio file, given as a local path, raw `bytes`, or a binary + object. + + Streams paths and file objects off the event loop. + + Args: + data: A local file (as path), raw `bytes`, or a binary object. + + Returns: The URL of the uploaded audio file. + """ + + content, headers = _upload_request(data) + + return await async_api.upload_file( + self._client.http_client, + content, + headers=headers, + ) + + async def submit( + self, + data: AudioSource, + config: Optional[types.TranscriptionConfig] = None, + ) -> AsyncTranscript: + """ + Submits a transcription job without waiting for its completion. + + Args: + data: An URL, a local file (as path), raw `bytes`, or a binary object. + config: Transcription options and features. If `None` is given, the + transcriber's default configuration will be used. + + Returns: The queued `AsyncTranscript`. Await `wait_for_completion()` + to poll for the result. + """ + + config = config or self.config + + if isinstance(data, str) and urlparse(data).scheme in {"http", "https"}: + audio_url = data + else: + # Note: If uploading fails, it should raise an Exception to the user. + audio_url = await self.upload_file(data) + + request = types.TranscriptRequest( + audio_url=audio_url, + **config.raw.dict(exclude_none=True), + ) + + # No try-except - if there is an HTTP error raise it to the user + response = await async_api.create_transcript( + client=self._client.http_client, + request=request, + ) + + return AsyncTranscript.from_response(client=self._client, response=response) + + async def transcribe( + self, + data: AudioSource, + config: Optional[types.TranscriptionConfig] = None, + ) -> AsyncTranscript: + """ + Transcribes an audio file and waits for the result. Accepts a local + path, a URL, raw `bytes`, or a binary object. + + Args: + data: An URL, a local file (as path), raw `bytes`, or a binary object. + config: Transcription options and features. If `None` is given, the + transcriber's default configuration will be used. + + Returns: The completed `AsyncTranscript`. Check its `status`. A + server-side failure returns `TranscriptStatus.error` and does not + raise. + """ + + transcript = await self.submit(data=data, config=config) + + return await transcript.wait_for_completion() + + async def submit_group( + self, + data: List[AudioSource], + config: Optional[types.TranscriptionConfig] = None, + return_failures: bool = False, + max_concurrency: int = _DEFAULT_MAX_CONCURRENCY, + ) -> Union[ + List[AsyncTranscript], + Tuple[List[AsyncTranscript], List[types.AssemblyAIError]], + ]: + """ + Submits multiple transcription jobs without waiting for their completion. + + Args: + data: A list of local paths, URLs, raw `bytes`, or binary objects (can be mixed). + config: Transcription options and features. If `None` is given, the + transcriber's default configuration will be used. + return_failures: Return the errors instead of raising the first one. + max_concurrency: How many submissions run at once. + + Returns: The submitted transcripts, in the order of `data`. Also returns + the errors of the failed ones when `return_failures` is set. + """ + + return await self._gather( + data, + lambda item: self.submit(data=item, config=config), + return_failures=return_failures, + max_concurrency=max_concurrency, + ) + + async def transcribe_group( + self, + data: List[AudioSource], + config: Optional[types.TranscriptionConfig] = None, + return_failures: bool = False, + max_concurrency: int = _DEFAULT_MAX_CONCURRENCY, + ) -> Union[ + List[AsyncTranscript], + Tuple[List[AsyncTranscript], List[types.AssemblyAIError]], + ]: + """ + Transcribes a list of files and waits for all of them. Accepts local + paths, URLs, raw `bytes`, and binary objects. + + Args: + data: A list of local paths, URLs, raw `bytes`, or binary objects (can be mixed). + config: Transcription options and features. If `None` is given, the + transcriber's default configuration will be used. + return_failures: Return the errors instead of raising the first one. + max_concurrency: How many transcriptions run at once. + + Returns: The completed transcripts, in the order of `data`. Also returns + the errors of the failed ones when `return_failures` is set. + """ + + return await self._gather( + data, + lambda item: self.transcribe(data=item, config=config), + return_failures=return_failures, + max_concurrency=max_concurrency, + ) + + async def _gather( + self, + data: List[AudioSource], + operation: Callable[[AudioSource], Awaitable[AsyncTranscript]], + *, + return_failures: bool, + max_concurrency: int, + ) -> Union[ + List[AsyncTranscript], + Tuple[List[AsyncTranscript], List[types.AssemblyAIError]], + ]: + """ + Runs `operation` over `data`, in order, with at most + `max_concurrency` in flight. + + Awaits every item, even after an earlier failure, so no task is left + running. A failure is never dropped: it is collected or raised once the + batch settles. + """ + + if max_concurrency < 1: + raise ValueError("max_concurrency must be at least 1") + + semaphore = asyncio.Semaphore(max_concurrency) + + async def _run(item: AudioSource) -> AsyncTranscript: + async with semaphore: + return await operation(item) + + results = await asyncio.gather( + *(_run(item) for item in data), + return_exceptions=True, + ) + + transcripts: List[AsyncTranscript] = [] + failures: List[types.AssemblyAIError] = [] + + for result in results: + if isinstance(result, BaseException): + if not return_failures or not isinstance(result, types.AssemblyAIError): + raise result + failures.append(result) + else: + transcripts.append(result) + + if return_failures: + return transcripts, failures + + return transcripts + + async def get_by_id(self, transcript_id: str) -> AsyncTranscript: + """ + Fetch an existing transcript, waiting until it is completed. + + Args: + transcript_id: the id of the transcript to fetch + + Returns: The transcript identified by the given id. + """ + + transcript = AsyncTranscript(transcript_id=transcript_id, client=self._client) + + return await transcript.wait_for_completion() + + async def delete_by_id(self, transcript_id: str) -> AsyncTranscript: + """ + Delete an existing transcript. + + Args: + transcript_id: the id of the transcript to delete + + Returns: The deleted transcript, with relevant fields/attributes cleared. + """ + + response = await async_api.delete_transcript( + client=self._client.http_client, + transcript_id=transcript_id, + ) + + return AsyncTranscript.from_response(client=self._client, response=response) + + async def list_transcripts( + self, + params: Optional[types.ListTranscriptParameters] = None, + ) -> types.ListTranscriptResponse: + """ + Retrieve a list of transcripts that were created. Transcripts are sorted from newest to oldest. + + Args: + params: The parameters to filter the transcript list by. + + Returns: A page with a list of transcripts along with page details. + + To paginate over all pages, you can set the `ListTranscriptParameters.before_id` + to the `before_id` of the `prev_url`. Example: + ``` + async with aai.AsyncTranscriber() as transcriber: + params = aai.ListTranscriptParameters() + page = await transcriber.list_transcripts(params) + while page.page_details.before_id_of_prev_url is not None: + params.before_id = page.page_details.before_id_of_prev_url + page = await transcriber.list_transcripts(params) + ``` + """ + + return await async_api.list_transcripts( + client=self._client.http_client, + params=params, + ) + + async def aclose(self) -> None: + """ + Closes the HTTP connection pool. + + Leaves a client that was passed in alone. Its creator closes it. + """ + + if self._owns_client: + await self._client.aclose() + + async def __aenter__(self) -> Self: + return self + + async def __aexit__( + self, + exc_type: Optional[Type[BaseException]], + exc_value: Optional[BaseException], + traceback: Optional[TracebackType], + ) -> None: + await self.aclose() diff --git a/assemblyai/types.py b/assemblyai/types.py index 4b44f71..676d870 100644 --- a/assemblyai/types.py +++ b/assemblyai/types.py @@ -2686,9 +2686,11 @@ def __init__( self._source = source self._type = None - from . import Transcript + from . import AsyncTranscript, Transcript - if isinstance(source, Transcript): + # LeMUR is sync-only for now, but only a source's id travels to the + # API, so an `AsyncTranscript` works just as well. + if isinstance(source, (Transcript, AsyncTranscript)): self._type = LemurSourceType.transcript else: raise ValueError(f"Invalid source: {source}") diff --git a/tests/unit/test_async_transcriber.py b/tests/unit/test_async_transcriber.py new file mode 100644 index 0000000..ca11cd2 --- /dev/null +++ b/tests/unit/test_async_transcriber.py @@ -0,0 +1,913 @@ +import asyncio +import io +import json +import os +import re +from urllib.parse import urlencode + +import httpx +import pytest +from pytest_httpx import HTTPXMock + +import assemblyai as aai +from assemblyai import async_api, types +from assemblyai.api import ENDPOINT_TRANSCRIPT, ENDPOINT_UPLOAD +from assemblyai.async_transcriber import _upload_request +from tests.unit import factories + +pytestmark = pytest.mark.asyncio + +aai.settings.api_key = "test" + +TRANSCRIPT_URL = f"{aai.settings.base_url}{ENDPOINT_TRANSCRIPT}" +UPLOAD_URL = f"{aai.settings.base_url}{ENDPOINT_UPLOAD}" + + +@pytest.fixture(autouse=True) +def fast_polling(): + """Keeps the polling loop from spending seconds sleeping in tests.""" + + original = aai.settings.polling_interval + aai.settings.polling_interval = 0.001 + yield + aai.settings.polling_interval = original + + +def _completed_response(**overrides) -> dict: + response = factories.generate_dict_factory( + factories.TranscriptCompletedResponseFactory + )() + response.update(overrides) + + return response + + +def _processing_response(transcript_id: str) -> dict: + response = factories.generate_dict_factory( + factories.TranscriptProcessingResponseFactory + )() + response["id"] = transcript_id + + return response + + +def _mock_submit(httpx_mock: HTTPXMock, response: dict) -> None: + httpx_mock.add_response( + url=TRANSCRIPT_URL, + method="POST", + status_code=httpx.codes.OK, + json=response, + ) + + +async def _drain(content) -> bytes: + """Collects an upload body, which is either bytes or an async iterable.""" + + if isinstance(content, bytes): + return content + + return b"".join([chunk async for chunk in content]) + + +def _stub_create_transcript(monkeypatch, handler) -> None: + """ + Replaces the create-transcript request with `handler`. + + The group methods only orchestrate `submit`, so a stub at the transport + boundary tests the order and the concurrency limit without a mock HTTP + server. pytest-httpx cannot delay a response on every supported version. + """ + + monkeypatch.setattr(async_api, "create_transcript", handler) + + +def _mock_poll(httpx_mock: HTTPXMock, response: dict, **kwargs) -> None: + httpx_mock.add_response( + url=f"{TRANSCRIPT_URL}/{response['id']}", + method="GET", + status_code=httpx.codes.OK, + json=response, + **kwargs, + ) + + +async def test_transcribe_url_submits_and_polls(httpx_mock: HTTPXMock): + # Given a job that is queued on submission and completed when polled + completed = _completed_response() + _mock_submit(httpx_mock, _processing_response(completed["id"])) + _mock_poll(httpx_mock, completed) + + # When transcribing a URL + async with aai.AsyncTranscriber() as transcriber: + transcript = await transcriber.transcribe("https://example.org/audio.wav") + + # Then the completed transcript is returned + assert isinstance(transcript, aai.AsyncTranscript) + assert transcript.id == completed["id"] + assert transcript.status == aai.TranscriptStatus.completed + assert transcript.text == completed["text"] + assert transcript.words is not None + assert transcript.utterances is not None + + # And the URL was submitted as-is, without an upload + submission = json.loads(httpx_mock.get_requests()[0].read()) + assert submission["audio_url"] == "https://example.org/audio.wav" + assert len(httpx_mock.get_requests()) == 2 + + +async def test_transcribe_passes_config(httpx_mock: HTTPXMock): + # Given a completed job + completed = _completed_response() + _mock_submit(httpx_mock, completed) + + # When transcribing with a config + config = aai.TranscriptionConfig( + speech_models=["universal-3-5-pro"], + speaker_labels=True, + ) + async with aai.AsyncTranscriber(config=config) as transcriber: + await transcriber.submit("https://example.org/audio.wav") + + # Then the config travels in the submission body + submission = json.loads(httpx_mock.get_requests()[0].read()) + assert submission["speech_models"] == ["universal-3-5-pro"] + assert submission["speaker_labels"] is True + + +async def test_submit_does_not_poll(httpx_mock: HTTPXMock): + # Given a job that comes back as processing + transcript_id = "some-id" + _mock_submit(httpx_mock, _processing_response(transcript_id)) + + # When submitting without waiting + async with aai.AsyncTranscriber() as transcriber: + transcript = await transcriber.submit("https://example.org/audio.wav") + + # Then only the submission request was made + assert transcript.status == aai.TranscriptStatus.processing + assert len(httpx_mock.get_requests()) == 1 + + +async def test_wait_for_completion_polls_until_terminal(httpx_mock: HTTPXMock): + # Given a job that stays queued for two polls before completing + completed = _completed_response() + _mock_submit(httpx_mock, _processing_response(completed["id"])) + _mock_poll(httpx_mock, _processing_response(completed["id"])) + _mock_poll(httpx_mock, _processing_response(completed["id"])) + _mock_poll(httpx_mock, completed) + + # When transcribing + async with aai.AsyncTranscriber() as transcriber: + transcript = await transcriber.transcribe("https://example.org/audio.wav") + + # Then it polled until the status was terminal + assert transcript.status == aai.TranscriptStatus.completed + assert len(httpx_mock.get_requests()) == 4 + + +async def test_transcribe_surfaces_error_status(httpx_mock: HTTPXMock): + # Given a job that fails server-side + error_response = factories.generate_dict_factory( + factories.TranscriptErrorResponseFactory + )() + error_response["id"] = "error-id" + _mock_submit(httpx_mock, _processing_response("error-id")) + _mock_poll(httpx_mock, error_response) + + # When transcribing + async with aai.AsyncTranscriber() as transcriber: + transcript = await transcriber.transcribe("https://example.org/audio.wav") + + # Then the failure surfaces on the transcript rather than as an exception + assert transcript.status == aai.TranscriptStatus.error + assert transcript.error == "Aw, snap!" + + +async def test_submit_raises_on_http_error(httpx_mock: HTTPXMock): + # Given a submission that is rejected + httpx_mock.add_response( + url=TRANSCRIPT_URL, + method="POST", + status_code=httpx.codes.BAD_REQUEST, + json={"error": "something went wrong"}, + ) + + # When submitting, then a TranscriptError carries the server message + async with aai.AsyncTranscriber() as transcriber: + with pytest.raises(aai.TranscriptError) as exc_info: + await transcriber.submit("https://example.org/audio.wav") + + assert "something went wrong" in str(exc_info.value) + + +async def test_upload_file_from_bytes(httpx_mock: HTTPXMock): + # Given a mocked upload endpoint + audio = os.urandom(64) + httpx_mock.add_response( + url=UPLOAD_URL, + method="POST", + status_code=httpx.codes.OK, + json={"upload_url": "https://example.org/uploaded.wav"}, + match_content=audio, + ) + + # When uploading raw bytes + async with aai.AsyncTranscriber() as transcriber: + upload_url = await transcriber.upload_file(audio) + + # Then the audio is posted verbatim + assert upload_url == "https://example.org/uploaded.wav" + + +async def test_upload_file_from_path_streams_with_content_length( + httpx_mock: HTTPXMock, + tmp_path, +): + # Given a local audio file larger than one read chunk + audio = os.urandom(3 * 1024 * 1024) + audio_path = tmp_path / "audio.wav" + audio_path.write_bytes(audio) + + httpx_mock.add_response( + url=UPLOAD_URL, + method="POST", + status_code=httpx.codes.OK, + json={"upload_url": "https://example.org/uploaded.wav"}, + ) + + # When uploading it by path + async with aai.AsyncTranscriber() as transcriber: + upload_url = await transcriber.upload_file(str(audio_path)) + + # Then the request is sized rather than chunk-encoded + assert upload_url == "https://example.org/uploaded.wav" + request = httpx_mock.get_requests()[0] + assert request.headers["content-length"] == str(len(audio)) + assert "transfer-encoding" not in request.headers + + +async def test_upload_file_from_file_object(httpx_mock: HTTPXMock): + # Given an in-memory binary file object + audio = os.urandom(128) + httpx_mock.add_response( + url=UPLOAD_URL, + method="POST", + status_code=httpx.codes.OK, + json={"upload_url": "https://example.org/uploaded.wav"}, + ) + + # When uploading it + async with aai.AsyncTranscriber() as transcriber: + await transcriber.upload_file(io.BytesIO(audio)) + + # Then its length is known before the request starts + request = httpx_mock.get_requests()[0] + assert request.headers["content-length"] == str(len(audio)) + assert "transfer-encoding" not in request.headers + + +async def test_upload_file_from_partially_read_file_object(httpx_mock: HTTPXMock): + # Given a file object that has already been read from + audio = os.urandom(128) + stream = io.BytesIO(audio) + stream.read(28) + + httpx_mock.add_response( + url=UPLOAD_URL, + method="POST", + status_code=httpx.codes.OK, + json={"upload_url": "https://example.org/uploaded.wav"}, + ) + + # When uploading it + async with aai.AsyncTranscriber() as transcriber: + await transcriber.upload_file(stream) + + # Then Content-Length counts the remaining bytes only + request = httpx_mock.get_requests()[0] + assert request.headers["content-length"] == str(len(audio) - 28) + + +async def test_upload_file_from_pathlike(httpx_mock: HTTPXMock, tmp_path): + # Given a local file addressed by a `pathlib.Path` + audio = os.urandom(64) + audio_path = tmp_path / "audio.wav" + audio_path.write_bytes(audio) + + httpx_mock.add_response( + url=UPLOAD_URL, + method="POST", + status_code=httpx.codes.OK, + json={"upload_url": "https://example.org/uploaded.wav"}, + ) + + # When uploading it + async with aai.AsyncTranscriber() as transcriber: + await transcriber.upload_file(audio_path) + + request = httpx_mock.get_requests()[0] + assert request.headers["content-length"] == str(len(audio)) + + +async def test_upload_request_sends_bytes_unchanged(): + # Given raw audio bytes + audio = os.urandom(64) + + # When building the upload body + content, headers = _upload_request(audio) + + # Then httpx receives the bytes and derives Content-Length itself + assert await _drain(content) == audio + assert headers == {} + + +async def test_upload_request_streams_a_path(tmp_path): + # Given a file larger than one read chunk + audio = os.urandom(3 * 1024 * 1024) + audio_path = tmp_path / "audio.wav" + audio_path.write_bytes(audio) + + # When building the upload body from a path and from a PathLike + for source in (str(audio_path), audio_path): + content, headers = _upload_request(source) + + # Then every chunk arrives, in order, with a Content-Length + assert await _drain(content) == audio + assert headers == {"Content-Length": str(len(audio))} + + +async def test_upload_request_streams_a_file_object(): + # Given a binary file object + audio = os.urandom(128) + + # When building the upload body + content, headers = _upload_request(io.BytesIO(audio)) + + assert await _drain(content) == audio + assert headers == {"Content-Length": str(len(audio))} + + +async def test_upload_request_sends_the_remainder_of_a_read_file_object(): + # Given a file object that has already been read from + audio = os.urandom(128) + stream = io.BytesIO(audio) + stream.read(28) + + # When building the upload body + content, headers = _upload_request(stream) + + # Then only the remaining bytes are sent + assert await _drain(content) == audio[28:] + assert headers == {"Content-Length": str(len(audio) - 28)} + + +async def test_upload_request_omits_content_length_for_an_unsized_stream(): + # Given a stream that cannot report its size + audio = os.urandom(64) + read_fd, write_fd = os.pipe() + os.write(write_fd, audio) + os.close(write_fd) + + # When building the upload body + with open(read_fd, "rb") as pipe: + content, headers = _upload_request(pipe) + + # Then httpx falls back to chunked transfer encoding + assert headers == {} + assert await _drain(content) == audio + + +async def test_upload_file_rejects_unsupported_input(): + # When uploading something that is neither bytes, a path, nor a file object + async with aai.AsyncTranscriber() as transcriber: + with pytest.raises(TypeError): + await transcriber.upload_file(42) # type: ignore[arg-type] + + +async def test_transcribe_local_file_uploads_then_submits( + httpx_mock: HTTPXMock, + tmp_path, +): + # Given a local file and a completed job + audio_path = tmp_path / "audio.wav" + audio_path.write_bytes(os.urandom(32)) + + completed = _completed_response() + httpx_mock.add_response( + url=UPLOAD_URL, + method="POST", + status_code=httpx.codes.OK, + json={"upload_url": "https://example.org/uploaded.wav"}, + ) + _mock_submit(httpx_mock, _processing_response(completed["id"])) + _mock_poll(httpx_mock, completed) + + # When transcribing the local file + async with aai.AsyncTranscriber() as transcriber: + transcript = await transcriber.transcribe(str(audio_path)) + + # Then the uploaded URL is what gets submitted + assert transcript.status == aai.TranscriptStatus.completed + submission = json.loads(httpx_mock.get_requests()[1].read()) + assert submission["audio_url"] == "https://example.org/uploaded.wav" + + +async def test_get_sentences_and_paragraphs(httpx_mock: HTTPXMock): + # Given a completed transcript with sentences and paragraphs available + completed = _completed_response() + transcript_id = completed["id"] + + sentences = factories.generate_dict_factory(factories.SentencesResponseFactory)() + paragraphs = factories.generate_dict_factory(factories.ParagraphsResponseFactory)() + + _mock_submit(httpx_mock, completed) + httpx_mock.add_response( + url=f"{TRANSCRIPT_URL}/{transcript_id}/sentences", + method="GET", + status_code=httpx.codes.OK, + json=sentences, + ) + httpx_mock.add_response( + url=f"{TRANSCRIPT_URL}/{transcript_id}/paragraphs", + method="GET", + status_code=httpx.codes.OK, + json=paragraphs, + ) + + # When requesting both + async with aai.AsyncTranscriber() as transcriber: + transcript = await transcriber.submit("https://example.org/audio.wav") + got_sentences = await transcript.get_sentences() + got_paragraphs = await transcript.get_paragraphs() + + # Then they are parsed into the corresponding models + assert [s.text for s in got_sentences] == [ + s["text"] for s in sentences["sentences"] + ] + assert [p.text for p in got_paragraphs] == [ + p["text"] for p in paragraphs["paragraphs"] + ] + + +async def test_export_subtitles(httpx_mock: HTTPXMock): + # Given a completed transcript whose subtitle endpoints return text + completed = _completed_response() + transcript_id = completed["id"] + + _mock_submit(httpx_mock, completed) + httpx_mock.add_response( + url=f"{TRANSCRIPT_URL}/{transcript_id}/srt", + method="GET", + status_code=httpx.codes.OK, + text="srt-subtitles", + ) + httpx_mock.add_response( + url=f"{TRANSCRIPT_URL}/{transcript_id}/vtt?chars_per_caption=32", + method="GET", + status_code=httpx.codes.OK, + text="vtt-subtitles", + ) + + # When exporting subtitles, with and without a caption limit + async with aai.AsyncTranscriber() as transcriber: + transcript = await transcriber.submit("https://example.org/audio.wav") + srt = await transcript.export_subtitles_srt() + vtt = await transcript.export_subtitles_vtt(chars_per_caption=32) + + assert srt == "srt-subtitles" + assert vtt == "vtt-subtitles" + + +async def test_word_search(httpx_mock: HTTPXMock): + # Given a completed transcript and a word-search result + completed = _completed_response() + matches = factories.generate_dict_factory( + factories.WordSearchMatchResponseFactory + )() + + _mock_submit(httpx_mock, completed) + httpx_mock.add_response( + url=f"{TRANSCRIPT_URL}/{completed['id']}/word-search?{urlencode({'words': 'foo,bar'})}", + method="GET", + status_code=httpx.codes.OK, + json=matches, + ) + + # When searching for words + async with aai.AsyncTranscriber() as transcriber: + transcript = await transcriber.submit("https://example.org/audio.wav") + found = await transcript.word_search(["foo", "bar"]) + + assert [m.text for m in found] == [m["text"] for m in matches["matches"]] + + +async def test_list_transcripts(httpx_mock: HTTPXMock): + # Given a page of transcripts + page = factories.generate_dict_factory(factories.ListTranscriptResponse)() + httpx_mock.add_response( + url=re.compile(rf"^{re.escape(TRANSCRIPT_URL)}\?.*"), + method="GET", + status_code=httpx.codes.OK, + json=page, + ) + + # When listing them with parameters + async with aai.AsyncTranscriber() as transcriber: + result = await transcriber.list_transcripts( + aai.ListTranscriptParameters(limit=2) + ) + + assert isinstance(result, aai.ListTranscriptResponse) + assert len(result.transcripts) == len(page["transcripts"]) + assert httpx_mock.get_requests()[0].url.params["limit"] == "2" + + +async def test_get_by_id_waits_for_completion(httpx_mock: HTTPXMock): + # Given an existing transcript that is still processing + completed = _completed_response() + _mock_poll(httpx_mock, _processing_response(completed["id"])) + _mock_poll(httpx_mock, completed) + + # When fetching it by id + async with aai.AsyncTranscriber() as transcriber: + transcript = await transcriber.get_by_id(completed["id"]) + + assert transcript.status == aai.TranscriptStatus.completed + assert transcript.text == completed["text"] + + +async def test_delete_by_id(httpx_mock: HTTPXMock): + # Given a transcript that gets deleted + deleted = factories.generate_dict_factory( + factories.TranscriptDeletedResponseFactory + )() + httpx_mock.add_response( + url=f"{TRANSCRIPT_URL}/{deleted['id']}", + method="DELETE", + status_code=httpx.codes.OK, + json=deleted, + ) + + # When deleting it + async with aai.AsyncTranscriber() as transcriber: + transcript = await transcriber.delete_by_id(deleted["id"]) + + assert transcript.id == deleted["id"] + assert transcript.text == "Deleted by user." + + +async def test_transcribe_group_preserves_order(monkeypatch): + # Given three URLs whose jobs finish in a different order + urls = [f"https://example.org/{i}.wav" for i in range(3)] + responses = [ + types.TranscriptResponse.parse_obj(_completed_response(text=f"transcript {i}")) + for i in range(3) + ] + delays = [0.03, 0.01, 0.02] + + async def create_transcript(*, client, request): + index = urls.index(request.audio_url) + await asyncio.sleep(delays[index]) + + return responses[index] + + _stub_create_transcript(monkeypatch, create_transcript) + + # When submitting them as a group + async with aai.AsyncTranscriber() as transcriber: + transcripts = await transcriber.submit_group(urls) + + # Then the results follow the input order, not the completion order + assert [t.text for t in transcripts] == [ + "transcript 0", + "transcript 1", + "transcript 2", + ] + + +async def test_transcribe_group_raises_first_failure(monkeypatch): + # Given a batch in which every submission is rejected + async def create_transcript(*, client, request): + raise aai.TranscriptError("nope", 400) + + _stub_create_transcript(monkeypatch, create_transcript) + + # When submitting the group without asking for the failures + async with aai.AsyncTranscriber() as transcriber: + with pytest.raises(aai.TranscriptError): + await transcriber.submit_group( + ["https://example.org/a.wav", "https://example.org/b.wav"] + ) + + +async def test_transcribe_group_returns_failures(monkeypatch): + # Given a batch where one submission fails and one succeeds + completed = types.TranscriptResponse.parse_obj(_completed_response(text="ok")) + + async def create_transcript(*, client, request): + if request.audio_url.endswith("bad.wav"): + raise aai.TranscriptError("nope", 400) + + return completed + + _stub_create_transcript(monkeypatch, create_transcript) + + # When submitting with return_failures + async with aai.AsyncTranscriber() as transcriber: + transcripts, failures = await transcriber.submit_group( + ["https://example.org/bad.wav", "https://example.org/good.wav"], + return_failures=True, + ) + + # Then the successful transcript and the error come back together + assert [t.text for t in transcripts] == ["ok"] + assert len(failures) == 1 + assert isinstance(failures[0], aai.TranscriptError) + + +async def test_transcribe_group_limits_concurrency(monkeypatch): + # Given a batch of six jobs and a transport that counts concurrent calls + completed = types.TranscriptResponse.parse_obj(_completed_response()) + in_flight = 0 + peak = 0 + + async def create_transcript(*, client, request): + nonlocal in_flight, peak + in_flight += 1 + peak = max(peak, in_flight) + try: + await asyncio.sleep(0.01) + finally: + in_flight -= 1 + + return completed + + _stub_create_transcript(monkeypatch, create_transcript) + + # When submitting them with a concurrency limit of two + async with aai.AsyncTranscriber() as transcriber: + transcripts = await transcriber.submit_group( + [f"https://example.org/{i}.wav" for i in range(6)], + max_concurrency=2, + ) + + # Then two ran at a time, and all six finished + assert len(transcripts) == 6 + assert peak == 2 + + +async def test_transcriptions_run_concurrently(monkeypatch): + # Given a transport that takes 50ms per call + completed = types.TranscriptResponse.parse_obj(_completed_response()) + + async def create_transcript(*, client, request): + await asyncio.sleep(0.05) + + return completed + + _stub_create_transcript(monkeypatch, create_transcript) + + # When submitting four jobs together on one thread + async with aai.AsyncTranscriber() as transcriber: + loop = asyncio.get_event_loop() + started = loop.time() + await asyncio.gather( + *(transcriber.submit(f"https://example.org/{i}.wav") for i in range(4)) + ) + elapsed = loop.time() - started + + # Then they overlapped instead of running in sequence (4 x 50ms) + assert elapsed < 0.15 + + +async def test_transcribe_group_rejects_invalid_concurrency(): + async with aai.AsyncTranscriber() as transcriber: + with pytest.raises(ValueError): + await transcriber.submit_group( + ["https://example.org/a.wav"], max_concurrency=0 + ) + + +async def test_properties_require_a_fetched_response(): + # Given a transcript that has not been fetched + async with aai.AsyncTranscriber() as transcriber: + transcript = aai.AsyncTranscript( + transcript_id="some-id", client=transcriber.client + ) + + # Then reading its fields fails loudly + with pytest.raises(ValueError): + transcript.text + + with pytest.raises(ValueError): + transcript.config + + # But its id is still available + assert transcript.id == "some-id" + + +async def test_operations_require_a_transcript_id(): + # Given a transcript with no id at all + async with aai.AsyncTranscriber() as transcriber: + transcript = aai.AsyncTranscript(transcript_id=None, client=transcriber.client) + + with pytest.raises(ValueError): + await transcript.wait_for_completion() + + with pytest.raises(ValueError): + await transcript.get_sentences() + + +async def test_config_reflects_submitted_options(httpx_mock: HTTPXMock): + # Given a transcript that was created with speaker labels + completed = _completed_response(speaker_labels=True) + _mock_submit(httpx_mock, completed) + + # When reading the config back off the transcript + async with aai.AsyncTranscriber() as transcriber: + transcript = await transcriber.submit("https://example.org/audio.wav") + + assert transcript.config.speaker_labels is True + + +async def test_redacted_audio_polls_until_ready(httpx_mock: HTTPXMock): + # Given a transcript with PII audio redaction enabled + completed = _completed_response( + redact_pii=True, + redact_pii_audio=True, + redact_pii_policies=["person_name"], + ) + transcript_id = completed["id"] + redacted_url = "https://example.org/redacted.wav" + + _mock_submit(httpx_mock, completed) + httpx_mock.add_response( + url=f"{TRANSCRIPT_URL}/{transcript_id}/redacted-audio", + method="GET", + status_code=httpx.codes.ACCEPTED, + json={}, + ) + httpx_mock.add_response( + url=f"{TRANSCRIPT_URL}/{transcript_id}/redacted-audio", + method="GET", + status_code=httpx.codes.OK, + json={"status": "redacted_audio_ready", "redacted_audio_url": redacted_url}, + ) + + # When asking for the redacted audio URL twice + async with aai.AsyncTranscriber() as transcriber: + transcript = await transcriber.submit("https://example.org/audio.wav") + first = await transcript.get_redacted_audio_url() + second = await transcript.get_redacted_audio_url() + + # Then it polled past the 202 and cached the result + assert first == redacted_url + assert second == redacted_url + assert len(httpx_mock.get_requests()) == 3 + + +async def test_redacted_audio_requires_redaction_config(httpx_mock: HTTPXMock): + # Given a transcript that was not configured for PII audio redaction + _mock_submit(httpx_mock, _completed_response()) + + # Then asking for redacted audio is rejected before any request is made + async with aai.AsyncTranscriber() as transcriber: + transcript = await transcriber.submit("https://example.org/audio.wav") + + with pytest.raises(ValueError): + await transcript.get_redacted_audio_url() + + +async def test_save_redacted_audio(httpx_mock: HTTPXMock, tmp_path): + # Given a transcript whose redacted audio is ready + completed = _completed_response( + redact_pii=True, + redact_pii_audio=True, + redact_pii_policies=["person_name"], + ) + redacted_url = "https://example.org/redacted.wav" + audio = os.urandom(2048) + + _mock_submit(httpx_mock, completed) + httpx_mock.add_response( + url=f"{TRANSCRIPT_URL}/{completed['id']}/redacted-audio", + method="GET", + status_code=httpx.codes.OK, + json={"status": "redacted_audio_ready", "redacted_audio_url": redacted_url}, + ) + httpx_mock.add_response( + url=redacted_url, + method="GET", + status_code=httpx.codes.OK, + content=audio, + ) + + # When saving it to disk + target = tmp_path / "redacted.wav" + async with aai.AsyncTranscriber() as transcriber: + transcript = await transcriber.submit("https://example.org/audio.wav") + await transcript.save_redacted_audio(str(target)) + + assert target.read_bytes() == audio + + +async def test_save_redacted_audio_raises_when_unavailable( + httpx_mock: HTTPXMock, + tmp_path, +): + # Given a redacted audio URL that no longer serves the file + completed = _completed_response( + redact_pii=True, + redact_pii_audio=True, + redact_pii_policies=["person_name"], + ) + redacted_url = "https://example.org/redacted.wav" + + _mock_submit(httpx_mock, completed) + httpx_mock.add_response( + url=f"{TRANSCRIPT_URL}/{completed['id']}/redacted-audio", + method="GET", + status_code=httpx.codes.OK, + json={"status": "redacted_audio_ready", "redacted_audio_url": redacted_url}, + ) + httpx_mock.add_response( + url=redacted_url, + method="GET", + status_code=httpx.codes.NOT_FOUND, + ) + + async with aai.AsyncTranscriber() as transcriber: + transcript = await transcriber.submit("https://example.org/audio.wav") + + with pytest.raises(aai.types.RedactedAudioUnavailableError): + await transcript.save_redacted_audio(str(tmp_path / "redacted.wav")) + + +async def test_sends_authorization_and_user_agent(httpx_mock: HTTPXMock): + # Given any request + _mock_submit(httpx_mock, _completed_response()) + + async with aai.AsyncTranscriber() as transcriber: + await transcriber.submit("https://example.org/audio.wav") + + # Then it carries the same auth and user-agent as the sync client + request = httpx_mock.get_requests()[0] + assert request.headers["authorization"] == "test" + assert "AssemblyAI/1.0" in request.headers["user-agent"] + + +async def test_client_is_closed_when_owned(httpx_mock: HTTPXMock): + # Given a transcriber that created its own client + async with aai.AsyncTranscriber() as transcriber: + client = transcriber.client + + # Then leaving the context closed the pool + assert client.http_client.is_closed + + +async def test_supplied_client_is_not_closed(): + # Given a transcriber handed an explicit client + client = aai.AsyncClient(settings=aai.settings) + + async with aai.AsyncTranscriber(client=client) as transcriber: + assert transcriber.client is client + + # Then closing the transcriber leaves the caller's pool open + assert not client.http_client.is_closed + await client.aclose() + assert client.http_client.is_closed + + +async def test_client_records_last_response(httpx_mock: HTTPXMock): + # Given a completed submission + _mock_submit(httpx_mock, _completed_response()) + + async with aai.AsyncTranscriber() as transcriber: + assert transcriber.client.last_response is None + await transcriber.submit("https://example.org/audio.wav") + + # Then the client exposes the last response, like the sync client does + assert transcriber.client.last_response is not None + assert transcriber.client.last_response.status_code == httpx.codes.OK + + +async def test_client_requires_an_api_key(): + # Given settings without an API key + settings = aai.Settings(api_key=None) + + with pytest.raises(ValueError): + aai.AsyncClient(settings=settings) + + +async def test_async_transcript_is_a_lemur_source(httpx_mock: HTTPXMock): + # Given a completed async transcript + completed = _completed_response() + _mock_submit(httpx_mock, completed) + + async with aai.AsyncTranscriber() as transcriber: + transcript = await transcriber.submit("https://example.org/audio.wav") + + # Then it can be handed to LeMUR, which only needs its id + source = aai.LemurSource(transcript) + + assert source.source.id == completed["id"]