From a7a1c4d307d25feda0ad7f12d08997da42d19aff Mon Sep 17 00:00:00 2001 From: Adam Urpsis Date: Thu, 15 Oct 2026 19:06:27 -0600 Subject: [PATCH] feat(dictation): add rewrite_mode to DictationConfig Adds `DictationConfig.rewrite_mode` and the `DictationRewriteMode` enum (`fast` | `smart`), selecting the model tier that runs the Dictation API's transcript rewrite. Unset, the field is omitted and the server default (`fast`) applies. Co-Authored-By: Claude Opus 5.5 (1M context) Claude-Session: https://claude.ai/code/session_01TNBmFAJar33SLWN2rKnBzJ --- CLAUDE.md | 2 +- README.md | 10 ++++++++++ assemblyai/__init__.py | 2 ++ assemblyai/__version__.py | 2 +- assemblyai/dictation/__init__.py | 2 ++ assemblyai/dictation/v1/__init__.py | 2 ++ assemblyai/dictation/v1/models.py | 21 +++++++++++++++++++ tests/unit/test_dictation.py | 31 +++++++++++++++++++++++++++++ 8 files changed, 70 insertions(+), 2 deletions(-) diff --git a/CLAUDE.md b/CLAUDE.md index 7e055617..508b60e3 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -50,7 +50,7 @@ aai.settings.api_key = "your-key" - `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.DictationConfig` — Dictation options: `sample_rate`, `channels`, `language_codes`, `keyterms_prompt`, `llm_instruction`, `rewrite_mode` (`aai.DictationRewriteMode.fast` default | `smart`). 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 e243627c..466736a4 100644 --- a/README.md +++ b/README.md @@ -779,6 +779,15 @@ if result.llm_error: The instruction is capped at 2048 characters. +`rewrite_mode` picks the model tier that runs the rewrite. `fast` is the default; `smart` follows detailed formatting instructions more reliably, at higher latency: + +```python +config = aai.DictationConfig( + llm_instruction="Lowercase unless a name, no final punctuation.", + rewrite_mode=aai.DictationRewriteMode.smart, +) +``` +
@@ -794,6 +803,7 @@ config = aai.DictationConfig( stt_prompt="A doctor dictating a visit note.", # what the audio is about; ≤6000 chars keyterms_prompt=["AssemblyAI", "Universal"], # bias the decoder; ≤100 terms / 8000 chars llm_instruction="Fix punctuation only.", # optional LLM pass + rewrite_mode="smart", # rewrite model tier: "fast" (default) or "smart" ) result = aai.DictationTranscriber().transcribe_live("./note.wav", config=config) diff --git a/assemblyai/__init__.py b/assemblyai/__init__.py index a4761f17..5bb3963d 100644 --- a/assemblyai/__init__.py +++ b/assemblyai/__init__.py @@ -6,6 +6,7 @@ DictationConfig, DictationError, DictationResponse, + DictationRewriteMode, DictationTranscriber, DictationWord, ) @@ -139,6 +140,7 @@ def __getattr__(name: str): "DictationConfig", "DictationError", "DictationResponse", + "DictationRewriteMode", "DictationTranscriber", "DictationWord", "Entity", diff --git a/assemblyai/__version__.py b/assemblyai/__version__.py index 31e744e4..dc79f8f0 100644 --- a/assemblyai/__version__.py +++ b/assemblyai/__version__.py @@ -1 +1 @@ -__version__ = "1.6.3" +__version__ = "1.6.4" diff --git a/assemblyai/dictation/__init__.py b/assemblyai/dictation/__init__.py index 53f74425..edee5569 100644 --- a/assemblyai/dictation/__init__.py +++ b/assemblyai/dictation/__init__.py @@ -16,6 +16,7 @@ DictationError, DictationLiveSession, DictationResponse, + DictationRewriteMode, DictationTranscriber, DictationWord, ) @@ -29,6 +30,7 @@ "DictationError", "DictationLiveSession", "DictationResponse", + "DictationRewriteMode", "DictationTranscriber", "DictationWord", ] diff --git a/assemblyai/dictation/v1/__init__.py b/assemblyai/dictation/v1/__init__.py index 2a5c8b2f..219646d5 100644 --- a/assemblyai/dictation/v1/__init__.py +++ b/assemblyai/dictation/v1/__init__.py @@ -5,6 +5,7 @@ DictationConfig, DictationError, DictationResponse, + DictationRewriteMode, DictationWord, ) @@ -17,6 +18,7 @@ "DictationError", "DictationLiveSession", "DictationResponse", + "DictationRewriteMode", "DictationTranscriber", "DictationWord", ] diff --git a/assemblyai/dictation/v1/models.py b/assemblyai/dictation/v1/models.py index a0fc3886..1d6dc419 100644 --- a/assemblyai/dictation/v1/models.py +++ b/assemblyai/dictation/v1/models.py @@ -4,6 +4,7 @@ the transcribers raise, live here alongside the client that uses them. """ +from enum import Enum from typing import List, Optional try: @@ -23,6 +24,7 @@ "DictationConfig", "DictationError", "DictationResponse", + "DictationRewriteMode", "DictationWord", ] @@ -56,6 +58,20 @@ def __init__( _DICTATION_MAX_STT_PROMPT_LEN = 6000 +class DictationRewriteMode(str, Enum): + """Model tier that runs the transcript rewrite.""" + + fast = "fast" + "The default: the lowest-latency rewrite." + + smart = "smart" + """Follows detailed formatting instructions more reliably, at higher + latency.""" + + def __str__(self): + return self.value + + class DictationConfig(BaseModel): """ Options for a Dictation API request. @@ -108,6 +124,11 @@ class DictationConfig(BaseModel): "Format this as a SOAP note." The rewritten text comes back as `llm_response`; the raw transcript stays in `text`. Max 2048 characters.""" + rewrite_mode: Optional[DictationRewriteMode] = None + """Model tier that runs the rewrite: `fast` or `smart`. `smart` follows + detailed formatting instructions more reliably at higher latency. Defaults + to None, which leaves the server default (`fast`).""" + if pydantic_v2: model_config = ConfigDict(extra="forbid") diff --git a/tests/unit/test_dictation.py b/tests/unit/test_dictation.py index 0258f983..ba8fafe4 100644 --- a/tests/unit/test_dictation.py +++ b/tests/unit/test_dictation.py @@ -373,6 +373,7 @@ def test_config_json_excludes_unset_fields(httpx_mock: HTTPXMock): assert b'"language_codes"' not in body assert b'"sample_rate"' not in body assert b'"keyterms_prompt"' not in body + assert b'"rewrite_mode"' not in body def test_transcribe_live_sends_keyterms_prompt_and_language_codes( @@ -620,6 +621,36 @@ def test_llm_instruction_too_long_raises(): aai.DictationConfig(llm_instruction="x" * 3000) +def test_transcribe_live_sends_rewrite_mode(httpx_mock: HTTPXMock): + # Given a config asking for the smart rewrite tier + _mock_ok(httpx_mock) + config = aai.DictationConfig( + llm_instruction="Format this as a SOAP note.", + rewrite_mode=aai.DictationRewriteMode.smart, + ) + + # When streaming + aai.DictationTranscriber().transcribe_live(_chunks(b"RIFF"), config=config) + + # Then the config JSON carries the mode's wire value + body = httpx_mock.get_requests()[0].read() + assert b'"rewrite_mode": "smart"' in body + + +def test_rewrite_mode_accepts_the_wire_string(): + # Given rewrite_mode passed as a plain string + config = aai.DictationConfig(rewrite_mode="fast") + + # Then it is parsed into the enum + assert config.rewrite_mode is aai.DictationRewriteMode.fast + + +def test_rewrite_mode_rejects_an_unknown_tier(): + # Given a tier the API does not define + with pytest.raises(ValueError): + aai.DictationConfig(rewrite_mode="gemma") + + 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(