diff --git a/CLAUDE.md b/CLAUDE.md index a2bbc688..a1f8b5c2 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -48,6 +48,10 @@ aai.settings.api_key = "your-key" - `aai.AsyncSyncTranscriber` — Asyncio counterpart of `SyncTranscriber`. Same input types, config, result, and errors; `transcribe()`, `transcribe_live()` and `warm()` are coroutines; `open_live()` returns an `AsyncLiveSession`. Owns an HTTP pool: use `async with` or `await aclose()`, or pass an `aai.AsyncClient` to share one - `aai.SyncTranscriptionConfig` — Sync options: `model` (default `universal-3-5-pro`), `prompt`, `keyterms_prompt`, `conversation_context`, `language_codes`, `timestamps`, `sample_rate`, `channels` - `aai.SyncTranscriptResponse` — Sync result: `.text`, `.words` (`SyncWord` with `confidence` always, `start`/`end` only when `timestamps=True`), `.confidence`, `.audio_duration_ms`, `.session_id`, `.request_time_ms` +- `aai.DictationTranscriber` — Dictation API: audio uploaded as it is spoken over one live request, transcript (and optional LLM pass) out. Methods: `open_live()` (push-style `DictationLiveSession`: `write()`/`close()`/`result()`/`abort()`), `transcribe_live()` (iterable, file object, bytes or path), `warm()`. There is no buffered `transcribe()`: every entry point posts to `/v1/transcribe/live` +- `aai.AsyncDictationTranscriber` — Asyncio counterpart of `DictationTranscriber`; `open_live()` returns an `AsyncDictationLiveSession`. Owns an HTTP pool: use `async with` or `await aclose()` +- `aai.DictationConfig` — Dictation options: `sample_rate`, `channels`, `language_codes`, `keyterms_prompt`, `llm_instruction`. Rejects unknown fields +- `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`) - `assemblyai.streaming.v3.RealTimeTranscriber` — Real-time streaming with event-based API (threaded) - `assemblyai.streaming.v3.AsyncRealTimeTranscriber` — Asyncio-native counterpart; same options/events diff --git a/README.md b/README.md index e4119bfa..b7f6d5bc 100644 --- a/README.md +++ b/README.md @@ -104,10 +104,12 @@ aai.settings.api_key = f"{ASSEMBLYAI_API_KEY}" | `assemblyai.prerecorded.v2.AsyncTranscriber` | The same, from asyncio code | | `assemblyai.sync.v1.SyncTranscriber` | Short clips (≤120s, ≤40MB) where you want the transcript back in one request, at the lowest latency — from a file you already have (`transcribe()`) or uploaded while it is still being recorded (`open_live()`, `transcribe_live()`) | | `assemblyai.sync.v1.AsyncSyncTranscriber` | The same, from asyncio code | +| `assemblyai.dictation.v1.DictationTranscriber` | Dictation: short spoken notes uploaded as they are spoken, returned as one transcript with an optional LLM rewrite (`final_text`). Push audio from a callback (`open_live()`) or hand over an iterator, file object, bytes or path (`transcribe_live()`) | +| `assemblyai.dictation.v1.AsyncDictationTranscriber` | The same, from asyncio code | | `assemblyai.streaming.v3.RealTimeTranscriber` | Live audio (microphone, telephony, voice agents), transcribed as it arrives over a websocket session | | `assemblyai.streaming.v3.AsyncRealTimeTranscriber` | The same, from asyncio code | -The versioned path is the preferred import for new code. The prerecorded and sync classes are also available as top-level shortcuts (`aai.Transcriber`, `aai.SyncTranscriber`, …) — see [Migrating to 1.0](MIGRATION.md) for the details. +The versioned path is the preferred import for new code. The prerecorded, sync and dictation classes are also available as top-level shortcuts (`aai.Transcriber`, `aai.SyncTranscriber`, `aai.DictationTranscriber`, …) — see [Migrating to 1.0](MIGRATION.md) for the details. --- @@ -677,6 +679,205 @@ except aai.SyncTranscriptError as error: --- +### **Dictation Examples** + +`aai.DictationTranscriber` is for dictation: a person speaks a short note and gets it back as text, optionally cleaned up or reformatted by an LLM. It has one connection, the live upload — audio is sent as it is spoken, the service transcribes each speech segment as it lands, and when the speaker stops what remains is the final segment and the LLM pass. There is no job id and no polling. Audio must be WAV or raw 16-bit PCM, up to 120 s per request. + +The result is a `DictationResponse`: `text` is the raw transcript, `llm_response` the rewritten text when the LLM pass ran, and `final_text` is the one to show the user — the rewrite when there is one, the raw transcript otherwise. + +
+ Dictate from a microphone (`open_live()`) + +Most audio sources deliver chunks in a callback. `open_live()` is built for that: the request starts the moment the session opens, `write()` from the callback (safe from any thread, never blocks), leave the `with` block when the speaker stops, and `result()` returns the transcript. + +```python +import assemblyai as aai +import sounddevice as sd # pip install sounddevice + +aai.settings.api_key = "" + +RATE = 16000 +config = aai.DictationConfig(sample_rate=RATE, channels=1) # raw PCM needs both + +with aai.DictationTranscriber() as transcriber: + with transcriber.open_live(config) as session: + microphone = sd.RawInputStream( + samplerate=RATE, channels=1, dtype="int16", + callback=lambda data, *_: session.write(bytes(data)), + ) + with microphone: + input("Dictating, press Enter to stop... ") + # leaving the block ends the audio; the final segment is all that is left + print(session.result().final_text) +``` + +Leaving the `with` block calls `close()`. If the block raises, the session is aborted instead and `result()` raises. Call `abort()` yourself to drop a dictation you no longer want transcribed. + +This is **not** the streaming API. A session returns one finished transcript when the audio ends; use `aai.RealTimeTranscriber` when you need words back while the speaker is still talking. + +The server aborts an upload that goes silent for long, so keep writing until the speaker is done rather than pausing. Errors that would normally arrive at the end (bad key, rate limit, capacity, audio too large) can surface part-way through the upload; `result()` raises them as `DictationError`. + +
+ +
+ Dictate from an iterator, file object, bytes or path (`transcribe_live()`) + +When the audio already comes as something you can iterate, or as a file object that fills over time, hand it to `transcribe_live()` directly. It blocks until the audio ends and returns the transcript. `open_live()` is this method with a queue in front of it. + +Audio you already hold whole — raw `bytes`, or a local path — goes to `transcribe_live()` too. It travels the same connection as a single chunk; there is no separate buffered request. + +```python +import subprocess +import assemblyai as aai + +aai.settings.api_key = "" +config = aai.DictationConfig(sample_rate=16000, channels=1) + +# A file object that fills over time, e.g. the stdout of a recorder (sox): +recorder = subprocess.Popen( + ["rec", "-q", "-t", "raw", "-r", "16000", "-c", "1", "-b", "16", "-e", "signed", "-"], + stdout=subprocess.PIPE, +) +result = aai.DictationTranscriber().transcribe_live(recorder.stdout, config=config) + +# Or any generator of chunks: +def from_call(call): + while call.active: + yield call.read_frame() + +result = aai.DictationTranscriber().transcribe_live(from_call(call), config=config) + +# Or a recording you already have (WAV carries its own rate and channels): +result = aai.DictationTranscriber().transcribe_live("./note.wav") +print(result.final_text) +``` + +End the iterator to finish. Anything the producer raises propagates unchanged and drops the upload. URLs are not accepted; use `aai.Transcriber` for URL ingestion. + +
+ +
+ Rewrite the transcript with an LLM instruction + +Set `llm_instruction` to have the service run a follow-up pass over the transcript. The rewrite comes back in `llm_response`; the raw transcript stays in `text`. If the pass fails, `llm_error` says why and `final_text` falls back to `text`, so reading `final_text` is safe either way. + +```python +import assemblyai as aai + +aai.settings.api_key = "" + +config = aai.DictationConfig(llm_instruction="Format this as a SOAP note.") + +result = aai.DictationTranscriber().transcribe_live("./visit.wav", config=config) + +print(result.final_text) # the SOAP note +print(result.text) # what was actually said +if result.llm_error: + print("LLM pass failed:", result.llm_error) +``` + +The instruction is capped at 2048 characters. + +
+ +
+ Configure the dictation + +```python +import assemblyai as aai + +aai.settings.api_key = "" + +config = aai.DictationConfig( + language_codes=["en", "es"], # one code, or several for multilingual audio + keyterms_prompt=["AssemblyAI", "Universal"], # bias the decoder; ≤2048 chars total + llm_instruction="Fix punctuation only.", # optional LLM pass +) + +result = aai.DictationTranscriber().transcribe_live("./note.wav", config=config) +``` + +`sample_rate` and `channels` are required only for raw PCM; WAV carries them in its own header. A config passed to `transcribe_live()` or `open_live()` overrides the transcriber's default config for that call. `DictationConfig` rejects unknown fields, so a typo or a sync-only option surfaces as a validation error instead of a setting that quietly does nothing. + +
+ +
+ Pre-warm the connection + +A request that connects on demand pays the full DNS + TCP + TLS handshake before the first audio byte can leave. Call `warm()` as soon as you know audio is coming — when the user reaches for the record button — so the request that follows reuses the open connection. + +```python +import assemblyai as aai + +aai.settings.api_key = "" + +with aai.DictationTranscriber() as transcriber: + transcriber.warm() # fire when the user is about to speak + with transcriber.open_live(config) as session: + record_into(session.write) + print(session.result().final_text) +``` + +
+ +
+ Use it from asyncio (`AsyncDictationTranscriber`) + +`aai.AsyncDictationTranscriber` is the asyncio counterpart — same audio sources, config, result and errors, with `transcribe_live()` and `warm()` as coroutines. `open_live()` returns an `AsyncDictationLiveSession`: `write()` and `close()` are plain functions so a callback can call them, and `result()` and `abort()` are coroutines. The request runs as a task on the event loop. `write()` must be called on the loop's thread; from an audio library's capture thread, schedule it with `loop.call_soon_threadsafe(session.write, chunk)`. + +```python +import asyncio +import assemblyai as aai + +aai.settings.api_key = "" +config = aai.DictationConfig(sample_rate=16000, channels=1) + + +async def handle(websocket): + """Transcribes the PCM frames a browser sends over one websocket.""" + async with aai.AsyncDictationTranscriber() as transcriber: + async with transcriber.open_live(config) as session: + async for frame in websocket: + session.write(frame) + result = await session.result() + await websocket.send(result.final_text) +``` + +`transcribe_live()` also takes an async iterable of chunks, as well as a file object, plain iterable, bytes or path. A file object or path is read in a worker thread, so a blocking `read()` is fine; a plain iterable is consumed inline and must not block the loop. + +```python +async def frames(websocket): + async for frame in websocket: + yield frame + +result = await transcriber.transcribe_live(frames(websocket), config=config) +``` + +The transcriber owns an HTTP connection pool: use `async with`, or call `await transcriber.aclose()`. To share one pool between transcribers, pass an `aai.AsyncClient`, which stays yours to close. + +
+ +
+ Handle errors + +Failures raise `aai.DictationError` with the HTTP `status_code`, a machine-readable `error_code` (`bad_audio`, `audio_too_large`, `capacity_exceeded`, `inference_timeout`, …), and `retry_after` (seconds) on 429/503 responses. It is a separate class from `aai.SyncTranscriptError`. + +```python +import assemblyai as aai + +aai.settings.api_key = "" + +try: + result = aai.DictationTranscriber().transcribe_live("./note.wav") + print(result.final_text) +except aai.DictationError as error: + print(error.status_code, error.error_code, error.retry_after) +``` + +
+ +--- + ### **Asyncio Examples** `aai.AsyncTranscriber` is the asyncio counterpart of `aai.Transcriber`. Every method @@ -1393,6 +1594,11 @@ aai.settings.polling_interval = 10.0 # need to cover the length of the recording. aai.settings.sync_http_timeout = 60.0 aai.settings.sync_live_http_timeout = 180.0 + +# The Dictation API host and its per-operation timeout (default 300.0), sized +# to outlast the final segment's inference plus the LLM pass. +aai.settings.dictation_base_url = "https://dictation.assemblyai.com" +aai.settings.dictation_http_timeout = 300.0 ``` diff --git a/assemblyai/__init__.py b/assemblyai/__init__.py index 982f0ceb..3e680ae2 100644 --- a/assemblyai/__init__.py +++ b/assemblyai/__init__.py @@ -1,6 +1,14 @@ from .__version__ import __version__ from .async_client import AsyncClient from .client import Client +from .dictation.v1 import ( + AsyncDictationTranscriber, + DictationConfig, + DictationError, + DictationResponse, + DictationTranscriber, + DictationWord, +) from .prerecorded.v2 import AsyncTranscriber, AsyncTranscript from .sync.v1 import AsyncSyncTranscriber, SyncTranscriber from .transcriber import Transcriber, Transcript, TranscriptGroup @@ -96,6 +104,7 @@ def __getattr__(name: str): # types "AssemblyAIError", "AsyncClient", + "AsyncDictationTranscriber", "AsyncSyncTranscriber", "AsyncTranscriber", "AsyncTranscript", @@ -109,6 +118,11 @@ def __getattr__(name: str): "ContentSafetySeverityScore", "CustomFormattingRequest", "CustomFormattingResponse", + "DictationConfig", + "DictationError", + "DictationResponse", + "DictationTranscriber", + "DictationWord", "Entity", "EntityType", "IABLabelResult", diff --git a/assemblyai/__version__.py b/assemblyai/__version__.py index 3e8d9f94..5b601886 100644 --- a/assemblyai/__version__.py +++ b/assemblyai/__version__.py @@ -1 +1 @@ -__version__ = "1.4.0" +__version__ = "1.5.0" diff --git a/assemblyai/_audio.py b/assemblyai/_audio.py new file mode 100644 index 00000000..e0bab838 --- /dev/null +++ b/assemblyai/_audio.py @@ -0,0 +1,170 @@ +"""Audio-input helpers shared by the single-request transcription products. + +Both the sync API and the Dictation API take audio as a multipart upload +rather than a URL, so they read the same input shapes (path, bytes, binary +file object), decide a multipart `Content-Type` the same way, and serialize +their config to the same JSON part. The product packages wrap these with +their own config type and error wording. +""" + +from __future__ import annotations + +import os +from typing import Any, BinaryIO, Mapping, Optional, Tuple, Type, Union +from urllib.parse import urlparse + +AudioInput = Union[str, bytes, bytearray, "os.PathLike[str]", BinaryIO] + +# Extensions that signal raw S16LE PCM rather than a container format. +_PCM_SUFFIXES = (".pcm", ".raw") + +# Content-Type used when the extension is unknown or absent. +_DEFAULT_CONTENT_TYPE = "audio/wav" + + +def check_config(owner: str, config: Optional[Any], expected_type: Type[Any]) -> None: + """ + Raises unless `config` is an instance of `expected_type` or `None`. + + The job API's `TranscriptionConfig` is a different, non-interchangeable + type, so passing it here is a mistake worth naming. + + Args: + owner: the class to name in the message, e.g. `SyncTranscriber`. + config: the configuration to check. + expected_type: the configuration class `owner` accepts. + """ + + if config is not None and not isinstance(config, expected_type): + raise TypeError( + f"{owner} expects {expected_type.__name__}, got {type(config).__name__}. " + f"Use aai.{expected_type.__name__}." + ) + + +def _resolve_audio( + data: AudioInput, + *, + sample_rate: Optional[int], + channels: Optional[int], + owner: str, + config_name: str, + content_types: Mapping[str, str], +) -> Tuple[bytes, str, str]: + """ + Reads the audio input into bytes and decides its multipart Content-Type. + + PCM is selected when the source has a `.pcm`/`.raw` extension or when + `sample_rate`/`channels` are set on the config (the fields these APIs + require only for raw PCM) — and both must then be present. Any other + input takes the Content-Type `content_types` maps its extension to, and + `audio/wav` when the extension is unknown or absent. URLs are rejected — + neither API has URL ingestion. + + Args: + data: a local file path, raw bytes, or a binary file object. + sample_rate: the config's `sample_rate`, or `None`. + channels: the config's `channels`, or `None`. + owner: the class to name in error messages, e.g. `SyncTranscriber`. + config_name: the config class to name in error messages. + content_types: extension (lowercased, with the dot) → Content-Type. + + Returns: `(audio_bytes, filename, content_type)`. + """ + suffix = "" + filename: Optional[str] = None + + if isinstance(data, (bytes, bytearray)): + audio = bytes(data) + elif isinstance(data, (str, os.PathLike)): + path = os.fspath(data) + if urlparse(path).scheme in ("http", "https"): + raise ValueError( + f"{owner} does not accept URLs. Pass a local file path or " + "audio bytes, or use aai.Transcriber for URL/async transcription." + ) + with open(path, "rb") as f: + audio = f.read() + filename = os.path.basename(path) + suffix = os.path.splitext(path)[1].lower() + elif hasattr(data, "read"): + audio = data.read() + name = getattr(data, "name", None) + if name: + filename = os.path.basename(name) + suffix = os.path.splitext(name)[1].lower() + else: + raise TypeError(f"unsupported audio input type: {type(data).__name__}") + + resolved_filename, content_type = resolve_format( + suffix=suffix, + filename=filename, + sample_rate=sample_rate, + channels=channels, + config_name=config_name, + content_types=content_types, + ) + + return audio, resolved_filename, content_type + + +def resolve_format( + *, + suffix: str, + filename: Optional[str], + sample_rate: Optional[int], + channels: Optional[int], + config_name: str, + content_types: Mapping[str, str], +) -> Tuple[str, str]: + """ + Decides the multipart filename and Content-Type for the audio part. + + PCM is selected when `suffix` is a PCM extension or when `sample_rate` / + `channels` are set — the fields these APIs require only for raw PCM — and + both must then be present. Any other suffix takes the Content-Type + `content_types` maps it to, and `audio/wav` when it is unknown or absent. + Needs no audio bytes, so it serves a streamed upload as well as a buffered + one. + + Args: + suffix: the source's lowercased file extension (with the dot), or "". + filename: the name for the multipart part; defaulted when absent. + sample_rate: the config's `sample_rate`, or `None`. + channels: the config's `channels`, or `None`. + config_name: the config class to name in error messages. + content_types: extension (lowercased, with the dot) → Content-Type. + + Returns: `(filename, content_type)`. + """ + wants_pcm = sample_rate is not None or channels is not None + is_pcm = suffix in _PCM_SUFFIXES or wants_pcm + if is_pcm and (sample_rate is None or channels is None): + raise ValueError( + f"raw PCM audio requires both sample_rate and channels in {config_name}" + ) + + if is_pcm: + content_type = "audio/pcm" + else: + content_type = content_types.get(suffix, _DEFAULT_CONTENT_TYPE) + if not filename: + filename = "audio.pcm" if is_pcm else "audio.wav" + + return filename, content_type + + +def _config_to_json(config: Any, *, exclude: Tuple[str, ...] = ()) -> Optional[dict]: + """ + Serializes the config to the JSON `config` part. + + Args: + config: the pydantic config model. + exclude: field names to drop — options carried outside the body. + + Returns: the JSON-ready dict, or `None` when nothing is set. + """ + data = config.dict(exclude_none=True) + for field in exclude: + data.pop(field, None) + return data or None diff --git a/assemblyai/sync/v1/_multipart.py b/assemblyai/_multipart.py similarity index 80% rename from assemblyai/sync/v1/_multipart.py rename to assemblyai/_multipart.py index 7ad83786..e9e95fef 100644 --- a/assemblyai/sync/v1/_multipart.py +++ b/assemblyai/_multipart.py @@ -1,7 +1,15 @@ -"""Multipart encoding for streamed uploads. - -Kept apart from `_base.py` so the transport modules can use it: `_base.py` -imports `api.py`, so anything `api.py` needs must not live there. +"""Multipart encoding for streamed uploads, shared by the single-request products. + +The buffered upload hands httpx a ``files=`` mapping and lets it encode the +whole body up front. A streamed upload cannot: the audio does not exist yet +when the request starts. This module emits the ``multipart/form-data`` framing +itself so the ``config`` part can go out immediately and audio can follow as +it arrives, and it turns whatever the caller produces — a file object, an +iterable, an async iterable — into the byte chunks that framing carries. + +Product-neutral by design: the sync API and the Dictation API stream the same +body shape to their live endpoints, so each product package wraps this with +its own endpoint, config type and error class. """ from __future__ import annotations @@ -21,8 +29,7 @@ ) # Audio for a streamed upload, delivered in whatever pieces the producer has -# ready. A file object is read lazily. Audio the caller already holds whole -# belongs in `transcribe()`, which is faster for it. +# ready. A file object is read lazily, in bounded pieces. AudioChunks = Union[Iterable[bytes], BinaryIO] AsyncAudioChunks = Union[AsyncIterable[bytes], Iterable[bytes], BinaryIO] @@ -46,20 +53,24 @@ def _escape_form_param(value: str) -> str: return _FORM_PARAM_RE.sub(lambda m: _FORM_PARAM_REPLACEMENTS[m.group(0)], value) +class _Aborted(Exception): + """ + Raised inside a live session's producer when the session is aborted. + + Propagates out of the transport, which drops the connection, and is + swallowed by the session's `abort()`; it never reaches the caller. + """ + + class StreamingMultipartEncoder: """ Builds a `multipart/form-data` body one part at a time. - The buffered path hands httpx a `files=` mapping and lets it encode the - whole body up front. A streamed upload cannot: the audio does not exist yet - when the request starts. This emits the framing itself so the `config` part - can go out immediately and audio can follow as it arrives. - **Part order is significant and is enforced here by construction.** The server decodes audio as it lands, so it needs `sample_rate` and `channels` before the first audio byte; `config` therefore precedes `audio`. Audio first is rejected by the server rather than buffered, which is the whole - point of the endpoint. + point of a live endpoint. """ def __init__(self, boundary: Optional[str] = None) -> None: @@ -80,9 +91,9 @@ def config_part(self, config: Optional[dict]) -> bytes: """ Encodes the JSON `config` part. - Always emitted: the streaming endpoint rejects a body whose audio is - not preceded by a `config` part, so when the caller set no options an - empty object is sent. The buffered route omits the part instead. + Always emitted: a live endpoint rejects a body whose audio is not + preceded by a `config` part, so when the caller set no options an + empty object is sent. """ return ( @@ -148,7 +159,7 @@ def iter_chunks(data: AudioChunks) -> Iterator[bytes]: if hasattr(data, "read"): while True: - chunk = data.read(_STREAM_READ_SIZE) + chunk = data.read(_STREAM_READ_SIZE) # type: ignore[union-attr] if not chunk: return yield _as_bytes(chunk) diff --git a/assemblyai/dictation/__init__.py b/assemblyai/dictation/__init__.py new file mode 100644 index 00000000..53f74425 --- /dev/null +++ b/assemblyai/dictation/__init__.py @@ -0,0 +1,34 @@ +"""The Dictation API: audio in as it is spoken, transcript (and optional LLM +rewrite) out. + +The implementation lives in ``dictation/v1/``, mirroring the ``sync/v1/`` and +``streaming/v3/`` layout. This ``__init__`` re-exports the public surface so +``from assemblyai.dictation import DictationTranscriber`` works alongside the +versioned path. +""" + +from .v1 import ( + AsyncAudioSource, + AsyncDictationLiveSession, + AsyncDictationTranscriber, + AudioSource, + DictationConfig, + DictationError, + DictationLiveSession, + DictationResponse, + DictationTranscriber, + DictationWord, +) + +__all__ = [ + "AsyncAudioSource", + "AsyncDictationLiveSession", + "AsyncDictationTranscriber", + "AudioSource", + "DictationConfig", + "DictationError", + "DictationLiveSession", + "DictationResponse", + "DictationTranscriber", + "DictationWord", +] diff --git a/assemblyai/dictation/v1/__init__.py b/assemblyai/dictation/v1/__init__.py new file mode 100644 index 00000000..2a5c8b2f --- /dev/null +++ b/assemblyai/dictation/v1/__init__.py @@ -0,0 +1,22 @@ +from ._base import AsyncAudioSource, AudioSource +from .async_client import AsyncDictationLiveSession, AsyncDictationTranscriber +from .client import DictationLiveSession, DictationTranscriber +from .models import ( + DictationConfig, + DictationError, + DictationResponse, + DictationWord, +) + +__all__ = [ + "AsyncAudioSource", + "AsyncDictationLiveSession", + "AsyncDictationTranscriber", + "AudioSource", + "DictationConfig", + "DictationError", + "DictationLiveSession", + "DictationResponse", + "DictationTranscriber", + "DictationWord", +] diff --git a/assemblyai/dictation/v1/_base.py b/assemblyai/dictation/v1/_base.py new file mode 100644 index 00000000..9ff2b25c --- /dev/null +++ b/assemblyai/dictation/v1/_base.py @@ -0,0 +1,187 @@ +from __future__ import annotations + +import os +from typing import AsyncIterable, BinaryIO, Iterable, Mapping, Optional, Tuple, Union +from urllib.parse import urlparse + +from ... import _audio +from ... import client as _client +from ..._multipart import AsyncAudioChunks, AudioChunks +from . import api +from .models import DictationConfig, DictationResponse + +__all__ = [ + "AsyncAudioSource", + "AudioSource", + "_DictationTranscriberImpl", + "_config_to_json", + "check_config", + "resolve_source", +] + +# Audio for a dictation request. The live endpoint is the only one this +# client opens, so it takes every shape audio comes in: a producer that is +# still recording (an iterable of chunks, or a file object read as it fills) +# and audio that is already complete (bytes, or a local path), which goes out +# over the same connection as a single chunk. +AudioSource = Union[ + str, bytes, bytearray, "os.PathLike[str]", BinaryIO, Iterable[bytes] +] +AsyncAudioSource = Union[AudioSource, AsyncIterable[bytes]] + +# The Dictation API decodes raw PCM and the common container formats; the +# extension picks the Content-Type that tells the server which decoder to use. +# The live route accepts WAV and PCM today: a compressed format is posted with +# its true type so the server's rejection names it, rather than mislabelled as +# WAV and failing to decode. +_CONTENT_TYPES: Mapping[str, str] = { + ".wav": "audio/wav", + ".pcm": "audio/pcm", + ".raw": "audio/pcm", + ".mp3": "audio/mpeg", + ".aac": "audio/aac", + ".mp4": "audio/mp4", + ".m4a": "audio/x-m4a", + ".ogg": "audio/ogg", + ".opus": "audio/opus", + ".flac": "audio/flac", + ".webm": "audio/webm", +} + + +def check_config(owner: str, config: Optional[DictationConfig]) -> None: + """ + Raises unless `config` is a `DictationConfig` or `None`. + + `SyncTranscriptionConfig` and the job API's `TranscriptionConfig` are + different, non-interchangeable types, so passing one here is a mistake + worth naming. + + Args: + owner: the class to name in the message, e.g. `DictationTranscriber`. + config: the configuration to check. + """ + + _audio.check_config(owner, config, DictationConfig) + + +def _config_to_json(config: DictationConfig) -> Optional[dict]: + """Serializes the config to the JSON `config` part.""" + + return _audio._config_to_json(config) + + +def resolve_source( + data: object, + config: DictationConfig, + owner: str, + *, + allow_async: bool = False, +) -> Tuple[Union[AudioChunks, AsyncAudioChunks], str, str]: + """ + Turns any accepted audio source into chunks for the live upload, and + decides the audio part's filename and Content-Type. + + Complete audio — `bytes`, or a local path, which is read here — becomes a + single chunk. A file object or an iterable is handed on untouched and read + as the upload proceeds. PCM is selected when the source has a `.pcm`/`.raw` + extension or when `sample_rate`/`channels` are set on the config (the + fields the Dictation API requires only for raw PCM) — and both must then be + present. Any other extension maps to its container Content-Type, and audio + with no usable extension is posted as WAV. URLs are rejected — the + Dictation API has no URL ingestion. + + Args: + data: the audio source. + config: the configuration for this call. + owner: the class to name in error messages, so each client reports + its own name. + allow_async: whether an object exposing only `__aiter__` is acceptable, + i.e. whether the caller is `AsyncDictationTranscriber`. + + Returns: `(chunks, filename, content_type)`. + + Raises: + ValueError: for a URL, or raw PCM missing `sample_rate` or `channels`. + TypeError: for a source of an unsupported type, or an async iterable + handed to the threaded transcriber. + """ + suffix = "" + filename: Optional[str] = None + chunks: Union[AudioChunks, AsyncAudioChunks] + + if isinstance(data, (bytes, bytearray)): + chunks = (bytes(data),) + elif isinstance(data, (str, os.PathLike)): + path = os.fspath(data) + if urlparse(path).scheme in ("http", "https"): + raise ValueError( + f"{owner} does not accept URLs. Pass a local file path, audio " + "bytes, a binary file object or an iterable of audio chunks, or " + "use aai.Transcriber for URL transcription." + ) + with open(path, "rb") as f: + chunks = (f.read(),) + filename = os.path.basename(path) + suffix = os.path.splitext(path)[1].lower() + elif hasattr(data, "read"): + chunks = data # type: ignore[assignment] + name = getattr(data, "name", None) + if isinstance(name, str) and name: + filename = os.path.basename(name) + suffix = os.path.splitext(name)[1].lower() + elif hasattr(data, "__iter__"): + chunks = data # type: ignore[assignment] + elif hasattr(data, "__aiter__"): + if not allow_async: + raise TypeError( + "DictationTranscriber.transcribe_live() cannot consume an async " + "iterable. Use AsyncDictationTranscriber.transcribe_live(), or " + "hand it a plain iterable or file object." + ) + chunks = data # type: ignore[assignment] + else: + raise TypeError(f"unsupported audio source type: {type(data).__name__}") + + resolved_filename, content_type = _audio.resolve_format( + suffix=suffix, + filename=filename, + sample_rate=config.sample_rate, + channels=config.channels, + config_name="DictationConfig", + content_types=_CONTENT_TYPES, + ) + + return chunks, resolved_filename, content_type + + +class _DictationTranscriberImpl: + def __init__( + self, + *, + client: _client.Client, + config: DictationConfig, + owner: str, + ) -> None: + self._client = client + self.config = config + self._owner = owner + + def transcribe_live( + self, + *, + data: AudioSource, + config: Optional[DictationConfig], + ) -> DictationResponse: + config = config or self.config + chunks, filename, content_type = resolve_source(data, config, self._owner) + + return api.transcribe_live( + self._client.http_client, + base_url=self._client.settings.dictation_base_url, + chunks=chunks, # type: ignore[arg-type] + filename=filename, + audio_content_type=content_type, + config=_config_to_json(config), + timeout=self._client.settings.dictation_http_timeout, + ) diff --git a/assemblyai/dictation/v1/api.py b/assemblyai/dictation/v1/api.py new file mode 100644 index 00000000..65934bf0 --- /dev/null +++ b/assemblyai/dictation/v1/api.py @@ -0,0 +1,102 @@ +from typing import Iterator, Optional + +import httpx + +from ... import types +from ..._multipart import AudioChunks, StreamingMultipartEncoder, iter_chunks +from ...sync.v1.api import _error_from_response as _shared_error_from_response +from .models import DictationError, DictationResponse + +# The one endpoint this client posts audio to. The Dictation API serves it +# under /v1 only (no unprefixed alias), and also at `/v1/transcribe/stream`, +# the path it shipped under — but `live` is its name. +ENDPOINT_TRANSCRIBE_LIVE = "/v1/transcribe/live" +ENDPOINT_WARM = "/v1/warm" + + +def _error_from_response(response: httpx.Response) -> types.AssemblyAIError: + """ + Builds a `DictationError` from a non-200 response. + + The Dictation API returns the same error envelopes as the sync API — an + RFC 9457 problem-details body (`{"status", "title", "detail"}`), + `{"error", "error_code"}`, or a bare `{"detail"}` — so the parsing is + shared and only the exception class differs. + """ + + return _shared_error_from_response( + response, + DictationError, + "dictation transcription", + ) + + +def transcribe_live( + client: httpx.Client, + *, + base_url: str, + chunks: AudioChunks, + filename: str, + audio_content_type: str, + config: Optional[dict], + timeout: float, +) -> DictationResponse: + """ + Posts a dictation request whose audio is uploaded as it arrives. + + Sends the body with chunked transfer encoding — httpx frames an unsized + iterator that way — so the request can start before the audio exists. The + server transcribes each speech segment as it lands, leaving only the final + segment's inference (and the LLM pass, when one was asked for) to wait on + once the caller stops speaking. Audio that is already complete travels the + same way, as a single chunk: there is one request shape in this client. + + Args: + client: the HTTP client (carries the `Authorization` header). + base_url: the Dictation API base URL, e.g. `https://dictation.assemblyai.com`. + chunks: audio pieces (WAV container or S16LE PCM), in order. + filename: name for the audio multipart part. + audio_content_type: the audio part's Content-Type; selects the decoder. + config: the JSON `config` part. None sends an empty object: the live + endpoint requires the part ahead of the audio. + timeout: per-operation timeout in seconds, as for every httpx request: + it bounds connecting, each socket write and each read while waiting + for the response, not the request end to end. Time blocked in the + caller's producer is not counted, so it need not cover the + recording. + + Returns: the parsed dictation response. + + Raises: `DictationError` on any non-200 response — including one the + server sends while the upload is still in flight, which it may do for + auth, rate-limit, size and capacity failures. `TypeError` if the + producer yields something other than bytes. Any other exception the + producer raises propagates unchanged; the connection is dropped and no + transcript is returned. + """ + encoder = StreamingMultipartEncoder() + + def body() -> Iterator[bytes]: + # config first: the server needs sample_rate and channels before it can + # decode a single audio byte, and rejects audio that arrives first. + yield encoder.config_part(config) + encoder.audio_header( + filename, audio_content_type + ) + + for chunk in iter_chunks(chunks): + if chunk: + yield chunk + + yield encoder.closing() + + response = client.post( + base_url.rstrip("/") + ENDPOINT_TRANSCRIBE_LIVE, + content=body(), + headers={"Content-Type": encoder.content_type}, + timeout=timeout, + ) + + if response.status_code != httpx.codes.OK: + raise _error_from_response(response) + + return DictationResponse.parse_obj(response.json()) diff --git a/assemblyai/dictation/v1/async_api.py b/assemblyai/dictation/v1/async_api.py new file mode 100644 index 00000000..426eded1 --- /dev/null +++ b/assemblyai/dictation/v1/async_api.py @@ -0,0 +1,74 @@ +"""The asyncio counterpart of `api.py`. + +Calls the same endpoint as its sync twin and raises the same `DictationError` +through `api._error_from_response`. +""" + +from typing import AsyncIterator, Optional + +import httpx + +from ..._multipart import AsyncAudioChunks, StreamingMultipartEncoder, aiter_chunks +from .api import ENDPOINT_TRANSCRIBE_LIVE, _error_from_response +from .models import DictationResponse + +__all__ = ["transcribe_live"] + + +async def transcribe_live( + client: httpx.AsyncClient, + *, + base_url: str, + chunks: AsyncAudioChunks, + filename: str, + audio_content_type: str, + config: Optional[dict], + timeout: float, +) -> DictationResponse: + """ + Posts a dictation request whose audio is uploaded as it arrives. + + The asyncio counterpart of `api.transcribe_live`; same endpoint, same + chunked framing, same errors. + + Args: + client: the HTTP client (carries the `Authorization` header). + base_url: the Dictation API base URL, e.g. `https://dictation.assemblyai.com`. + chunks: audio pieces (WAV container or S16LE PCM), in order. + filename: name for the audio multipart part. + audio_content_type: the audio part's Content-Type; selects the decoder. + config: the JSON `config` part. None sends an empty object: the live + endpoint requires the part ahead of the audio. + timeout: per-operation timeout in seconds; see `api.transcribe_live`. + + Returns: the parsed dictation response. + + Raises: `DictationError` on any non-200 response — including one the + server sends while the upload is still in flight. + """ + encoder = StreamingMultipartEncoder() + + async def body() -> AsyncIterator[bytes]: + # config first: the server needs sample_rate and channels before it can + # decode a single audio byte, and rejects audio that arrives first. + yield encoder.config_part(config) + encoder.audio_header( + filename, audio_content_type + ) + + async for chunk in aiter_chunks(chunks): + if chunk: + yield chunk + + yield encoder.closing() + + response = await client.post( + base_url.rstrip("/") + ENDPOINT_TRANSCRIBE_LIVE, + content=body(), + headers={"Content-Type": encoder.content_type}, + timeout=timeout, + ) + + if response.status_code != httpx.codes.OK: + raise _error_from_response(response) + + return DictationResponse.parse_obj(response.json()) diff --git a/assemblyai/dictation/v1/async_client.py b/assemblyai/dictation/v1/async_client.py new file mode 100644 index 00000000..b4a381a2 --- /dev/null +++ b/assemblyai/dictation/v1/async_client.py @@ -0,0 +1,408 @@ +"""The asyncio counterpart of `client.py`.""" + +from __future__ import annotations + +import asyncio +import functools +from types import TracebackType +from typing import Any, AsyncIterator, Callable, Optional, Type, TypeVar + +import httpx +from typing_extensions import Self + +from ... import async_client as _async_client +from ..._multipart import _Aborted, _as_bytes +from . import api, async_api +from ._base import AsyncAudioSource, _config_to_json, check_config, resolve_source +from .models import DictationConfig, DictationResponse + +_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) + + +class AsyncDictationLiveSession: + """ + A live dictation that audio is pushed into, for asyncio code. + + Returned by `AsyncDictationTranscriber.open_live()`. The request runs as a + task on the current event loop from the moment the session opens; + `write()` hands it audio, `close()` ends the audio, and `await result()` + waits for the transcript. Built for callback-driven sources — a WebRTC + track, a websocket handler receiving frames, a telephony media stream — + where the audio arrives in a callback rather than from an async iterator + you can hand to `transcribe_live()`. + + `write()` and `close()` are plain functions so a callback can call them, + and must run on the event loop's thread. From another thread — an audio + library's capture thread, say — schedule them with + `loop.call_soon_threadsafe(session.write, chunk)`. + + Everything `transcribe_live()` says about the 120 s audio cap, the + server-side silence abort and errors surfacing mid-upload applies here + unchanged. + + Example: + ```python + async with aai.AsyncDictationTranscriber() as transcriber: + async with transcriber.open_live(config) as session: + async for frame in websocket: # audio frames from a browser + session.write(frame) + print((await session.result()).final_text) + ``` + """ + + def __init__( + self, + start: Callable[[AsyncIterator[bytes]], asyncio.Task[DictationResponse]], + ) -> None: + """ + Args: + start: begins the upload from the given producer and returns the + task that will hold its outcome. Supplied by the transcriber. + """ + self._queue: "asyncio.Queue[Optional[bytes]]" = asyncio.Queue() + self._closed = False + self._aborted = False + self._task = start(self._chunks()) + + async def _chunks(self) -> AsyncIterator[bytes]: + while True: + chunk = await self._queue.get() + if chunk is None: + if self._aborted: + raise _Aborted() + return + yield chunk + + @property + def closed(self) -> bool: + """Whether the audio has ended, by `close()`, `result()` or `abort()`.""" + + return self._closed + + def write(self, chunk: bytes) -> None: + """ + Queues a piece of audio for upload. Never blocks. + + Must be called on the event loop's thread; see the class docstring for + calling from elsewhere. + + Args: + chunk: `bytes`, `bytearray` or `memoryview` of audio, in order. + + Raises: + TypeError: if `chunk` is not bytes-like. + RuntimeError: if the session is closed. + """ + if self._closed: + raise RuntimeError( + "the live session is closed; no more audio can be written" + ) + + self._queue.put_nowait(_as_bytes(chunk)) + + def close(self) -> None: + """ + Ends the audio. The server transcribes what was sent; `result()` + returns it. Idempotent, and never blocks. + """ + if not self._closed: + self._closed = True + self._queue.put_nowait(None) + + async def result(self) -> DictationResponse: + """ + Waits for the transcript, ending the audio first if it is still open. + + Raises: + DictationError: if the request failed, including a rejection the + server sent while the upload was still in flight. + RuntimeError: if the session was aborted. + """ + if self._aborted: + raise RuntimeError("the live session was aborted; there is no result") + + self.close() + + return await self._task + + async def abort(self) -> None: + """ + Drops the request without a transcript. + + The connection is closed and nothing the server may still return is + kept; `result()` raises afterwards. Idempotent, and a no-op once the + request has completed. Waits until the task has let go of the + connection. After `close()` the upload is already complete, so + aborting then waits out the in-flight response instead. + """ + if self._aborted or self._task.done(): + return + + self._aborted = True + if not self._closed: + self._closed = True + self._queue.put_nowait(None) + + try: + await self._task + except Exception: + pass + + async def __aenter__(self) -> "AsyncDictationLiveSession": + return self + + async def __aexit__( + self, + exc_type: Optional[Type[BaseException]], + exc_value: Optional[BaseException], + traceback: Optional[TracebackType], + ) -> None: + """Ends the audio on a clean exit; aborts if the block raised.""" + + if exc_type is not None: + await self.abort() + else: + self.close() + + +class AsyncDictationTranscriber: + """ + The asyncio counterpart of `DictationTranscriber`: audio in, transcript + out, one request — without blocking the event loop. + + Same audio sources (an async or plain iterable of chunks, a file object, + bytes or a local path — no URLs), same `DictationConfig`, same + `DictationResponse` and `DictationError`, with `transcribe_live()` and + `warm()` as coroutines and `open_live()` returning an + `AsyncDictationLiveSession`. Use it in asyncio code (FastAPI, aiohttp, + voice agents), where the threaded `DictationTranscriber` would block the + loop. + + 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.AsyncDictationTranscriber() as transcriber: + async with transcriber.open_live() as session: + async for frame in websocket: + session.write(frame) + print((await session.result()).final_text) + + asyncio.run(main()) + ``` + """ + + def __init__( + self, + *, + client: Optional[_async_client.AsyncClient] = None, + config: Optional[DictationConfig] = None, + api_key: Optional[str] = None, + ) -> None: + """ + Creates an `AsyncDictationTranscriber`. + + 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: Default dictation options. Per-call `config` overrides it. + api_key: The API key to authenticate with. The transcriber builds + its own `AsyncClient` from it and closes that client on + `aclose()`. Given alongside `client`, it takes precedence: the + transcriber builds and owns a client made from a copy of that + client's settings with the key replaced, and the given client is + left untouched and stays the caller's to close. + + Raises: + TypeError: if `config` is not a `DictationConfig`. + """ + check_config(type(self).__name__, config) + + self._owns_client = client is None or api_key is not None + self._client = _async_client._resolve_client(client, api_key) + self.config = config or DictationConfig() + + @property + def client(self) -> _async_client.AsyncClient: + """The `AsyncClient` this transcriber sends requests with.""" + + return self._client + + async def transcribe_live( + self, + data: AsyncAudioSource, + config: Optional[DictationConfig] = None, + ) -> DictationResponse: + """ + Transcribes audio uploaded as it is produced. + + The asyncio counterpart of `DictationTranscriber.transcribe_live`. + Starts the request immediately and uploads chunks as they arrive, so + authorization, the upload and every speech segment but the last + resolve while the caller is still recording. Audio that is already + complete — bytes, or a local path, which is read off the event loop — + is accepted as well and sent as a single chunk over the same + connection. + + The caller must keep producing: an upload that goes silent for long + enough is aborted server-side. Stop by ending the iterator, not by + pausing it. The service caps a request at 120 s of audio. + + Args: + data: An async iterable of audio chunks, a plain iterable, a file + object, raw audio bytes, or a local file path. A file object is + read in a worker thread, so a blocking `read()` is fine; a + plain iterable is consumed inline and so must not block. Raw + PCM also requires `sample_rate` and `channels` on the config. + config: Options for this call. If `None`, the transcriber's default + configuration is used. + + Raises: + TypeError: if `config` is not a `DictationConfig`; if `data` is of + an unsupported type; or if a chunk is not bytes (a file opened + in text mode, say). + ValueError: for a URL, or raw PCM without both `sample_rate` and + `channels`. + DictationError: if the request fails. Auth, rate-limit, size and + capacity failures can surface part-way through the upload. + Exception: anything the producer raises mid-upload propagates + unchanged. The connection is dropped and no transcript is + returned. + + Example: + ```python + async def mic_chunks(): + while recording: + yield await stream.read(4096) + + async with aai.AsyncDictationTranscriber() as transcriber: + result = await transcriber.transcribe_live(mic_chunks()) + ``` + """ + check_config(type(self).__name__, config) + + config = config or self.config + # Resolving may read a whole file for a path source; keep that off + # the loop. + chunks, filename, content_type = await _run_in_thread( + functools.partial( + resolve_source, data, config, type(self).__name__, allow_async=True + ) + ) + + return await async_api.transcribe_live( + self._client.http_client, + base_url=self._client.settings.dictation_base_url, + chunks=chunks, + filename=filename, + audio_content_type=content_type, + config=_config_to_json(config), + timeout=self._client.settings.dictation_http_timeout, + ) + + def open_live( + self, + config: Optional[DictationConfig] = None, + ) -> AsyncDictationLiveSession: + """ + Opens a live dictation that audio is pushed into. + + The push-style counterpart of `transcribe_live()`, for sources that + deliver audio through a callback rather than an async iterator. The + request starts as a task on the running event loop immediately; call + `session.write(chunk)` as audio arrives, then `await session.result()` + for the transcript once the speaker stops. See + `AsyncDictationLiveSession`. + + Not a coroutine, so it can be used directly as `async with + transcriber.open_live(config) as session:`. Must be called while the + event loop is running. + + Args: + config: Options for this call. If `None`, the transcriber's default + configuration is used. Raw PCM requires `sample_rate` and + `channels`. + + Raises: + TypeError: if `config` is not a `DictationConfig`. + RuntimeError: if no event loop is running. + """ + check_config(type(self).__name__, config) + loop = asyncio.get_running_loop() + + return AsyncDictationLiveSession( + lambda chunks: loop.create_task(self.transcribe_live(chunks, config=config)) + ) + + async def warm(self) -> bool: + """ + Opens the connection to the Dictation API ahead of time. + + A request that opens its connection on demand pays the full DNS + TCP + + TLS handshake before the first audio byte can leave — one network + round trip that, for a distant client, is a noticeable share of a + short dictation. Awaiting `warm()` as soon as you know audio is coming + — e.g. via `asyncio.create_task(transcriber.warm())` when the user + reaches for the record button — spends that setup early: the next + request reuses the already-open connection. + + The warmed connection is reused while it stays in the HTTP pool — + `settings.keepalive_expiry` seconds (httpx's 5s default unless + raised). Call `warm()` shortly before the request, or raise + `keepalive_expiry` so a single call covers a longer pause. `warm()` is + idempotent and cheap, so calling it again to refresh the connection is + fine. + + Returns: + True once the connection is open (any HTTP response — even a + non-200 — means the socket is established); False if the + connection could not be opened (transport error). + """ + settings = self._client.settings + url = settings.dictation_base_url.rstrip("/") + api.ENDPOINT_WARM + try: + await self._client.http_client.get( + url, + timeout=min(settings.dictation_http_timeout, 10.0), + ) + except httpx.HTTPError: + return False + return True + + 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/dictation/v1/client.py b/assemblyai/dictation/v1/client.py new file mode 100644 index 00000000..ec025697 --- /dev/null +++ b/assemblyai/dictation/v1/client.py @@ -0,0 +1,377 @@ +from __future__ import annotations + +import concurrent.futures +import os +import queue +from typing import Any, Callable, Iterator, Optional + +import httpx + +from ... import client as _client +from ..._multipart import _Aborted, _as_bytes +from . import api +from ._base import AudioSource, _DictationTranscriberImpl, check_config +from .models import DictationConfig, DictationResponse + + +class DictationLiveSession: + """ + A live dictation that audio is pushed into. + + Returned by `DictationTranscriber.open_live()`. The request starts on one + of the transcriber's worker threads the moment the session opens; + `write()` hands it audio, `close()` ends the audio, and `result()` waits + for the transcript. Built for callback-driven sources — a microphone + library, a WebRTC track, a telephony media stream — where the audio + arrives in a callback rather than from an iterator you can hand to + `transcribe_live()`. + + Everything `transcribe_live()` says about the 120 s audio cap, the + server-side silence abort and errors surfacing mid-upload applies here + unchanged. + + Example: + ```python + import sounddevice as sd + + config = aai.DictationConfig(sample_rate=16000, channels=1) + + with aai.DictationTranscriber() as transcriber: + with transcriber.open_live(config) as session: + stream = sd.RawInputStream( + samplerate=16000, channels=1, dtype="int16", + callback=lambda data, *_: session.write(bytes(data)), + ) + with stream: + input("Dictating, press Enter to stop... ") + print(session.result().final_text) + ``` + """ + + def __init__( + self, + start: Callable[ + [Iterator[bytes]], concurrent.futures.Future[DictationResponse] + ], + ) -> None: + """ + Args: + start: begins the upload from the given producer and returns the + future that will hold its outcome. Supplied by the transcriber. + """ + self._queue: "queue.Queue[Optional[bytes]]" = queue.Queue() + self._closed = False + self._aborted = False + self._future = start(self._chunks()) + + def _chunks(self) -> Iterator[bytes]: + while True: + chunk = self._queue.get() + if chunk is None: + if self._aborted: + raise _Aborted() + return + yield chunk + + @property + def closed(self) -> bool: + """Whether the audio has ended, by `close()`, `result()` or `abort()`.""" + + return self._closed + + def write(self, chunk: bytes) -> None: + """ + Queues a piece of audio for upload. + + Safe to call from any thread, including an audio library's capture + callback; it never blocks. + + Args: + chunk: `bytes`, `bytearray` or `memoryview` of audio, in order. + + Raises: + TypeError: if `chunk` is not bytes-like. + RuntimeError: if the session is closed. + """ + if self._closed: + raise RuntimeError( + "the live session is closed; no more audio can be written" + ) + + self._queue.put(_as_bytes(chunk)) + + def close(self) -> None: + """ + Ends the audio. The server transcribes what was sent; `result()` + returns it. Idempotent, and never blocks. + """ + if not self._closed: + self._closed = True + self._queue.put(None) + + def result(self) -> DictationResponse: + """ + Waits for the transcript, ending the audio first if it is still open. + + Raises: + DictationError: if the request failed, including a rejection the + server sent while the upload was still in flight. + RuntimeError: if the session was aborted. + """ + if self._aborted: + raise RuntimeError("the live session was aborted; there is no result") + + self.close() + + return self._future.result() + + def abort(self) -> None: + """ + Drops the request without a transcript. + + The connection is closed and nothing the server may still return is + kept; `result()` raises afterwards. Idempotent, and a no-op once the + request has completed. Blocks until the worker thread has let go of + the connection. After `close()` the upload is already complete, so + aborting then waits out the in-flight response instead. + """ + if self._aborted or self._future.done(): + return + + self._aborted = True + if not self._closed: + self._closed = True + self._queue.put(None) + + try: + self._future.result() + except Exception: + pass + + def __enter__(self) -> "DictationLiveSession": + return self + + def __exit__(self, exc_type: Any, *_exc: Any) -> None: + """Ends the audio on a clean exit; aborts if the block raised.""" + + if exc_type is not None: + self.abort() + else: + self.close() + + +class DictationTranscriber: + """ + Transcribes dictated audio as it is spoken: audio in, transcript out, one + request. + + Targets the Dictation API (`dictation.assemblyai.com`), tuned for short + dictated clips. The audio is uploaded while it is still being produced — + from a callback with `open_live()`, or from an iterator or file object + with `transcribe_live()` — so the upload and every speech segment but the + last are done by the time the speaker stops. Audio you already hold whole + (bytes, or a local path) goes to `transcribe_live()` too; it travels the + same connection as a single chunk. Like `SyncTranscriber` there is no job + id and no polling, and no URL ingestion. Beyond the transcript it can run + a follow-up LLM pass: set `llm_instruction` on the config and read + `result.final_text`. + + Example: + ```python + import assemblyai as aai + + aai.settings.api_key = "your-key" + + with aai.DictationTranscriber() as transcriber: + with transcriber.open_live() as session: + start_capture(on_audio=session.write) + wait_until_done() + print(session.result().final_text) + ``` + """ + + def __init__( + self, + *, + client: Optional[_client.Client] = None, + config: Optional[DictationConfig] = None, + max_workers: Optional[int] = None, + api_key: Optional[str] = None, + ) -> None: + """ + Creates a `DictationTranscriber`. + + Args: + client: The HTTP client to use. Defaults to the shared default client. + config: Default dictation options. Per-call `config` overrides it. + max_workers: Thread pool size for `open_live()` sessions, each of + which runs its request on a worker thread. Defaults to the CPU + count minus one. + api_key: The API key to authenticate with. Builds a `Client` for this + transcriber. Given alongside `client`, it takes precedence: the + transcriber builds its own client from a copy of that client's + settings with the key replaced, and the given client is left + untouched. + + Raises: + TypeError: if `config` is not a `DictationConfig`. + """ + check_config(type(self).__name__, config) + + self._client = _client._resolve_client(client, api_key) + self._impl = _DictationTranscriberImpl( + client=self._client, + config=config or DictationConfig(), + owner=type(self).__name__, + ) + + if not max_workers: + cpu_count = os.cpu_count() + max_workers = max(1, cpu_count - 1) if cpu_count else 1 + + self._executor = concurrent.futures.ThreadPoolExecutor( + max_workers=max_workers, + ) + + @property + def config(self) -> DictationConfig: + """The default configuration of the `DictationTranscriber`.""" + return self._impl.config + + @config.setter + def config(self, config: DictationConfig) -> None: + check_config(type(self).__name__, config) + + self._impl.config = config + + def transcribe_live( + self, + data: AudioSource, + config: Optional[DictationConfig] = None, + ) -> DictationResponse: + """ + Transcribes audio uploaded as it is produced. + + Starts the request immediately and uploads chunks as they arrive, so + authorization, the upload and every speech segment but the last + resolve while the caller is still recording. What is left to wait for + once they stop is the final segment, and the LLM pass if one was asked + for. Audio that is already complete — bytes, or a local path — is + accepted as well and sent as a single chunk over the same connection. + + The caller must keep producing: an upload that goes silent for long + enough is aborted server-side. Stop by ending the iterator, not by + pausing it. The service caps a request at 120 s of audio. + + Args: + data: An iterable of audio chunks, a binary file object read as it + fills, raw audio bytes, or a local file path. Raw PCM also + requires `sample_rate` and `channels` on the config. + config: Options for this call. If `None`, the transcriber's default + configuration is used. + + Raises: + TypeError: if `config` is not a `DictationConfig`; if `data` is an + async iterable or an unsupported type; or if a chunk is not + bytes (a file opened in text mode, say). + ValueError: for a URL, or raw PCM without both `sample_rate` and + `channels`. + DictationError: if the request fails. Auth, rate-limit, size and + capacity failures can surface part-way through the upload. + Exception: anything the producer raises mid-upload propagates + unchanged. The connection is dropped and no transcript is + returned. + + Example: + ```python + def mic_chunks(): + while recording: + yield stream.read(4096) + + result = aai.DictationTranscriber().transcribe_live(mic_chunks()) + ``` + """ + check_config(type(self).__name__, config) + + return self._impl.transcribe_live(data=data, config=config) + + def open_live( + self, + config: Optional[DictationConfig] = None, + ) -> DictationLiveSession: + """ + Opens a live dictation that audio is pushed into. + + The push-style counterpart of `transcribe_live()`, for sources that + deliver audio through a callback rather than an iterator. The request + starts on a worker thread immediately; call `session.write(chunk)` from + the callback, then `session.result()` for the transcript once the + speaker stops. See `DictationLiveSession`. + + Args: + config: Options for this call. If `None`, the transcriber's default + configuration is used. Raw PCM requires `sample_rate` and + `channels`. + + Raises: + TypeError: if `config` is not a `DictationConfig`. + + Example: + ```python + with transcriber.open_live(config) as session: + start_capture(on_audio=session.write) + wait_until_done() + result = session.result() + ``` + """ + check_config(type(self).__name__, config) + + return DictationLiveSession( + lambda chunks: self._executor.submit( + self._impl.transcribe_live, data=chunks, config=config + ) + ) + + def warm(self) -> bool: + """ + Opens the connection to the Dictation API ahead of time. + + A request that opens its connection on demand pays the full DNS + TCP + + TLS handshake before the first audio byte can leave — one network + round trip that, for a distant client, is a noticeable share of a + short dictation. Calling `warm()` as soon as you know audio is coming + — when the user reaches for the record button, say — spends that setup + early: the next request reuses the already-open connection. + + The warmed connection is reused while it stays in the HTTP pool — + `settings.keepalive_expiry` seconds (httpx's 5s default unless + raised). Call `warm()` shortly before the request, or raise + `keepalive_expiry` so a single call covers a longer pause. `warm()` is + idempotent and cheap, so calling it again to refresh the connection is + fine. + + Returns: + True once the connection is open (any HTTP response — even a + non-200 — means the socket is established); False if the + connection could not be opened (transport error). + """ + settings = self._client.settings + url = settings.dictation_base_url.rstrip("/") + api.ENDPOINT_WARM + try: + self._client.http_client.get( + url, + timeout=min(settings.dictation_http_timeout, 10.0), + ) + except httpx.HTTPError: + return False + return True + + def close(self) -> None: + """Shuts down the worker-thread pool that `open_live()` sessions run on.""" + self._executor.shutdown(wait=False) + + def __enter__(self) -> "DictationTranscriber": + return self + + def __exit__(self, *_exc: Any) -> None: + self.close() diff --git a/assemblyai/dictation/v1/models.py b/assemblyai/dictation/v1/models.py new file mode 100644 index 00000000..6d5c0f30 --- /dev/null +++ b/assemblyai/dictation/v1/models.py @@ -0,0 +1,185 @@ +"""Models for the Dictation API. + +The request and response shapes of `dictation.assemblyai.com`, plus the error +the transcribers raise, live here alongside the client that uses them. +""" + +from typing import List, Optional + +try: + # pydantic v2 import + from pydantic import BaseModel, ConfigDict, Field, field_validator + + pydantic_v2 = True +except ImportError: + # pydantic v1 import (fallback for Python < 3.14) + from pydantic import BaseModel, Field, validator + + pydantic_v2 = False + +from ...types import AssemblyAIError + +__all__ = [ + "DictationConfig", + "DictationError", + "DictationResponse", + "DictationWord", +] + + +class DictationError(AssemblyAIError): + """ + Error raised when a Dictation API request fails. + + Carries a machine-readable `error_code` — the snake_cased problem-details + `title` from the server (e.g. `bad_audio`, `audio_too_large`, + `capacity_exceeded`, `inference_timeout`) — when present, and + `retry_after` (seconds) for 429/503 responses that include a + `Retry-After` header. + """ + + def __init__( + self, + message: str, + status_code: Optional[int] = None, + error_code: Optional[str] = None, + retry_after: Optional[int] = None, + ): + super().__init__(message, status_code) + self.error_code = error_code + self.retry_after = retry_after + + +_DICTATION_MAX_KEYTERMS_PROMPT_LEN = 2048 +_DICTATION_MAX_LLM_INSTRUCTION_LEN = 2048 + + +class DictationConfig(BaseModel): + """ + Options for a Dictation API request. + + `language_codes` and `keyterms_prompt` shape the transcript; + `llm_instruction` asks the server to run a follow-up LLM pass over it + (the rewrite lands in `DictationResponse.llm_response`). `sample_rate` + and `channels` are required only for raw PCM audio — container formats + carry them in their own headers. + + Unknown fields are rejected rather than silently dropped: the Dictation + API accepts this exact set, so a typo or a sync-only option surfaces as a + validation error instead of a config that quietly does nothing. + """ + + sample_rate: Optional[int] = None + """Source sample rate in Hz. Setting either this or `channels` marks the + audio as raw 16-bit PCM, and both are then required. Leave both unset for + WAV and the other container formats, which carry the rate in their own + headers.""" + + channels: Optional[int] = None + """Channel count (1 mono, 2 stereo). Setting either this or `sample_rate` + marks the audio as raw 16-bit PCM, and both are then required. Leave both + unset for WAV and the other container formats, which carry the channel + count in their own headers.""" + + language_codes: Optional[List[str]] = None + """ISO 639-1 codes for the language(s) of the audio — a single-element + list (e.g. `["es"]`) for monolingual audio, or several codes (e.g. + `["en", "es"]`) for multilingual audio. Defaults to None, which leaves + the language to the server's default.""" + + keyterms_prompt: Optional[List[str]] = None + "Keyterms biasing the decoder. Whitespace is stripped and empty terms dropped. Max 2048 characters total." + + llm_instruction: Optional[str] = Field( + default=None, max_length=_DICTATION_MAX_LLM_INSTRUCTION_LEN + ) + """Instruction for a follow-up LLM pass over the transcript, e.g. + "Format this as a SOAP note." The rewritten text comes back as + `llm_response`; the raw transcript stays in `text`. Max 2048 characters.""" + + if pydantic_v2: + model_config = ConfigDict(extra="forbid") + + @field_validator("keyterms_prompt") + @classmethod + def _normalize_keyterms_prompt(cls, v): + if not v: + return None + terms = [t.strip() for t in v if t and t.strip()] + total = sum(len(t) for t in terms) + if total > _DICTATION_MAX_KEYTERMS_PROMPT_LEN: + raise ValueError( + f"keyterms_prompt exceeds {_DICTATION_MAX_KEYTERMS_PROMPT_LEN} characters (got {total})" + ) + return terms or None + + else: + + class Config: + extra = "forbid" + + @validator("keyterms_prompt") + def _normalize_keyterms_prompt(cls, v): + if not v: + return None + terms = [t.strip() for t in v if t and t.strip()] + total = sum(len(t) for t in terms) + if total > _DICTATION_MAX_KEYTERMS_PROMPT_LEN: + raise ValueError( + f"keyterms_prompt exceeds {_DICTATION_MAX_KEYTERMS_PROMPT_LEN} characters (got {total})" + ) + return terms or None + + +class DictationWord(BaseModel): + """A single word in a dictation transcript.""" + + text: str + "The text of the word." + + confidence: float + "Word confidence in the range 0-1." + + +class DictationResponse(BaseModel): + """The result of a Dictation API request.""" + + text: str + "The raw transcript text, before any LLM pass." + + words: List[DictationWord] = Field(default_factory=list) + "Per-word text and confidence." + + confidence: float + "Overall transcript confidence in the range 0-1." + + llm_response: Optional[str] = None + """The transcript rewritten by the LLM pass `llm_instruction` asked for. + `None` when no instruction was sent or the pass failed.""" + + llm_error: Optional[str] = None + "Why the LLM pass failed, when it did. `None` otherwise." + + audio_duration_ms: int + "Total audio duration in milliseconds." + + session_id: str + "Server-generated UUID for this request. Record it to correlate with support." + + request_time_ms: Optional[float] = None + """End-to-end server-side request time in milliseconds: auth, multipart + parse, decode, inference, the LLM pass, and serialization.""" + + sync_time_ms: Optional[float] = None + "Time in milliseconds spent transcribing, excluding the LLM pass." + + @property + def final_text(self) -> str: + """ + The text to show the user: the LLM rewrite when there is one. + + Falls back to the raw transcript when no `llm_instruction` was sent + or the LLM pass failed, so reading it is safe either way. + """ + + return self.llm_response if self.llm_response is not None else self.text diff --git a/assemblyai/sync/v1/__init__.py b/assemblyai/sync/v1/__init__.py index de668acf..db2b2ed8 100644 --- a/assemblyai/sync/v1/__init__.py +++ b/assemblyai/sync/v1/__init__.py @@ -1,3 +1,4 @@ +from ..._multipart import AsyncAudioChunks, AudioChunks from ...types import ( SyncSpeechModel, SyncTranscriptError, @@ -6,7 +7,6 @@ SyncWord, ) from ._base import AudioInput -from ._multipart import AsyncAudioChunks, AudioChunks from .async_client import AsyncLiveSession, AsyncSyncTranscriber from .client import LiveSession, SyncTranscriber diff --git a/assemblyai/sync/v1/_base.py b/assemblyai/sync/v1/_base.py index 8507a1c9..a26a27ae 100644 --- a/assemblyai/sync/v1/_base.py +++ b/assemblyai/sync/v1/_base.py @@ -1,18 +1,31 @@ from __future__ import annotations import os -from typing import BinaryIO, Optional, Tuple, Union -from urllib.parse import urlparse +from typing import Mapping, Optional, Tuple +from ... import _audio, types from ... import client as _client -from ... import types +from ..._audio import _PCM_SUFFIXES, AudioInput +from ..._multipart import AudioChunks, _Aborted from . import api -from ._multipart import AudioChunks -AudioInput = Union[str, bytes, bytearray, "os.PathLike[str]", BinaryIO] - -# Extensions that signal raw S16LE PCM rather than a WAV container. -_PCM_SUFFIXES = (".pcm", ".raw") +__all__ = [ + "_PCM_SUFFIXES", + "AudioChunks", + "AudioInput", + "_Aborted", + "_SyncTranscriberImpl", + "_config_to_json", + "_resolve_audio", + "check_chunks", + "check_config", + "resolve_format", + "stream_filename", +] + +# The sync API decodes a WAV container or raw PCM; every other extension is +# posted as WAV and left to the server to sniff. +_CONTENT_TYPES: Mapping[str, str] = {} def check_config(owner: str, config: Optional[types.SyncTranscriptionConfig]) -> None: @@ -27,11 +40,7 @@ def check_config(owner: str, config: Optional[types.SyncTranscriptionConfig]) -> config: the configuration to check. """ - if config is not None and not isinstance(config, types.SyncTranscriptionConfig): - raise TypeError( - f"{owner} expects SyncTranscriptionConfig, got {type(config).__name__}. " - "Use aai.SyncTranscriptionConfig." - ) + _audio.check_config(owner, config, types.SyncTranscriptionConfig) def _resolve_audio( @@ -49,34 +58,15 @@ def _resolve_audio( Returns: `(audio_bytes, filename, content_type)`. """ - suffix = "" - filename: Optional[str] = None - if isinstance(data, (bytes, bytearray)): - audio = bytes(data) - elif isinstance(data, (str, os.PathLike)): - path = os.fspath(data) - if urlparse(path).scheme in ("http", "https"): - raise ValueError( - "SyncTranscriber does not accept URLs. Pass a local file path or " - "audio bytes, or use aai.Transcriber for URL/async transcription." - ) - with open(path, "rb") as f: - audio = f.read() - filename = os.path.basename(path) - suffix = os.path.splitext(path)[1].lower() - elif hasattr(data, "read"): - audio = data.read() - name = getattr(data, "name", None) - if name: - filename = os.path.basename(name) - suffix = os.path.splitext(name)[1].lower() - else: - raise TypeError(f"unsupported audio input type: {type(data).__name__}") - - resolved_filename, content_type = resolve_format(config, suffix, filename) - - return audio, resolved_filename, content_type + return _audio._resolve_audio( + data, + sample_rate=config.sample_rate, + channels=config.channels, + owner="SyncTranscriber", + config_name="SyncTranscriptionConfig", + content_types=_CONTENT_TYPES, + ) def resolve_format( @@ -100,35 +90,21 @@ def resolve_format( Returns: `(filename, content_type)`. """ - wants_pcm = config.sample_rate is not None or config.channels is not None - is_pcm = suffix in _PCM_SUFFIXES or wants_pcm - if is_pcm and (config.sample_rate is None or config.channels is None): - raise ValueError( - "raw PCM audio requires both sample_rate and channels in " - "SyncTranscriptionConfig" - ) - content_type = "audio/pcm" if is_pcm else "audio/wav" - if not filename: - filename = "audio.pcm" if is_pcm else "audio.wav" - - return filename, content_type + return _audio.resolve_format( + suffix=suffix, + filename=filename, + sample_rate=config.sample_rate, + channels=config.channels, + config_name="SyncTranscriptionConfig", + content_types=_CONTENT_TYPES, + ) def _config_to_json(config: types.SyncTranscriptionConfig) -> Optional[dict]: """Serializes the config to the JSON `config` part, dropping the routing model.""" - data = config.dict(exclude_none=True) - data.pop("model", None) - return data or None - -class _Aborted(Exception): - """ - Raised inside a live session's producer when the session is aborted. - - Propagates out of the transport, which drops the connection, and is - swallowed by `abort()`; it never reaches the caller. - """ + return _audio._config_to_json(config, exclude=("model",)) def check_chunks(data: object, *, allow_async: bool = False) -> None: diff --git a/assemblyai/sync/v1/api.py b/assemblyai/sync/v1/api.py index 8eb18e1e..88134070 100644 --- a/assemblyai/sync/v1/api.py +++ b/assemblyai/sync/v1/api.py @@ -1,9 +1,9 @@ -from typing import Iterator, Optional +from typing import Callable, Iterator, Optional import httpx from ... import types -from ._multipart import AudioChunks, StreamingMultipartEncoder, iter_chunks +from ..._multipart import AudioChunks, StreamingMultipartEncoder, iter_chunks # Canonical paths since the sync API gained a /v1 prefix (#18103); the # unprefixed routes remain served for SDK versions that predate it. @@ -19,15 +19,28 @@ MODEL_HEADER = "X-AAI-Model" -def _error_from_response(response: httpx.Response) -> types.SyncTranscriptError: +_ErrorFactory = Callable[..., types.AssemblyAIError] + + +def _error_from_response( + response: httpx.Response, + error_cls: _ErrorFactory = types.SyncTranscriptError, + what: str = "sync transcription", +) -> types.AssemblyAIError: """ - Builds a `SyncTranscriptError` from a non-200 response. + Builds an error from a non-200 response. The service returns an RFC 9457 problem-details envelope (`{"status", "title", "detail"}`); `error_code` is the snake_cased `title` (e.g. `"Audio Too Large"` -> `audio_too_large`). Older envelopes (`{"error_code", "message"}`, `{"detail"}`, and `{"error"}`) are still accepted; a bare `error` string carries no `error_code`. + + Args: + response: the non-200 response. + error_cls: the exception class to build, taking `message`, + `status_code`, `error_code`, and `retry_after`. + what: the operation to name in the fallback message. """ error_code: Optional[str] = None message: Optional[str] = None @@ -48,7 +61,7 @@ def _error_from_response(response: httpx.Response) -> types.SyncTranscriptError: message = response.text or None if not message: - message = f"sync transcription failed with status {response.status_code}" + message = f"{what} failed with status {response.status_code}" retry_after_header = response.headers.get("retry-after") retry_after = ( @@ -57,7 +70,7 @@ def _error_from_response(response: httpx.Response) -> types.SyncTranscriptError: else None ) - return types.SyncTranscriptError( + return error_cls( message, status_code=response.status_code, error_code=error_code, diff --git a/assemblyai/sync/v1/async_api.py b/assemblyai/sync/v1/async_api.py index 05f94978..b6530111 100644 --- a/assemblyai/sync/v1/async_api.py +++ b/assemblyai/sync/v1/async_api.py @@ -9,7 +9,7 @@ import httpx from ... import types -from ._multipart import AsyncAudioChunks, StreamingMultipartEncoder, aiter_chunks +from ..._multipart import AsyncAudioChunks, StreamingMultipartEncoder, aiter_chunks from .api import ( ENDPOINT_TRANSCRIBE_LIVE, MODEL_HEADER, diff --git a/assemblyai/sync/v1/async_client.py b/assemblyai/sync/v1/async_client.py index a78c5aa2..40ade287 100644 --- a/assemblyai/sync/v1/async_client.py +++ b/assemblyai/sync/v1/async_client.py @@ -11,6 +11,7 @@ from ... import async_client as _async_client from ... import types +from ..._multipart import AsyncAudioChunks, _as_bytes from . import api, async_api from ._base import ( AudioInput, @@ -21,7 +22,6 @@ check_config, stream_filename, ) -from ._multipart import AsyncAudioChunks, _as_bytes _T = TypeVar("_T") diff --git a/assemblyai/sync/v1/client.py b/assemblyai/sync/v1/client.py index 1eeeb2b6..5a688506 100644 --- a/assemblyai/sync/v1/client.py +++ b/assemblyai/sync/v1/client.py @@ -9,6 +9,7 @@ from ... import client as _client from ... import types +from ..._multipart import _as_bytes from . import api from ._base import ( AudioChunks, @@ -17,7 +18,6 @@ _SyncTranscriberImpl, check_config, ) -from ._multipart import _as_bytes class LiveSession: diff --git a/assemblyai/types.py b/assemblyai/types.py index 36ce1969..11ae861f 100644 --- a/assemblyai/types.py +++ b/assemblyai/types.py @@ -122,6 +122,16 @@ class Settings(BaseSettings): sync_live_http_timeout: float = 180.0 "The HTTP timeout for streamed synchronous transcription requests (`transcribe_live`). Like every httpx timeout it bounds each socket operation, not the request as a whole: connecting, each write of a chunk, and each read while waiting for the response after the last byte. Time spent waiting on the caller's producer is not counted, so it need not cover the recording. The sync API caps audio at 120 seconds and independently aborts an upload that goes silent." + dictation_base_url: str = "https://dictation.assemblyai.com" + "The base URL for the Dictation API (used by `DictationTranscriber`)" + + dictation_http_timeout: float = 300.0 + """The HTTP timeout for Dictation API requests. Like every httpx timeout it + bounds each socket operation — connecting, each write, each read while + waiting for the response — not the request end to end, so time spent + producing audio is not counted. Sized to outlast the final segment's + inference plus the LLM pass over the transcript.""" + polling_interval: float = Field(default=3.0, gt=0.0) "The default polling interval for long-running requests (e.g. polling the `Transcript`'s status)" diff --git a/tests/unit/test_audio_helpers.py b/tests/unit/test_audio_helpers.py new file mode 100644 index 00000000..ac227f6f --- /dev/null +++ b/tests/unit/test_audio_helpers.py @@ -0,0 +1,174 @@ +"""Tests for the audio helpers the sync and Dictation APIs share. + +``assemblyai/_audio.py`` holds the input reading, Content-Type selection, and +config serialization both single-request products need. Each product package +wraps it with its own config type and error wording, so the wrappers must keep +the exact names, signatures, and wire behavior their callers rely on. +""" + +import assemblyai as aai +from assemblyai import _audio +from assemblyai.dictation.v1 import _base as dictation_base +from assemblyai.sync.v1 import _base as sync_base + + +def test_sync_base_still_exposes_the_helper_names(): + """Every name ``sync/v1/_base.py`` published is still importable from it.""" + for name in ( + "AudioInput", + "_PCM_SUFFIXES", + "_resolve_audio", + "_config_to_json", + "check_config", + "_SyncTranscriberImpl", + ): + assert hasattr(sync_base, name), f"assemblyai.sync.v1._base.{name} is gone" + + +def test_sync_audio_input_is_the_shared_alias(): + # Given the shared alias, Then sync re-exports the same object + assert sync_base.AudioInput is _audio.AudioInput + + +def test_sync_resolve_audio_defaults_unknown_extensions_to_wav(tmp_path): + # Given an MP3 file and a default sync config + audio_file = tmp_path / "call.mp3" + audio_file.write_bytes(b"ID3fake-mp3-bytes") + + # When resolving it for the sync API + audio, filename, content_type = sync_base._resolve_audio( + str(audio_file), aai.SyncTranscriptionConfig() + ) + + # Then the sync API still posts every container as WAV + assert audio == b"ID3fake-mp3-bytes" + assert filename == "call.mp3" + assert content_type == "audio/wav" + + +def test_dictation_resolve_source_maps_the_extension(tmp_path): + # Given the same MP3 file and a default dictation config + audio_file = tmp_path / "call.mp3" + audio_file.write_bytes(b"ID3fake-mp3-bytes") + + # When resolving it for the Dictation API + chunks, filename, content_type = dictation_base.resolve_source( + str(audio_file), aai.DictationConfig(), "DictationTranscriber" + ) + + # Then the file is read whole as one chunk and the extension picks the + # container Content-Type + assert list(chunks) == [b"ID3fake-mp3-bytes"] + assert filename == "call.mp3" + assert content_type == "audio/mpeg" + + +def test_resolve_audio_selects_pcm_for_both_products(): + # Given raw bytes with sample_rate + channels + sync_config = aai.SyncTranscriptionConfig(sample_rate=16000, channels=1) + dictation_config = aai.DictationConfig(sample_rate=16000, channels=1) + + # When resolving for either product + _, sync_name, sync_type = sync_base._resolve_audio(b"\x00\x01", sync_config) + _, dictation_name, dictation_type = dictation_base.resolve_source( + b"\x00\x01", dictation_config, "DictationTranscriber" + ) + + # Then both route the audio as raw PCM under the same default filename + assert (sync_name, sync_type) == ("audio.pcm", "audio/pcm") + assert (dictation_name, dictation_type) == ("audio.pcm", "audio/pcm") + + +def test_sync_config_to_json_still_drops_the_routing_model(): + # Given a sync config with a prompt + config = aai.SyncTranscriptionConfig(prompt="Transcribe verbatim.") + + # When serializing the config part + data = sync_base._config_to_json(config) + + # Then the routing model stays out of the body + assert data == {"prompt": "Transcribe verbatim."} + + +def test_dictation_config_to_json_keeps_every_set_field(): + # Given a dictation config with two fields set + config = aai.DictationConfig(language_codes=["en"], llm_instruction="Summarize.") + + # When serializing the config part + data = dictation_base._config_to_json(config) + + # Then both survive and the unset fields are dropped + assert data == {"language_codes": ["en"], "llm_instruction": "Summarize."} + + +def test_config_to_json_returns_none_when_nothing_is_set(): + # Given empty configs + # When serializing, Then there is no config part to send + assert sync_base._config_to_json(aai.SyncTranscriptionConfig()) is None + assert dictation_base._config_to_json(aai.DictationConfig()) is None + + +def test_check_config_names_the_expected_type(): + # Given a config of the wrong product + # When checking it, Then the message names the owner and expected type + try: + sync_base.check_config("SyncTranscriber", aai.DictationConfig()) + except TypeError as error: + assert "SyncTranscriber expects SyncTranscriptionConfig" in str(error) + assert "Use aai.SyncTranscriptionConfig." in str(error) + else: + raise AssertionError("check_config accepted the wrong config type") + + try: + dictation_base.check_config( + "DictationTranscriber", aai.SyncTranscriptionConfig() + ) + except TypeError as error: + assert "DictationTranscriber expects DictationConfig" in str(error) + else: + raise AssertionError("check_config accepted the wrong config type") + + +def test_check_config_accepts_none(): + # Given no config at all, Then neither product objects + assert sync_base.check_config("SyncTranscriber", None) is None + assert dictation_base.check_config("DictationTranscriber", None) is None + + +def test_resolve_source_names_the_calling_client_in_url_errors(): + # Given each dictation client's own class name + for owner in ("DictationTranscriber", "AsyncDictationTranscriber"): + # When a URL is passed, Then the error names that client, not a fixed one + try: + dictation_base.resolve_source( + "https://example.com/audio.wav", aai.DictationConfig(), owner + ) + except ValueError as error: + assert str(error).startswith(f"{owner} does not accept URLs.") + else: + raise AssertionError("resolve_source accepted a URL") + + +def test_resolve_format_serves_the_streamed_path_without_audio_bytes(): + # Given only a suffix and the PCM fields, no audio + filename, content_type = _audio.resolve_format( + suffix="", + filename=None, + sample_rate=16000, + channels=1, + config_name="DictationConfig", + content_types={}, + ) + + # Then the part is named and typed as PCM without reading anything + assert (filename, content_type) == ("audio.pcm", "audio/pcm") + + # And an unknown suffix falls back to WAV under the given name + assert _audio.resolve_format( + suffix=".xyz", + filename="clip.xyz", + sample_rate=None, + channels=None, + config_name="DictationConfig", + content_types={}, + ) == ("clip.xyz", "audio/wav") diff --git a/tests/unit/test_dictation.py b/tests/unit/test_dictation.py new file mode 100644 index 00000000..ebff0e22 --- /dev/null +++ b/tests/unit/test_dictation.py @@ -0,0 +1,972 @@ +"""Tests for the threaded dictation client. + +Every entry point posts to the live route. Body-shape assertions go through +`httpx.MockTransport`, which consumes the request stream the way a real +transport does. `pytest_httpx` records the request without reading it on older +httpx, so a body read after the call sees a spent generator or a closed file. +""" + +import email +import io +import threading +from typing import List, Optional + +import httpx +import pytest +from pytest_httpx import HTTPXMock + +import assemblyai as aai +from assemblyai._multipart import _STREAM_READ_SIZE +from assemblyai.dictation.v1 import api + +aai.settings.api_key = "test" + +LIVE_URL = f"{aai.settings.dictation_base_url}/v1/transcribe/live" +BUFFERED_URL = f"{aai.settings.dictation_base_url}/v1/transcribe" +WARM_URL = f"{aai.settings.dictation_base_url}/v1/warm" + +_OK_RESPONSE = { + "text": "patient reports mild headache", + "words": [ + {"text": "patient", "confidence": 0.9}, + {"text": "reports", "confidence": 0.95}, + ], + "confidence": 0.92, + "llm_response": None, + "llm_error": None, + "audio_duration_ms": 400, + "session_id": "eb92c4ff-4bbb-429f-9b99-7279d7fe738f", + "request_time_ms": 243.7, + "sync_time_ms": 180.2, +} + + +def _mock_ok(httpx_mock: HTTPXMock, **overrides) -> None: + httpx_mock.add_response( + url=LIVE_URL, + method="POST", + status_code=httpx.codes.OK, + json={**_OK_RESPONSE, **overrides}, + ) + + +def _chunks(*pieces: bytes): + for piece in pieces: + yield piece + + +def _parts(request: httpx.Request): + """Parses a recorded request body into its multipart parts, in order.""" + + raw = b"Content-Type: " + request.headers["content-type"].encode() + b"\r\n\r\n" + message = email.message_from_bytes(raw + request.content) + + return [ + (part.get_param("name", header="content-disposition"), part) + for part in message.get_payload() + ] + + +def _send(chunks, config: Optional[dict] = None, **kwargs) -> List[httpx.Request]: + """Runs `api.transcribe_live` against a transport that records the body.""" + + seen: List[httpx.Request] = [] + + def handler(request: httpx.Request) -> httpx.Response: + seen.append(request) + return httpx.Response(httpx.codes.OK, json=_OK_RESPONSE) + + with httpx.Client(transport=httpx.MockTransport(handler)) as client: + api.transcribe_live( + client, + base_url=aai.settings.dictation_base_url, + chunks=chunks, + filename=kwargs.get("filename", "audio.wav"), + audio_content_type=kwargs.get("audio_content_type", "audio/wav"), + config=config, + timeout=30.0, + ) + + return seen + + +# --- public surface --------------------------------------------------------- + + +def test_public_names_are_the_dictation_models_classes(): + # Given the module the dictation models live in + from assemblyai.dictation.v1 import models + + # Then the top-level and versioned names are that module's classes + assert aai.DictationConfig is models.DictationConfig + assert aai.DictationError is models.DictationError + assert aai.DictationResponse is models.DictationResponse + assert aai.DictationWord is models.DictationWord + + +def test_there_is_no_buffered_transcribe(): + # Given the two clients + # Then neither exposes the buffered entry points: the live route is the + # only one the Dictation API markets, and the only one this client opens + for cls in (aai.DictationTranscriber, aai.AsyncDictationTranscriber): + assert not hasattr(cls, "transcribe") + assert not hasattr(cls, "transcribe_async") + assert hasattr(cls, "transcribe_live") + assert hasattr(cls, "open_live") + assert not hasattr(api, "ENDPOINT_TRANSCRIBE") + assert not hasattr(api, "transcribe") + + +def test_sessions_are_exported_from_the_dictation_packages(): + from assemblyai import dictation + from assemblyai.dictation import v1 + + for module in (dictation, v1): + assert module.DictationLiveSession is v1.client.DictationLiveSession + assert ( + module.AsyncDictationLiveSession + is v1.async_client.AsyncDictationLiveSession + ) + + +# --- body shape ------------------------------------------------------------- + + +def test_body_sends_config_before_audio(): + # Given a config alongside the audio + # When the body is built + request = _send(_chunks(b"RIFF", b"fake"), config={"language_codes": ["en"]})[0] + + # Then config precedes audio: the server decodes as bytes land and rejects + # audio that arrives first + assert [name for name, _ in _parts(request)] == ["config", "audio"] + + +def test_body_carries_config_json_and_joined_audio(): + # Given audio produced in several pieces + # When the body is built + request = _send(_chunks(b"RIFF", b"fake", b"wav"), config={"llm_instruction": "x"})[ + 0 + ] + + # Then the pieces arrive as one audio part and the config part is JSON + parts = dict(_parts(request)) + assert parts["audio"].get_payload(decode=True) == b"RIFFfakewav" + assert parts["config"].get_payload(decode=True) == b'{"llm_instruction": "x"}' + assert parts["config"].get_content_type() == "application/json" + + +def test_body_sends_an_empty_config_part_when_there_is_no_config(): + # Given no config + # When the body is built + request = _send(_chunks(b"RIFF"))[0] + + # Then a config part still precedes the audio, as an empty object: the live + # endpoint rejects audio that no config part came before + parts = _parts(request) + assert [name for name, _ in parts] == ["config", "audio"] + assert dict(parts)["config"].get_payload(decode=True) == b"{}" + + +def test_body_skips_empty_chunks(): + # Given a producer that yields an empty piece between real ones + # When the body is built + request = _send(_chunks(b"RIFF", b"", b"wav"))[0] + + # Then it is dropped rather than ending the chunked body early + assert dict(_parts(request))["audio"].get_payload(decode=True) == b"RIFFwav" + + +def test_body_is_uploaded_chunked(): + # Given audio that does not exist yet when the request starts + # When the body is built + request = _send(_chunks(b"RIFF", b"more"))[0] + + # Then it is framed without a length, which is what lets the request start + assert request.headers.get("transfer-encoding") == "chunked" + assert "content-length" not in request.headers + + +class _SpyFile(io.BytesIO): + """A file object that records the size of every read asked of it.""" + + def __init__(self, payload: bytes) -> None: + super().__init__(payload) + self.reads: List[int] = [] + self.threads: List[threading.Thread] = [] + + def read(self, size: int = -1) -> bytes: + self.reads.append(size) + self.threads.append(threading.current_thread()) + return super().read(size) + + +def test_body_reads_a_file_object_in_bounded_pieces(): + # Given a file object holding more than two read sizes of audio + payload = b"RIFF" + bytes(2 * _STREAM_READ_SIZE + 100) + source = _SpyFile(payload) + + # When it is streamed + request = _send(source, filename="note.wav")[0] + + # Then it is read in fixed pieces until exhausted, never slurped whole + assert source.reads == [_STREAM_READ_SIZE] * 4 + part = dict(_parts(request))["audio"] + assert part.get_payload(decode=True) == payload + assert part.get_filename() == "note.wav" + + +def test_body_escapes_the_filename(): + # Given a file whose name would otherwise break the part header + # When the body is built + request = _send(_chunks(b"RIFF"), filename='we"ird\r\nname.wav')[0] + + # Then the name is escaped the way httpx escapes its own multipart parts + # and the audio part still parses + parts = dict(_parts(request)) + assert parts["audio"].get_filename() == "we%22ird%0D%0Aname.wav" + assert parts["audio"].get_payload(decode=True) == b"RIFF" + + +def test_body_accepts_bytearray_and_memoryview_chunks(): + # Given a producer handing out buffer objects rather than bytes + # When the body is built + request = _send(iter([bytearray(b"RIFF"), memoryview(b"fake")]))[0] + + # Then they are sent as bytes + assert dict(_parts(request))["audio"].get_payload(decode=True) == b"RIFFfake" + + +def test_body_rejects_text_chunks(): + # Given a producer yielding str + # When the body is built, then the mistake is named rather than failing + # inside the transport + with pytest.raises(TypeError, match="binary mode"): + _send(iter(["RIFF", "fake"])) + + +def test_body_rejects_a_text_mode_file(): + # Given a file object opened without "b" + # When the body is built, then the mistake is named + with pytest.raises(TypeError, match="binary mode"): + _send(io.StringIO("RIFFfake")) + + +def test_body_hits_the_live_endpoint_without_a_model_header(): + # Given a streamed upload + # When it is sent + request = _send(_chunks(b"RIFF"))[0] + + # Then it lands on the live route, and the Dictation API picks its own + # model, so no routing header is sent + assert str(request.url) == LIVE_URL + assert "X-AAI-Model" not in request.headers + + +# --- transcriber behaviour -------------------------------------------------- + + +def test_transcribe_live_parses_response(httpx_mock: HTTPXMock): + # Given a mocked live endpoint + _mock_ok(httpx_mock) + + # When streaming audio chunks + result = aai.DictationTranscriber().transcribe_live(_chunks(b"RIFF", b"fake")) + + # Then the response is parsed into a DictationResponse + assert isinstance(result, aai.DictationResponse) + assert result.text == "patient reports mild headache" + assert result.session_id == _OK_RESPONSE["session_id"] + assert result.words[1].text == "reports" + assert result.words[1].confidence == 0.95 + assert result.confidence == 0.92 + assert result.audio_duration_ms == 400 + assert result.request_time_ms == 243.7 + assert result.sync_time_ms == 180.2 + + +def test_transcribe_live_sends_raw_api_key_to_the_dictation_host(httpx_mock: HTTPXMock): + # Given a mocked live endpoint + _mock_ok(httpx_mock) + + # When streaming + aai.DictationTranscriber().transcribe_live(_chunks(b"RIFF")) + + # Then the request goes to the dictation host, not the sync host, with the + # API key sent raw (no Bearer prefix) and no routing header + request = httpx_mock.get_requests()[0] + assert str(request.url) == LIVE_URL + assert request.headers["authorization"] == "test" + assert "X-AAI-Model" not in request.headers + + +def test_complete_audio_goes_over_the_live_connection(httpx_mock: HTTPXMock): + """Bytes are a stream whose pieces are all ready at once. + + There is one request shape in this client: a complete clip takes the live + endpoint and the chunked framing rather than a second, buffered path. + """ + _mock_ok(httpx_mock) + + result = aai.DictationTranscriber().transcribe_live(b"RIFFfake-wav-bytes") + + assert result.text == "patient reports mild headache" + request = httpx_mock.get_requests()[0] + assert str(request.url) == LIVE_URL + assert "content-length" not in request.headers + body = request.read() + assert body.index(b'name="config"') < body.index(b'name="audio"') + assert b"RIFFfake-wav-bytes" in body + assert b'filename="audio.wav"' in body + assert b"Content-Type: audio/wav" in body + + +def test_transcribe_live_uses_the_dictation_timeout(httpx_mock: HTTPXMock): + # Given a mocked live endpoint + _mock_ok(httpx_mock) + + # When streaming audio + aai.DictationTranscriber().transcribe_live(_chunks(b"RIFF")) + + # Then the request gets the dictation budget, which must outlast the final + # segment and the LLM pass + timeout = httpx_mock.get_requests()[0].extensions["timeout"] + assert timeout["read"] == aai.settings.dictation_http_timeout + + +def test_transcribe_live_uses_default_config_and_per_call_override( + httpx_mock: HTTPXMock, +): + # Given a transcriber with a default config + _mock_ok(httpx_mock) + _mock_ok(httpx_mock) + default = aai.DictationConfig(llm_instruction="default instruction") + transcriber = aai.DictationTranscriber(config=default) + + # When streaming without a per-call config, and with an override + transcriber.transcribe_live(_chunks(b"RIFF")) + override = aai.DictationConfig(llm_instruction="override instruction") + transcriber.transcribe_live(_chunks(b"RIFF"), config=override) + + # Then the default applies to the first call and the override to the second + first, second = (request.read() for request in httpx_mock.get_requests()) + assert b"default instruction" in first + assert b"override instruction" in second + + +def test_config_json_excludes_unset_fields(httpx_mock: HTTPXMock): + # Given a config that sets only llm_instruction + _mock_ok(httpx_mock) + config = aai.DictationConfig(llm_instruction="Format this as a SOAP note.") + + # When streaming + aai.DictationTranscriber().transcribe_live(_chunks(b"RIFF"), config=config) + + # Then the config JSON carries that field and omits the unset ones + body = httpx_mock.get_requests()[0].read() + assert b"Format this as a SOAP note." in body + assert b'"language_codes"' not in body + assert b'"sample_rate"' not in body + assert b'"keyterms_prompt"' not in body + + +def test_transcribe_live_sends_keyterms_prompt_and_language_codes( + httpx_mock: HTTPXMock, +): + # Given a mocked live endpoint + _mock_ok(httpx_mock) + + # When streaming with keyterms and languages + config = aai.DictationConfig( + keyterms_prompt=["AssemblyAI", " Universal ", ""], + language_codes=["en", "es"], + ) + aai.DictationTranscriber().transcribe_live(_chunks(b"RIFF"), config=config) + + # Then the config part carries both, with keyterms normalized + body = httpx_mock.get_requests()[0].read() + assert b'"AssemblyAI"' in body + assert b'"Universal"' in body # whitespace stripped, empty term dropped + assert b'"language_codes"' in body + assert b'"es"' in body + + +def test_transcribe_live_marks_pcm_from_config(httpx_mock: HTTPXMock): + # Given a config carrying the fields only raw PCM needs + _mock_ok(httpx_mock) + config = aai.DictationConfig(sample_rate=16000, channels=1) + + # When streaming + aai.DictationTranscriber().transcribe_live(_chunks(b"\x00\x01"), config=config) + + # Then the audio part selects the PCM decoder and the config carries both + body = httpx_mock.get_requests()[0].read() + assert b"Content-Type: audio/pcm" in body + assert b'filename="audio.pcm"' in body + assert b'"sample_rate"' in body + assert b'"channels"' in body + + +def test_transcribe_live_requires_both_pcm_fields(): + # Given a config with only half of what raw PCM needs + config = aai.DictationConfig(sample_rate=16000) + + # When streaming, then it is rejected before any request is made + with pytest.raises(ValueError, match="sample_rate and channels"): + aai.DictationTranscriber().transcribe_live(_chunks(b"\x00"), config=config) + + +def test_transcribe_live_pcm_path_requires_rate_and_channels(tmp_path): + # Given a .pcm file and a config with neither rate nor channels + audio_file = tmp_path / "note.pcm" + audio_file.write_bytes(b"\x00\x01" * 100) + + # When streaming it, then the extension alone routes it as PCM and the + # missing fields are caught locally + with pytest.raises(ValueError, match="sample_rate and channels"): + aai.DictationTranscriber().transcribe_live(str(audio_file)) + + +def test_transcribe_live_reads_a_path_and_maps_its_extension( + httpx_mock: HTTPXMock, tmp_path +): + # Given a local MP3 file + _mock_ok(httpx_mock) + audio_file = tmp_path / "note.mp3" + audio_file.write_bytes(b"ID3fake-mp3-bytes") + + # When streaming the path + aai.DictationTranscriber().transcribe_live(str(audio_file)) + + # Then the file is read and sent whole, with its true Content-Type and + # name, so a format the live route does not decode is rejected by name + body = httpx_mock.get_requests()[0].read() + assert b"ID3fake-mp3-bytes" in body + assert b"Content-Type: audio/mpeg" in body + assert b'filename="note.mp3"' in body + + +def test_transcribe_live_accepts_a_pathlike(httpx_mock: HTTPXMock, tmp_path): + # Given a local WAV file as a Path + _mock_ok(httpx_mock) + audio_file = tmp_path / "note.wav" + audio_file.write_bytes(b"RIFFfake-wav-bytes") + + # When streaming the Path object + result = aai.DictationTranscriber().transcribe_live(audio_file) + + # Then it succeeds and ships a WAV audio part under the file's name + assert result.text == "patient reports mild headache" + body = httpx_mock.get_requests()[0].read() + assert b"Content-Type: audio/wav" in body + assert b'filename="note.wav"' in body + + +def test_transcribe_live_names_a_file_object(httpx_mock: HTTPXMock): + # Given a file object carrying a name + _mock_ok(httpx_mock) + stream = io.BytesIO(b"RIFF") + stream.name = "/tmp/note.flac" + + # When it is streamed + aai.DictationTranscriber().transcribe_live(stream) + + # Then the audio part takes its name and type from the file + body = httpx_mock.get_requests()[0].read() + assert b'filename="note.flac"' in body + assert b"Content-Type: audio/flac" in body + + +def test_transcribe_live_rejects_url(): + # Given an http URL as input + transcriber = aai.DictationTranscriber() + + # When streaming, then it is rejected with a pointer to Transcriber, and + # the message names this client + with pytest.raises(ValueError, match="DictationTranscriber does not accept URLs"): + transcriber.transcribe_live("https://example.com/audio.wav") + + +def test_transcribe_live_rejects_an_async_iterable(): + # Given an async producer handed to the threaded transcriber + async def chunks(): + yield b"RIFF" + + # When streaming it, then it is turned away before any request is made + # and pointed at the transcriber that can drive it + with pytest.raises(TypeError, match="AsyncDictationTranscriber"): + aai.DictationTranscriber().transcribe_live(chunks()) + + +def test_transcribe_live_rejects_an_unsupported_source(): + # Given something that is neither audio nor a producer of it + with pytest.raises(TypeError, match="unsupported audio source type: int"): + aai.DictationTranscriber().transcribe_live(42) + + +def test_transcribe_live_rejects_other_products_configs(): + # Given the sync API's and the job API's config types + transcriber = aai.DictationTranscriber() + + # When streaming with either, then the mismatch is named + with pytest.raises(TypeError, match="expects DictationConfig"): + transcriber.transcribe_live( + _chunks(b"RIFF"), config=aai.SyncTranscriptionConfig() + ) + with pytest.raises(TypeError, match="expects DictationConfig"): + transcriber.transcribe_live(_chunks(b"RIFF"), config=aai.TranscriptionConfig()) + + +def test_final_text_prefers_llm_response(httpx_mock: HTTPXMock): + # Given a response carrying an LLM rewrite + _mock_ok(httpx_mock, llm_response="S: Mild headache.") + + # When streaming with an llm_instruction + config = aai.DictationConfig(llm_instruction="Format this as a SOAP note.") + result = aai.DictationTranscriber().transcribe_live(_chunks(b"RIFF"), config=config) + + # Then final_text is the rewrite and the raw transcript is still there + assert result.final_text == "S: Mild headache." + assert result.text == "patient reports mild headache" + + +def test_final_text_falls_back_to_text(httpx_mock: HTTPXMock): + # Given a response whose LLM pass failed + _mock_ok(httpx_mock, llm_error="llm timed out") + + # When streaming + result = aai.DictationTranscriber().transcribe_live(_chunks(b"RIFF")) + + # Then final_text falls back to the raw transcript + assert result.final_text == "patient reports mild headache" + assert result.llm_error == "llm timed out" + + +def test_response_ignores_unknown_fields(httpx_mock: HTTPXMock): + # Given a server that added a field this SDK version predates + _mock_ok(httpx_mock, auth_time_ms=1.4) + + # When streaming, then the extra key is ignored instead of failing + result = aai.DictationTranscriber().transcribe_live(_chunks(b"RIFF")) + assert result.text == "patient reports mild headache" + + +def test_config_rejects_unknown_fields(): + # Given an option the Dictation API does not accept + # When building the config, then it is rejected rather than dropped + with pytest.raises(ValueError): + aai.DictationConfig(prompt="Transcribe verbatim.") + + +def test_keyterms_prompt_too_long_raises(): + # Given a keyterms_prompt exceeding the 2048-char cap + with pytest.raises(ValueError, match="keyterms_prompt exceeds"): + aai.DictationConfig(keyterms_prompt=["x" * 3000]) + + +def test_llm_instruction_too_long_raises(): + # Given an llm_instruction exceeding the 2048-char cap + with pytest.raises(ValueError): + aai.DictationConfig(llm_instruction="x" * 3000) + + +def test_error_envelope_maps_to_dictation_error(httpx_mock: HTTPXMock): + # Given a server rejecting the audio with an {error, error_code} body + httpx_mock.add_response( + url=LIVE_URL, + method="POST", + status_code=415, + json={"error": "unsupported media type", "error_code": "bad_audio"}, + ) + + # When streaming, then a DictationError carries code + status + with pytest.raises(aai.DictationError) as exc_info: + aai.DictationTranscriber().transcribe_live(_chunks(b"RIFF")) + + error = exc_info.value + assert error.status_code == 415 + assert error.error_code == "bad_audio" + assert "unsupported media type" in str(error) + # And it is not the sync API's error class + assert not isinstance(error, aai.SyncTranscriptError) + + +def test_problem_details_envelope_maps_to_dictation_error(httpx_mock: HTTPXMock): + # Given the server rejects oversized audio mid-upload with a problem-details body + httpx_mock.add_response( + url=LIVE_URL, + method="POST", + status_code=413, + json={"status": 413, "title": "Audio Too Large", "detail": "too long"}, + ) + + # When streaming, then the snake_cased title becomes the error_code + with pytest.raises(aai.DictationError) as exc_info: + aai.DictationTranscriber().transcribe_live(_chunks(b"RIFF")) + + error = exc_info.value + assert error.status_code == 413 + assert error.error_code == "audio_too_large" + assert "too long" in str(error) + + +def test_detail_only_envelope_on_invalid_key(httpx_mock: HTTPXMock): + # Given an invalid key, which the Dictation API answers 404 + httpx_mock.add_response( + url=LIVE_URL, + method="POST", + status_code=404, + json={"detail": "Invalid API key"}, + ) + + # When streaming, then the detail becomes the message and error_code + # stays absent + with pytest.raises(aai.DictationError) as exc_info: + aai.DictationTranscriber().transcribe_live(_chunks(b"RIFF")) + + error = exc_info.value + assert error.status_code == 404 + assert error.error_code is None + assert "Invalid API key" in str(error) + + +def test_rate_limit_surfaces_retry_after(httpx_mock: HTTPXMock): + # Given a rate-limit response with a Retry-After header + httpx_mock.add_response( + url=LIVE_URL, + method="POST", + status_code=429, + json={ + "status": 429, + "title": "Too Many Requests", + "detail": "Too many requests", + }, + headers={"Retry-After": "5"}, + ) + + # When streaming, then retry_after and the snake_cased title are parsed + with pytest.raises(aai.DictationError) as exc_info: + aai.DictationTranscriber().transcribe_live(_chunks(b"RIFF")) + + error = exc_info.value + assert error.status_code == 429 + assert error.error_code == "too_many_requests" + assert error.retry_after == 5 + + +def test_warm_opens_connection(httpx_mock: HTTPXMock): + # Given a mocked warm endpoint + httpx_mock.add_response(url=WARM_URL, method="GET", status_code=httpx.codes.OK) + + # When warming the transcriber + warmed = aai.DictationTranscriber().warm() + + # Then it returns True and probes the dictation warm route + assert warmed is True + request = httpx_mock.get_requests()[0] + assert str(request.url) == WARM_URL + assert request.method == "GET" + assert "X-AAI-Model" not in request.headers + + +def test_warm_returns_true_on_non_200(httpx_mock: HTTPXMock): + # Given a warm route that the load balancer answers with a 404 + httpx_mock.add_response(url=WARM_URL, method="GET", status_code=404) + + # When warming, then the socket is still established, so warm() is True + assert aai.DictationTranscriber().warm() is True + + +def test_warm_returns_false_on_transport_error(httpx_mock: HTTPXMock): + # Given the dictation host is unreachable + httpx_mock.add_exception(httpx.ConnectError("connection refused")) + + # When warming, then the failure is swallowed and reported as False + assert aai.DictationTranscriber().warm() is False + + +def test_api_key_constructor_builds_own_client(httpx_mock: HTTPXMock): + # Given a transcriber constructed with an explicit key + _mock_ok(httpx_mock) + + # When streaming + aai.DictationTranscriber(api_key="per-call-key").transcribe_live(_chunks(b"RIFF")) + + # Then that key authenticates the request + assert httpx_mock.get_requests()[0].headers["authorization"] == "per-call-key" + + +def test_rejects_sync_transcription_config(): + # Given the sync API's config, a different non-interchangeable type + # When constructing, then the mistake is named + with pytest.raises(TypeError, match="expects DictationConfig"): + aai.DictationTranscriber(config=aai.SyncTranscriptionConfig()) + + +def test_config_setter_validates_and_round_trips(): + # Given a transcriber and a new config + transcriber = aai.DictationTranscriber() + config = aai.DictationConfig(language_codes=["de"]) + + # When assigning it + transcriber.config = config + + # Then it is the transcriber's default + assert transcriber.config is config + + # And a wrong type is rejected + with pytest.raises(TypeError, match="expects DictationConfig"): + transcriber.config = aai.SyncTranscriptionConfig() + + +def test_context_manager_returns_self_and_closes(): + # Given a transcriber used as a context manager + with aai.DictationTranscriber() as transcriber: + # Then the bound value is the transcriber itself + assert isinstance(transcriber, aai.DictationTranscriber) + + # And leaving the block shuts the worker pool down + assert transcriber._executor._shutdown is True + + +def test_dictation_settings_defaults(): + # Given default settings + settings = aai.Settings() + + # Then the dictation host and the 300s per-operation timeout are in place + assert settings.dictation_base_url == "https://dictation.assemblyai.com" + assert settings.dictation_http_timeout == 300.0 + + +def test_the_buffered_endpoint_is_never_requested(httpx_mock: HTTPXMock): + """No entry point on this client posts to `/v1/transcribe`. + + The service still serves the buffered route, but the live route is the + only one this client opens, whichever way audio is submitted. + """ + for _ in range(4): # one per request; older pytest-httpx has no is_reusable + _mock_ok(httpx_mock) + + with aai.DictationTranscriber() as transcriber: + transcriber.transcribe_live(b"RIFFfake-wav-bytes") + transcriber.transcribe_live([b"RIFFfake", b"-wav-bytes"]) + transcriber.transcribe_live(io.BytesIO(b"RIFFfake-wav-bytes")) + with transcriber.open_live(aai.DictationConfig()) as session: + session.write(b"RIFFfake-wav-bytes") + session.result() + + requested = {str(request.url) for request in httpx_mock.get_requests()} + assert requested == {LIVE_URL} + + +# --- push-style sessions ------------------------------------------------------ + + +def _fake_live(monkeypatch, error: Optional[Exception] = None) -> dict: + """ + Replaces the transport with a fake that drains the producer the way a real + transport would, recording what it saw. An abort raised by the producer + propagates from `list()` exactly as it would from httpx. + """ + seen: dict = {"called": threading.Event(), "completed": False} + + def fake(client, **kwargs): + seen["called"].set() + seen["config"] = kwargs["config"] + seen["chunks"] = list(kwargs["chunks"]) + seen["completed"] = True + if error: + raise error + return aai.DictationResponse.parse_obj(_OK_RESPONSE) + + monkeypatch.setattr(api, "transcribe_live", fake) + return seen + + +def test_open_live_uploads_written_chunks_in_order(monkeypatch): + # Given a session fed from a callback-style producer + seen = _fake_live(monkeypatch) + + with aai.DictationTranscriber() as transcriber: + session = transcriber.open_live( + aai.DictationConfig(sample_rate=16000, channels=1) + ) + for piece in (b"\x00\x01", bytearray(b"\x02\x03"), memoryview(b"\x04\x05")): + session.write(piece) + + # When the audio ends and the result is awaited + result = session.result() + + # Then the chunks went out in order, as bytes, with the config, and the + # transcript came back parsed + assert seen["chunks"] == [b"\x00\x01", b"\x02\x03", b"\x04\x05"] + assert seen["config"] == {"sample_rate": 16000, "channels": 1} + assert result.text == "patient reports mild headache" + + +def test_open_live_starts_the_request_before_any_audio(monkeypatch): + # Given a session that has just been opened + seen = _fake_live(monkeypatch) + + with aai.DictationTranscriber() as transcriber: + session = transcriber.open_live() + + # Then the request is already in flight, which is the point: the + # connection and config go out while the speaker is still talking + assert seen["called"].wait(timeout=2.0) + assert not seen["completed"] + + session.result() + + +def test_open_live_uses_the_transcribers_default_config(monkeypatch): + # Given a transcriber with a default config and a session opened without one + seen = _fake_live(monkeypatch) + + with aai.DictationTranscriber( + config=aai.DictationConfig(llm_instruction="Summarize.") + ) as transcriber: + session = transcriber.open_live() + session.result() + + # Then the default config rode along + assert seen["config"] == {"llm_instruction": "Summarize."} + + +def test_live_session_context_manager_ends_the_audio(monkeypatch): + # Given audio written inside a with-block + _fake_live(monkeypatch) + + with aai.DictationTranscriber() as transcriber: + with transcriber.open_live() as session: + session.write(b"RIFF") + assert not session.closed + + # Then leaving the block closes the audio and the result is waiting + assert session.closed + assert session.result().text == "patient reports mild headache" + + +def test_live_session_rejects_writes_after_close(monkeypatch): + # Given a closed session + _fake_live(monkeypatch) + + with aai.DictationTranscriber() as transcriber: + session = transcriber.open_live() + session.close() + + # When more audio arrives, then it is refused rather than dropped silently + with pytest.raises(RuntimeError, match="closed"): + session.write(b"RIFF") + + session.result() + + +def test_live_session_rejects_text_at_write_time(monkeypatch): + # Given a producer handing out str + _fake_live(monkeypatch) + + with aai.DictationTranscriber() as transcriber: + session = transcriber.open_live() + + # When it is written, then the mistake is named immediately, not when + # the result is collected + with pytest.raises(TypeError, match="bytes"): + session.write("RIFF") + + session.result() + + +def test_live_session_abort_drops_the_request(monkeypatch): + # Given a session mid-upload + seen = _fake_live(monkeypatch) + + with aai.DictationTranscriber() as transcriber: + session = transcriber.open_live() + session.write(b"RIFF") + assert seen["called"].wait(timeout=2.0) + + # When it is aborted + session.abort() + + # Then the transport never completed the upload, there is no result, + # and a second abort is harmless + assert not seen["completed"] + assert session.closed + with pytest.raises(RuntimeError, match="aborted"): + session.result() + session.abort() + + +def test_live_session_aborts_when_the_block_raises(monkeypatch): + # Given a with-block that fails part-way through recording + seen = _fake_live(monkeypatch) + + with aai.DictationTranscriber() as transcriber: + with pytest.raises(ValueError): + with transcriber.open_live() as session: + session.write(b"RIFF") + raise ValueError("microphone unplugged") + + # Then the upload was dropped rather than transcribed + assert not seen["completed"] + with pytest.raises(RuntimeError, match="aborted"): + session.result() + + +def test_live_session_abort_after_completion_is_a_no_op(monkeypatch): + # Given a session whose result has already arrived + _fake_live(monkeypatch) + + with aai.DictationTranscriber() as transcriber: + session = transcriber.open_live() + result = session.result() + + # When it is aborted anyway, then nothing changes + session.abort() + assert session.result() is result + + +def test_live_session_rejects_other_products_configs(): + # Given the sync API's config type + with aai.DictationTranscriber() as transcriber: + with pytest.raises(TypeError, match="expects DictationConfig"): + transcriber.open_live(aai.SyncTranscriptionConfig()) + + +def test_live_session_roundtrip_over_http(httpx_mock: HTTPXMock): + # Given the real transport against a mocked endpoint + _mock_ok(httpx_mock) + + with aai.DictationTranscriber() as transcriber: + with transcriber.open_live() as session: + session.write(b"RIFF") + session.write(b"fake") + + # Then the request hit the live route and parsed like the pull path + result = session.result() + + assert result.text == "patient reports mild headache" + request = httpx_mock.get_requests()[0] + assert str(request.url) == LIVE_URL + assert request.headers.get("transfer-encoding") == "chunked" + + +def test_live_session_surfaces_server_errors(httpx_mock: HTTPXMock): + # Given a rejection the server may send mid-upload + httpx_mock.add_response( + url=LIVE_URL, + method="POST", + status_code=httpx.codes.SERVICE_UNAVAILABLE, + json={"status": 503, "title": "Capacity Exceeded", "detail": "no capacity"}, + ) + + with aai.DictationTranscriber() as transcriber: + session = transcriber.open_live() + session.write(b"RIFF") + + # When the result is collected, then it is the same error the pull path raises + with pytest.raises(aai.DictationError) as exc: + session.result() + + assert exc.value.error_code == "capacity_exceeded" diff --git a/tests/unit/test_dictation_async.py b/tests/unit/test_dictation_async.py new file mode 100644 index 00000000..ae839a9a --- /dev/null +++ b/tests/unit/test_dictation_async.py @@ -0,0 +1,689 @@ +"""Tests for the asyncio dictation client. + +The framing itself is covered by `test_dictation.py`; these pin the asyncio +twin's behaviour: same route, same parsing, same errors, plus the async +producer shapes and the `AsyncDictationLiveSession`. +""" + +import asyncio +import email +import io +import threading +from typing import List, Optional + +import httpx +import pytest +from pytest_httpx import HTTPXMock + +import assemblyai as aai +from assemblyai._multipart import _STREAM_READ_SIZE, aiter_chunks +from assemblyai.dictation.v1 import async_api + +pytestmark = pytest.mark.asyncio + +aai.settings.api_key = "test" + +LIVE_URL = f"{aai.settings.dictation_base_url}/v1/transcribe/live" +WARM_URL = f"{aai.settings.dictation_base_url}/v1/warm" + +_OK_RESPONSE = { + "text": "patient reports mild headache", + "words": [ + {"text": "patient", "confidence": 0.9}, + {"text": "reports", "confidence": 0.95}, + ], + "confidence": 0.92, + "llm_response": None, + "llm_error": None, + "audio_duration_ms": 400, + "session_id": "eb92c4ff-4bbb-429f-9b99-7279d7fe738f", + "request_time_ms": 243.7, + "sync_time_ms": 180.2, +} + + +def _mock_ok(httpx_mock: HTTPXMock, **overrides) -> None: + httpx_mock.add_response( + url=LIVE_URL, + method="POST", + status_code=httpx.codes.OK, + json={**_OK_RESPONSE, **overrides}, + ) + + +def _chunks(*pieces: bytes): + for piece in pieces: + yield piece + + +async def _achunks(*pieces: bytes): + for piece in pieces: + yield piece + + +def _parts(request: httpx.Request): + """Parses a recorded request body into its multipart parts, in order.""" + + raw = b"Content-Type: " + request.headers["content-type"].encode() + b"\r\n\r\n" + message = email.message_from_bytes(raw + request.content) + + return [ + (part.get_param("name", header="content-disposition"), part) + for part in message.get_payload() + ] + + +async def _asend(chunks, config: Optional[dict] = None) -> List[httpx.Request]: + """Runs `async_api.transcribe_live` against a transport that records the body.""" + + seen: List[httpx.Request] = [] + + def handler(request: httpx.Request) -> httpx.Response: + seen.append(request) + return httpx.Response(httpx.codes.OK, json=_OK_RESPONSE) + + async with httpx.AsyncClient(transport=httpx.MockTransport(handler)) as client: + await async_api.transcribe_live( + client, + base_url=aai.settings.dictation_base_url, + chunks=chunks, + filename="audio.wav", + audio_content_type="audio/wav", + config=config, + timeout=30.0, + ) + + return seen + + +def _capture(monkeypatch) -> dict: + """ + Replaces the async transport with a fake that records what the client + handed it, draining the chunks. Body-shape assertions on a recorded + request are not possible on the httpx floors — `pytest_httpx` there does + not read an async request's stream — so the client-level tests inspect + the transport call instead; the framing itself is covered by `_asend`. + """ + seen: dict = {"calls": []} + + async def fake(client, **kwargs): + kwargs["chunks"] = [chunk async for chunk in aiter_chunks(kwargs["chunks"])] + seen["calls"].append(kwargs) + return aai.DictationResponse.parse_obj(_OK_RESPONSE) + + monkeypatch.setattr(async_api, "transcribe_live", fake) + return seen + + +class _SpyFile(io.BytesIO): + """A file object that records the thread of every read asked of it.""" + + def __init__(self, payload: bytes) -> None: + super().__init__(payload) + self.reads: List[int] = [] + self.threads: List[threading.Thread] = [] + + def read(self, size: int = -1) -> bytes: + self.reads.append(size) + self.threads.append(threading.current_thread()) + return super().read(size) + + +# --- body shape ------------------------------------------------------------- + + +async def test_body_matches_the_threaded_twin(): + # Given the same audio through the async transport + # When the body is built + request = (await _asend(_achunks(b"RIFF", b"fake", b"wav"), config={"a": 1}))[0] + + # Then it is framed identically: config first, chunked, no model header + parts = dict(_parts(request)) + assert [name for name, _ in _parts(request)] == ["config", "audio"] + assert parts["audio"].get_payload(decode=True) == b"RIFFfakewav" + assert parts["config"].get_payload(decode=True) == b'{"a": 1}' + assert request.headers.get("transfer-encoding") == "chunked" + assert str(request.url) == LIVE_URL + assert "X-AAI-Model" not in request.headers + + +async def test_body_sends_an_empty_config_part_when_there_is_no_config(): + request = (await _asend(_achunks(b"RIFF")))[0] + + parts = _parts(request) + assert [name for name, _ in parts] == ["config", "audio"] + assert dict(parts)["config"].get_payload(decode=True) == b"{}" + + +async def test_body_accepts_a_sync_iterable(): + # Given a plain iterable handed to the async transport + request = (await _asend(_chunks(b"RIFF", b"fake")))[0] + + # Then one producer can feed both transports + assert dict(_parts(request))["audio"].get_payload(decode=True) == b"RIFFfake" + + +async def test_body_reads_a_file_object_off_the_event_loop(): + # Given a file object handed to the async transport + source = _SpyFile(b"RIFF" + bytes(_STREAM_READ_SIZE + 1)) + + # When it is streamed + request = (await _asend(source))[0] + + # Then every read ran in a worker thread, so a blocking read cannot stall + # the loop, and the body is intact + assert source.reads == [_STREAM_READ_SIZE] * 3 + assert all(t is not threading.main_thread() for t in source.threads) + assert dict(_parts(request))["audio"].get_payload(decode=True) == source.getvalue() + + +# --- transcriber behaviour -------------------------------------------------- + + +async def test_transcribe_live_parses_response(httpx_mock: HTTPXMock): + # Given a mocked live endpoint + _mock_ok(httpx_mock) + + # When streaming from an async producer + async with aai.AsyncDictationTranscriber() as transcriber: + result = await transcriber.transcribe_live(_achunks(b"RIFF", b"fake")) + + # Then the response is parsed into a DictationResponse + assert isinstance(result, aai.DictationResponse) + assert result.text == "patient reports mild headache" + assert result.session_id == _OK_RESPONSE["session_id"] + assert result.words[1].text == "reports" + assert result.request_time_ms == 243.7 + assert result.sync_time_ms == 180.2 + + +async def test_transcribe_live_posts_to_the_dictation_host(httpx_mock: HTTPXMock): + # Given a mocked live endpoint + _mock_ok(httpx_mock) + + # When streaming + async with aai.AsyncDictationTranscriber() as transcriber: + await transcriber.transcribe_live(_achunks(b"RIFF")) + + # Then the request goes to the dictation host with the raw key, no + # routing header, and the dictation timeout + request = httpx_mock.get_requests()[0] + assert str(request.url) == LIVE_URL + assert request.headers["authorization"] == "test" + assert "X-AAI-Model" not in request.headers + assert request.extensions["timeout"]["read"] == aai.settings.dictation_http_timeout + + +async def test_complete_audio_goes_over_the_live_connection(httpx_mock: HTTPXMock): + # Given audio the caller already holds whole + _mock_ok(httpx_mock) + + # When streaming it + async with aai.AsyncDictationTranscriber() as transcriber: + await transcriber.transcribe_live(b"RIFFfake-wav-bytes") + + # Then it takes the live route, framed without a length + request = httpx_mock.get_requests()[0] + assert str(request.url) == LIVE_URL + assert "content-length" not in request.headers + assert request.headers["content-type"].startswith("multipart/form-data") + + +async def test_complete_audio_is_one_chunk(monkeypatch): + # Given audio the caller already holds whole + seen = _capture(monkeypatch) + + # When streaming it + async with aai.AsyncDictationTranscriber() as transcriber: + await transcriber.transcribe_live(b"RIFFfake-wav-bytes") + + # Then the transport receives it as a single chunk, typed as WAV + (call,) = seen["calls"] + assert call["chunks"] == [b"RIFFfake-wav-bytes"] + assert call["filename"] == "audio.wav" + assert call["audio_content_type"] == "audio/wav" + assert call["config"] is None + + +async def test_transcribe_live_reads_a_path_off_the_loop(monkeypatch, tmp_path): + # Given a local MP3 file + seen = _capture(monkeypatch) + audio_file = tmp_path / "note.mp3" + audio_file.write_bytes(b"ID3fake-mp3-bytes") + + # When streaming the path + async with aai.AsyncDictationTranscriber() as transcriber: + await transcriber.transcribe_live(str(audio_file)) + + # Then the file is sent whole with its true Content-Type and name + (call,) = seen["calls"] + assert call["chunks"] == [b"ID3fake-mp3-bytes"] + assert call["audio_content_type"] == "audio/mpeg" + assert call["filename"] == "note.mp3" + + +async def test_transcribe_live_uses_default_config_and_per_call_override(monkeypatch): + # Given a transcriber with a default config + seen = _capture(monkeypatch) + default = aai.DictationConfig(llm_instruction="default instruction") + + async with aai.AsyncDictationTranscriber(config=default) as transcriber: + # When streaming without a per-call config, and with an override + await transcriber.transcribe_live(_achunks(b"RIFF")) + override = aai.DictationConfig(llm_instruction="override instruction") + await transcriber.transcribe_live(_achunks(b"RIFF"), config=override) + + # Then the default applies to the first call and the override to the second + first, second = seen["calls"] + assert first["config"] == {"llm_instruction": "default instruction"} + assert second["config"] == {"llm_instruction": "override instruction"} + + +async def test_transcribe_live_marks_pcm_from_config(monkeypatch): + # Given a config carrying the fields only raw PCM needs + seen = _capture(monkeypatch) + config = aai.DictationConfig(sample_rate=16000, channels=1) + + # When streaming + async with aai.AsyncDictationTranscriber() as transcriber: + await transcriber.transcribe_live(_achunks(b"\x00\x01"), config=config) + + # Then the audio part selects the PCM decoder and the config carries both + (call,) = seen["calls"] + assert call["audio_content_type"] == "audio/pcm" + assert call["filename"] == "audio.pcm" + assert call["config"] == {"sample_rate": 16000, "channels": 1} + + +async def test_transcribe_live_requires_both_pcm_fields(): + # Given a config with only half of what raw PCM needs + config = aai.DictationConfig(sample_rate=16000) + + # When streaming, then it fails locally before any request + async with aai.AsyncDictationTranscriber() as transcriber: + with pytest.raises(ValueError, match="sample_rate and channels"): + await transcriber.transcribe_live(_achunks(b"\x00"), config=config) + + +async def test_transcribe_live_rejects_url(): + # Given an http URL as input + async with aai.AsyncDictationTranscriber() as transcriber: + # When streaming, then the message names this client, not its + # threaded twin + with pytest.raises( + ValueError, match="AsyncDictationTranscriber does not accept URLs" + ): + await transcriber.transcribe_live("https://example.com/audio.wav") + + +async def test_transcribe_live_rejects_an_unsupported_source(): + async with aai.AsyncDictationTranscriber() as transcriber: + with pytest.raises(TypeError, match="unsupported audio source type"): + await transcriber.transcribe_live(object()) + + +async def test_transcribe_live_gather_runs_concurrently(httpx_mock: HTTPXMock): + # Given a mocked live endpoint answering twice + _mock_ok(httpx_mock) + _mock_ok(httpx_mock) + + # When fanning two clips out with asyncio.gather on one transcriber + async with aai.AsyncDictationTranscriber() as transcriber: + results = await asyncio.gather( + transcriber.transcribe_live(b"RIFFone"), + transcriber.transcribe_live(_achunks(b"RIFFtwo")), + ) + + # Then both finish and parse + assert [result.text for result in results] == [ + "patient reports mild headache", + "patient reports mild headache", + ] + + +async def test_final_text_prefers_llm_response(httpx_mock: HTTPXMock): + # Given a response carrying an LLM rewrite + _mock_ok(httpx_mock, llm_response="S: Mild headache.") + + # When streaming with an llm_instruction + config = aai.DictationConfig(llm_instruction="Format this as a SOAP note.") + async with aai.AsyncDictationTranscriber() as transcriber: + result = await transcriber.transcribe_live(_achunks(b"RIFF"), config=config) + + # Then final_text is the rewrite and the raw transcript is still there + assert result.final_text == "S: Mild headache." + assert result.text == "patient reports mild headache" + + +async def test_error_envelope_maps_to_dictation_error(httpx_mock: HTTPXMock): + # Given a server rejecting the audio with an {error, error_code} body + httpx_mock.add_response( + url=LIVE_URL, + method="POST", + status_code=415, + json={"error": "unsupported media type", "error_code": "bad_audio"}, + ) + + # When streaming, then a DictationError carries code + status + async with aai.AsyncDictationTranscriber() as transcriber: + with pytest.raises(aai.DictationError) as exc_info: + await transcriber.transcribe_live(_achunks(b"RIFF")) + + error = exc_info.value + assert error.status_code == 415 + assert error.error_code == "bad_audio" + assert not isinstance(error, aai.SyncTranscriptError) + + +async def test_rate_limit_surfaces_retry_after(httpx_mock: HTTPXMock): + # Given a rate-limit response with a Retry-After header + httpx_mock.add_response( + url=LIVE_URL, + method="POST", + status_code=429, + json={"status": 429, "title": "Too Many Requests", "detail": "slow down"}, + headers={"Retry-After": "5"}, + ) + + # When streaming, then retry_after and the snake_cased title are parsed + async with aai.AsyncDictationTranscriber() as transcriber: + with pytest.raises(aai.DictationError) as exc_info: + await transcriber.transcribe_live(_achunks(b"RIFF")) + + error = exc_info.value + assert error.status_code == 429 + assert error.error_code == "too_many_requests" + assert error.retry_after == 5 + + +async def test_warm_opens_connection(httpx_mock: HTTPXMock): + # Given a mocked warm endpoint + httpx_mock.add_response(url=WARM_URL, method="GET", status_code=httpx.codes.OK) + + # When warming the transcriber + async with aai.AsyncDictationTranscriber() as transcriber: + warmed = await transcriber.warm() + + # Then it returns True and probes the dictation warm route + assert warmed is True + request = httpx_mock.get_requests()[0] + assert str(request.url) == WARM_URL + assert request.method == "GET" + assert "X-AAI-Model" not in request.headers + + +async def test_warm_returns_true_on_non_200(httpx_mock: HTTPXMock): + httpx_mock.add_response(url=WARM_URL, method="GET", status_code=404) + + async with aai.AsyncDictationTranscriber() as transcriber: + assert await transcriber.warm() is True + + +async def test_warm_returns_false_on_transport_error(httpx_mock: HTTPXMock): + httpx_mock.add_exception(httpx.ConnectError("connection refused")) + + async with aai.AsyncDictationTranscriber() as transcriber: + assert await transcriber.warm() is False + + +async def test_api_key_constructor_builds_own_client(httpx_mock: HTTPXMock): + # Given a transcriber constructed with an explicit key + _mock_ok(httpx_mock) + + # When streaming + async with aai.AsyncDictationTranscriber(api_key="per-call-key") as transcriber: + await transcriber.transcribe_live(_achunks(b"RIFF")) + + # Then that key authenticates the request + assert httpx_mock.get_requests()[0].headers["authorization"] == "per-call-key" + + +async def test_rejects_sync_transcription_config(): + with pytest.raises(TypeError, match="expects DictationConfig"): + aai.AsyncDictationTranscriber(config=aai.SyncTranscriptionConfig()) + + +async def test_context_manager_closes_owned_client(): + # Given a transcriber that created its own client + async with aai.AsyncDictationTranscriber() as transcriber: + assert isinstance(transcriber, aai.AsyncDictationTranscriber) + assert not transcriber.client.http_client.is_closed + + # Then leaving the block closes the owned connection pool + assert transcriber.client.http_client.is_closed + + +async def test_aclose_leaves_shared_client_open(): + # Given a transcriber built on a caller-owned client + async with aai.AsyncClient(settings=aai.settings) as client: + transcriber = aai.AsyncDictationTranscriber(client=client) + + # When closing the transcriber + await transcriber.aclose() + + # Then the shared pool stays open — its creator closes it + assert not client.http_client.is_closed + + assert client.http_client.is_closed + + +# --- push-style sessions ------------------------------------------------------ + + +def _fake_live(monkeypatch, error: Optional[Exception] = None) -> dict: + """Replaces the async transport with a fake that drains the producer.""" + + seen: dict = {"called": False, "completed": False} + + async def fake(client, **kwargs): + seen["called"] = True + seen["config"] = kwargs["config"] + seen["chunks"] = [chunk async for chunk in kwargs["chunks"]] + seen["completed"] = True + if error: + raise error + return aai.DictationResponse.parse_obj(_OK_RESPONSE) + + monkeypatch.setattr(async_api, "transcribe_live", fake) + return seen + + +async def test_open_live_uploads_written_chunks_in_order(monkeypatch): + # Given an asyncio session fed from a callback-style producer + seen = _fake_live(monkeypatch) + + async with aai.AsyncDictationTranscriber() as transcriber: + session = transcriber.open_live( + aai.DictationConfig(sample_rate=16000, channels=1) + ) + for piece in (b"\x00\x01", bytearray(b"\x02\x03")): + session.write(piece) + + # When the audio ends and the result is awaited + result = await session.result() + + # Then it matches the threaded twin + assert seen["chunks"] == [b"\x00\x01", b"\x02\x03"] + assert seen["config"] == {"sample_rate": 16000, "channels": 1} + assert result.text == "patient reports mild headache" + + +async def test_open_live_starts_the_request_before_any_audio(monkeypatch): + # Given a session that has just been opened + seen = _fake_live(monkeypatch) + + async with aai.AsyncDictationTranscriber() as transcriber: + session = transcriber.open_live() + + # Then the request reaches the transport with the audio still open, + # which is the point: the connection and config go out while the + # speaker is still talking. The source is resolved in a worker thread + # first, so give the task a moment rather than a single loop tick. + for _ in range(200): + if seen["called"]: + break + await asyncio.sleep(0.01) + assert seen["called"] + assert not seen["completed"] + + await session.result() + + +async def test_live_session_context_manager_ends_the_audio(monkeypatch): + # Given audio written inside an async with-block + _fake_live(monkeypatch) + + async with aai.AsyncDictationTranscriber() as transcriber: + async with transcriber.open_live() as session: + session.write(b"RIFF") + assert not session.closed + + # Then leaving the block closes the audio and the result is waiting + assert session.closed + assert (await session.result()).text == "patient reports mild headache" + + +async def test_live_session_rejects_writes_after_close(monkeypatch): + # Given a closed session + _fake_live(monkeypatch) + + async with aai.AsyncDictationTranscriber() as transcriber: + session = transcriber.open_live() + session.close() + + # When more audio arrives, then it is refused + with pytest.raises(RuntimeError, match="closed"): + session.write(b"RIFF") + + await session.result() + + +async def test_live_session_rejects_text_at_write_time(monkeypatch): + _fake_live(monkeypatch) + + async with aai.AsyncDictationTranscriber() as transcriber: + session = transcriber.open_live() + + with pytest.raises(TypeError, match="bytes"): + session.write("RIFF") + + await session.result() + + +async def test_live_session_abort_drops_the_request(monkeypatch): + # Given a session mid-upload + seen = _fake_live(monkeypatch) + + async with aai.AsyncDictationTranscriber() as transcriber: + session = transcriber.open_live() + session.write(b"RIFF") + await asyncio.sleep(0) # let the task start consuming + + # When it is aborted + await session.abort() + + # Then the transport never completed the upload, there is no result, + # and a second abort is harmless + assert not seen["completed"] + assert session.closed + with pytest.raises(RuntimeError, match="aborted"): + await session.result() + await session.abort() + + +async def test_live_session_aborts_when_the_block_raises(monkeypatch): + # Given an async with-block that fails mid-recording + seen = _fake_live(monkeypatch) + + async with aai.AsyncDictationTranscriber() as transcriber: + with pytest.raises(ValueError): + async with transcriber.open_live() as session: + session.write(b"RIFF") + raise ValueError("call dropped") + + # Then the upload was dropped rather than transcribed + assert not seen["completed"] + with pytest.raises(RuntimeError, match="aborted"): + await session.result() + + +async def test_live_session_abort_after_completion_is_a_no_op(monkeypatch): + _fake_live(monkeypatch) + + async with aai.AsyncDictationTranscriber() as transcriber: + session = transcriber.open_live() + result = await session.result() + + await session.abort() + assert await session.result() is result + + +async def test_live_session_rejects_other_products_configs(): + async with aai.AsyncDictationTranscriber() as transcriber: + with pytest.raises(TypeError, match="expects DictationConfig"): + transcriber.open_live(aai.SyncTranscriptionConfig()) + + +async def test_live_session_roundtrip_over_http(httpx_mock: HTTPXMock): + # Given the real transport against a mocked endpoint + _mock_ok(httpx_mock) + + async with aai.AsyncDictationTranscriber() as transcriber: + async with transcriber.open_live() as session: + session.write(b"RIFF") + + result = await session.result() + + # Then the request hit the live route and parsed like the pull path + assert result.text == "patient reports mild headache" + request = httpx_mock.get_requests()[0] + assert str(request.url) == LIVE_URL + assert request.headers.get("transfer-encoding") == "chunked" + + +async def test_live_session_surfaces_server_errors(httpx_mock: HTTPXMock): + # Given a rejection + httpx_mock.add_response( + url=LIVE_URL, + method="POST", + status_code=httpx.codes.TOO_MANY_REQUESTS, + json={"status": 429, "title": "Rate Limited", "detail": "slow down"}, + headers={"Retry-After": "3"}, + ) + + async with aai.AsyncDictationTranscriber() as transcriber: + session = transcriber.open_live() + session.write(b"RIFF") + + # When the result is collected, then it is the same error the pull path raises + with pytest.raises(aai.DictationError) as exc: + await session.result() + + assert exc.value.status_code == 429 + assert exc.value.retry_after == 3 + + +async def test_the_buffered_endpoint_is_never_requested(httpx_mock: HTTPXMock): + """No entry point on the asyncio client posts to `/v1/transcribe`.""" + for _ in range(4): # one per request; older pytest-httpx has no is_reusable + _mock_ok(httpx_mock) + + async with aai.AsyncDictationTranscriber() as transcriber: + await transcriber.transcribe_live(b"RIFFfake-wav-bytes") + await transcriber.transcribe_live(_achunks(b"RIFFfake", b"-wav-bytes")) + await transcriber.transcribe_live(io.BytesIO(b"RIFFfake-wav-bytes")) + async with transcriber.open_live() as session: + session.write(b"RIFFfake-wav-bytes") + await session.result() + + requested = {str(request.url) for request in httpx_mock.get_requests()} + assert requested == {LIVE_URL} + + +def test_open_live_needs_a_running_loop(): + # Given no event loop + transcriber = aai.AsyncDictationTranscriber() + + # When a session is opened, then the mistake is named rather than deferred + with pytest.raises(RuntimeError, match="no running event loop"): + transcriber.open_live() diff --git a/tests/unit/test_sync_live.py b/tests/unit/test_sync_live.py index 510ba49e..c7cbaebf 100644 --- a/tests/unit/test_sync_live.py +++ b/tests/unit/test_sync_live.py @@ -17,8 +17,8 @@ from pytest_httpx import HTTPXMock import assemblyai as aai +from assemblyai._multipart import _STREAM_READ_SIZE from assemblyai.sync.v1 import api, async_api -from assemblyai.sync.v1._multipart import _STREAM_READ_SIZE aai.settings.api_key = "test"