feat(dictation): add rewrite_mode to DictationConfig

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

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,
+)
+```
+
 </details>
 
 <details>
@@ -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)

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

assemblyai/__version__.py

@@ -1 +1 @@
-__version__ = "1.6.3"
+__version__ = "1.6.4"

assemblyai/dictation/__init__.py

@@ -16,6 +16,7 @@
     DictationError,
     DictationLiveSession,
     DictationResponse,
+    DictationRewriteMode,
     DictationTranscriber,
     DictationWord,
 )
@@ -29,6 +30,7 @@
     "DictationError",
     "DictationLiveSession",
     "DictationResponse",
+    "DictationRewriteMode",
     "DictationTranscriber",
     "DictationWord",
 ]

assemblyai/dictation/v1/__init__.py

@@ -5,6 +5,7 @@
     DictationConfig,
     DictationError,
     DictationResponse,
+    DictationRewriteMode,
     DictationWord,
 )
 
@@ -17,6 +18,7 @@
     "DictationError",
     "DictationLiveSession",
     "DictationResponse",
+    "DictationRewriteMode",
     "DictationTranscriber",
     "DictationWord",
 ]

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

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(