Skip to content

Commit bba995d

Browse files
committed
Add assembly stream --save-audio PATH to tee streamed PCM to a WAV
Tee exactly the bytes sent to the streaming API to a 16-bit mono WAV while transcribing, without altering the live transcript. This lets a downstream consumer (e.g. an ensemble that compares live turns against an async re-transcribe) keep the audio without owning capture itself. - New streaming/record.py: tee_wav() writes each chunk at the source's true rate and yields it onward; the header is patched on exhaustion or early close (Ctrl-C) so a partial recording is still a valid WAV. validate_target() rejects a missing parent dir up front, before credentials. - StreamSession gains save_audio; the single-source path tees in stream_one. - Rejected combinations (clear usage errors): --system-audio/--system-audio-only (two streams can't share one file), --from-stdin (batch is many sources), and --show-code (generated SDK code doesn't tee). https://claude.ai/code/session_01MiPAW6mr1pYQuAE123HGxD
1 parent 30a4bb8 commit bba995d

7 files changed

Lines changed: 265 additions & 2 deletions

File tree

aai_cli/commands/stream/__init__.py

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -33,6 +33,7 @@
3333
("Stream a list of files in turn", "ls *.wav | assembly stream --from-stdin"),
3434
("Stream the hosted sample", "assembly stream --sample"),
3535
("Label speakers in the live transcript", "assembly stream --speaker-labels"),
36+
("Save a WAV of the audio while streaming", "assembly stream --save-audio out.wav"),
3637
(
3738
"Boost domain terms with keyterm prompts",
3839
'assembly stream --keyterms-prompt "AssemblyAI" --keyterms-prompt "Claude"',
@@ -82,6 +83,16 @@ def stream(
8283
help="macOS only: stream system/app audio without the microphone",
8384
rich_help_panel=help_panels.OPT_CAPTURE,
8485
),
86+
save_audio: Path | None = typer.Option(
87+
None,
88+
"--save-audio",
89+
help="Tee the streamed PCM to PATH as a 16-bit mono WAV while transcribing",
90+
rich_help_panel=help_panels.OPT_CAPTURE,
91+
dir_okay=False,
92+
# Click guardrail; flipping it changes no behavior a unit test can observe
93+
# (and the writable check is a no-op under the test runner's root uid).
94+
writable=True, # pragma: no mutate
95+
),
8596
# model & input
8697
speech_model: SpeechModel = typer.Option(
8798
DEFAULT_SPEECH_MODEL,
@@ -355,5 +366,6 @@ def stream(
355366
config_file=config_file,
356367
output_field=output_field,
357368
show_code=show_code,
369+
save_audio=save_audio,
358370
)
359371
run_with_options(ctx, stream_exec.run_stream, opts, json=json_out)

aai_cli/commands/stream/_exec.py

Lines changed: 21 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -22,7 +22,7 @@
2222
from aai_cli.core import choices, client, config_builder, stdio, youtube
2323
from aai_cli.core.errors import UsageError, mutually_exclusive
2424
from aai_cli.core.microphone import MicrophoneSource
25-
from aai_cli.streaming import turn_presets
25+
from aai_cli.streaming import record, turn_presets
2626
from aai_cli.streaming.macos import MacSystemAudioSource
2727
from aai_cli.streaming.render import StreamRenderer
2828
from aai_cli.streaming.session import (
@@ -85,6 +85,7 @@ class StreamOptions:
8585
config_file: Path | None
8686
output_field: choices.TextOrJson | None
8787
show_code: bool
88+
save_audio: Path | None
8889

8990
def source_options(self) -> SourceOptions:
9091
"""The audio-input subset, in the shape the validation/dispatch helpers read."""
@@ -245,6 +246,11 @@ def _collect_batch_sources(opts: StreamOptions, *, text_mode: bool) -> list[str]
245246
("--show-code", opts.show_code),
246247
suggestion="--show-code renders one source; pass a single file or URL.",
247248
)
249+
mutually_exclusive(
250+
("--from-stdin", True),
251+
("--save-audio", opts.save_audio is not None),
252+
suggestion="--save-audio tees one stream; run a single source to record it.",
253+
)
248254
mutually_exclusive(
249255
("--llm", bool(opts.llm_prompt)),
250256
("-o text", text_mode),
@@ -305,12 +311,25 @@ def run_stream(opts: StreamOptions, state: AppState, *, json_mode: bool) -> None
305311
base_flags = opts.base_flags()
306312

307313
if opts.show_code:
314+
if opts.save_audio is not None:
315+
raise UsageError(
316+
"--save-audio cannot be combined with --show-code; the generated SDK "
317+
"code does not tee audio to disk."
318+
)
308319
_print_show_code(opts, sources, base_flags, text_mode=text_mode)
309320
return
310321

311322
# Validate the requested sources (including that a local file exists) before
312323
# credentials, so a typo'd path reads as "file not found" — not as a login.
313324
validate_sources(sources, has_llm=bool(opts.llm_prompt), text_mode=text_mode)
325+
if opts.save_audio is not None:
326+
if sources.from_system_audio:
327+
raise UsageError(
328+
"--save-audio cannot be combined with --system-audio; the mic and system "
329+
"streams can't share one file.",
330+
suggestion="Record a single source (mic, file, URL, or - on stdin).",
331+
)
332+
record.validate_target(opts.save_audio)
314333
if sources.from_file and not sources.from_stdin:
315334
client.resolve_audio_source(sources.source, sample=sources.sample)
316335
api_key = state.resolve_api_key()
@@ -326,6 +345,7 @@ def run_stream(opts: StreamOptions, state: AppState, *, json_mode: bool) -> None
326345
llm_prompts=llm_prompts,
327346
model=opts.model,
328347
max_tokens=opts.max_tokens,
348+
save_audio=opts.save_audio,
329349
llm_interval=opts.llm_interval,
330350
)
331351
_dispatch(session, sources)

aai_cli/streaming/record.py

Lines changed: 67 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,67 @@
1+
"""Tee streamed PCM to a WAV file — backs `assembly stream --save-audio PATH`.
2+
3+
The whole point is a verbatim recording of exactly the bytes sent to the streaming
4+
API, so a caller (e.g. an ensemble that compares the live turns against an async
5+
re-transcribe) can keep the audio without owning capture itself. The tee never alters
6+
what's transcribed: it writes each chunk to disk and yields it onward unchanged.
7+
"""
8+
9+
from __future__ import annotations
10+
11+
import wave
12+
from collections.abc import Generator, Iterable
13+
from pathlib import Path
14+
15+
from aai_cli.core.errors import CLIError
16+
from aai_cli.streaming.sources import PCM16_SAMPLE_WIDTH_BYTES
17+
18+
19+
def validate_target(path: Path) -> None:
20+
"""Reject a ``--save-audio`` path whose parent directory is missing, before streaming.
21+
22+
Run before credentials/audio are opened so a bad path reads as a path error up
23+
front, not after a session has already started recording into the void.
24+
"""
25+
parent = path.parent
26+
if not parent.is_dir():
27+
raise CLIError(
28+
f"Cannot save audio to {path}: {parent} is not a directory.",
29+
error_type="save_audio_path",
30+
exit_code=2,
31+
suggestion="Create the directory first, or pass a path under an existing one.",
32+
)
33+
34+
35+
def tee_wav(audio: Iterable[bytes], path: Path, *, rate: int) -> Generator[bytes, None, None]:
36+
"""Yield every PCM16 chunk from ``audio`` unchanged while writing it to ``path`` as WAV.
37+
38+
The recording is mono 16-bit PCM at ``rate`` — the same shape the streaming API
39+
receives. The header's length fields are patched when the iterable is exhausted or
40+
closed early (Ctrl-C raises ``GeneratorExit`` at the ``yield``), so even an
41+
interrupted run leaves a valid, playable WAV of the audio captured so far.
42+
"""
43+
try:
44+
# Open the handle ourselves (rather than letting wave.open(str) do it): a bad
45+
# path then fails here cleanly, with no half-built Wave_write whose __del__ would
46+
# later raise an "ignored in __del__" warning during GC.
47+
handle = path.open("wb")
48+
except OSError as exc:
49+
raise CLIError(
50+
f"Cannot open {path} for writing: {exc}",
51+
error_type="save_audio_path",
52+
exit_code=2,
53+
) from exc
54+
try:
55+
# The Wave_write context manager closes (flushes + patches the length fields from
56+
# what was actually written) on exit, so the file is a valid WAV even when the
57+
# generator is closed mid-stream (Ctrl-C). The outer finally then closes the
58+
# handle we opened — after the patch — since wave only closes handles it opened.
59+
with wave.open(handle, "wb") as wav:
60+
wav.setnchannels(1)
61+
wav.setsampwidth(PCM16_SAMPLE_WIDTH_BYTES)
62+
wav.setframerate(rate)
63+
for chunk in audio:
64+
wav.writeframesraw(chunk)
65+
yield chunk
66+
finally:
67+
handle.close()

aai_cli/streaming/session.py

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@
1717
UsageError,
1818
mutually_exclusive,
1919
)
20+
from aai_cli.streaming import record
2021
from aai_cli.streaming.render import StreamRenderer, speaker_prefix
2122
from aai_cli.ui import output
2223
from aai_cli.ui.follow import FollowRenderer
@@ -137,6 +138,9 @@ class StreamSession:
137138
llm_prompts: list[str]
138139
model: str
139140
max_tokens: int
141+
# When set, tee the streamed PCM to this path as a WAV (see record.tee_wav). Only
142+
# the single-source path sets it — the parallel/batch callers reject --save-audio.
143+
save_audio: Path | None = None
140144
# Seconds between --llm summary refreshes; <=0 re-runs the chain on every turn.
141145
llm_interval: float = 0.0
142146
# Monotonic clock, injectable so the interval throttle is deterministic in tests.
@@ -242,6 +246,9 @@ def _maybe_summarize(self, *, final: bool = False) -> None:
242246
def stream_one(
243247
self, audio: Iterable[bytes], rate: int, *, source_label: str | None = None
244248
) -> None:
249+
if self.save_audio is not None:
250+
# Tee verbatim to disk at the source's true rate before it hits the wire.
251+
audio = record.tee_wav(audio, self.save_audio, rate=rate)
245252
flags = self.base_flags | {"sample_rate": rate}
246253
if source_label == "you":
247254
# The microphone captures you alone, so never diarize it into separate

tests/__snapshots__/test_snapshots_help_run.ambr

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -712,6 +712,9 @@
712712
│ --system-audio-only macOS only: stream │
713713
│ system/app audio without │
714714
│ the microphone │
715+
│ --save-audio FILE Tee the streamed PCM to │
716+
│ PATH as a 16-bit mono WAV │
717+
│ while transcribing │
715718
╰──────────────────────────────────────────────────────────────────────────────╯
716719
╭─ Model & Language ───────────────────────────────────────────────────────────╮
717720
│ --speech-model [universal-streaming-m Streaming speech model │
@@ -813,6 +816,8 @@
813816
$ assembly stream --sample
814817
Label speakers in the live transcript
815818
$ assembly stream --speaker-labels
819+
Save a WAV of the audio while streaming
820+
$ assembly stream --save-audio out.wav
816821
Boost domain terms with keyterm prompts
817822
$ assembly stream --keyterms-prompt "AssemblyAI" --keyterms-prompt "Claude"
818823
Summarize action items live as you talk

tests/test_stream_exec.py

Lines changed: 86 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -9,14 +9,16 @@
99
from __future__ import annotations
1010

1111
import dataclasses
12+
import wave
13+
from pathlib import Path
1214

1315
import pytest
1416

1517
from aai_cli.app.context import AppState
1618
from aai_cli.commands.stream import DEFAULT_SPEECH_MODEL
1719
from aai_cli.commands.stream import _exec as stream_exec
1820
from aai_cli.core import config, llm
19-
from aai_cli.core.errors import UsageError
21+
from aai_cli.core.errors import CLIError, UsageError
2022
from aai_cli.streaming.turn_presets import TurnDetectionPreset
2123

2224
# The CLI's flag defaults, as data. Tests override per-case with dataclasses.replace.
@@ -60,6 +62,7 @@
6062
config_file=None,
6163
output_field=None,
6264
show_code=False,
65+
save_audio=None,
6366
)
6467

6568

@@ -170,6 +173,7 @@ def test_stream_options_are_immutable():
170173
{"from_stdin": True, "device": 2}, # mic-only capture flags
171174
{"from_stdin": True, "sample_rate": 44100},
172175
{"from_stdin": True, "show_code": True}, # renders one source
176+
{"from_stdin": True, "save_audio": Path("out.wav")}, # tees one stream
173177
],
174178
)
175179
def test_from_stdin_rejects_incompatible_flags(overrides):
@@ -222,3 +226,84 @@ def fake_stream_batch(sources, *, make_session, open_source, renderer, json_mode
222226
dataclasses.replace(DEFAULTS, from_stdin=True), AppState(), json_mode=True
223227
)
224228
assert seen["sources"] == ["a.wav", "b.wav"]
229+
230+
231+
# --- --save-audio (tee the streamed PCM to a WAV) --------------------------
232+
class RecordingMic(FakeMic):
233+
"""A mic that yields known PCM so the tee'd WAV's contents can be asserted."""
234+
235+
PCM = b"\x01\x02\x03\x04\x05\x06\x07\x08"
236+
237+
def __iter__(self):
238+
return iter([self.PCM])
239+
240+
241+
def test_save_audio_tees_streamed_pcm_to_a_wav(monkeypatch, tmp_path):
242+
# The bytes the streaming API receives are also written to --save-audio, verbatim,
243+
# as a 16-bit mono WAV at the source's sample rate.
244+
config.set_api_key("default", "sk_live")
245+
out = tmp_path / "rec.wav"
246+
247+
def fake_stream_audio(api_key, source, *, params, **_kwargs):
248+
# Draining the iterable is what drives the tee — mirror the real SDK consuming it.
249+
sent = b"".join(source)
250+
assert sent == RecordingMic.PCM # the API still sees the unaltered audio
251+
252+
monkeypatch.setattr(stream_exec.client, "stream_audio", fake_stream_audio)
253+
monkeypatch.setattr(stream_exec, "MicrophoneSource", RecordingMic)
254+
255+
stream_exec.run_stream(
256+
dataclasses.replace(DEFAULTS, save_audio=out), AppState(), json_mode=True
257+
)
258+
259+
assert out.is_file()
260+
with wave.open(str(out), "rb") as w:
261+
assert w.getnchannels() == 1
262+
assert w.getsampwidth() == 2
263+
assert w.getframerate() == 16000 # FakeMic's reported rate
264+
assert w.readframes(w.getnframes()) == RecordingMic.PCM
265+
266+
267+
def test_save_audio_not_written_when_flag_unset(monkeypatch, tmp_path):
268+
# Without --save-audio, the default run leaves no stray WAV behind (kills a mutant
269+
# that tees unconditionally).
270+
config.set_api_key("default", "sk_live")
271+
monkeypatch.setattr(stream_exec.client, "stream_audio", lambda *a, **k: b"".join(a[1]))
272+
monkeypatch.setattr(stream_exec, "MicrophoneSource", RecordingMic)
273+
274+
stream_exec.run_stream(DEFAULTS, AppState(), json_mode=True)
275+
276+
assert list(tmp_path.glob("*.wav")) == []
277+
278+
279+
def test_save_audio_rejects_system_audio():
280+
# The mic + system streams can't share one file, so the combo is a usage error
281+
# (raised before credentials).
282+
with pytest.raises(UsageError):
283+
stream_exec.run_stream(
284+
dataclasses.replace(DEFAULTS, save_audio=Path("rec.wav"), system_audio=True),
285+
AppState(),
286+
json_mode=False,
287+
)
288+
289+
290+
def test_save_audio_rejects_show_code():
291+
# --show-code emits SDK code that doesn't tee audio, so the combo is rejected.
292+
with pytest.raises(UsageError):
293+
stream_exec.run_stream(
294+
dataclasses.replace(DEFAULTS, save_audio=Path("rec.wav"), show_code=True),
295+
AppState(),
296+
json_mode=False,
297+
)
298+
299+
300+
def test_save_audio_rejects_missing_parent_dir(tmp_path):
301+
# A path under a directory that doesn't exist is a clean path error, before auth.
302+
config.set_api_key("default", "sk_live")
303+
with pytest.raises(CLIError) as excinfo:
304+
stream_exec.run_stream(
305+
dataclasses.replace(DEFAULTS, save_audio=tmp_path / "nope" / "rec.wav"),
306+
AppState(),
307+
json_mode=False,
308+
)
309+
assert excinfo.value.error_type == "save_audio_path"

tests/test_streaming_record.py

Lines changed: 67 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,67 @@
1+
"""Unit tests for aai_cli.streaming.record — the --save-audio WAV tee."""
2+
3+
from __future__ import annotations
4+
5+
import wave
6+
7+
import pytest
8+
9+
from aai_cli.core.errors import CLIError
10+
from aai_cli.streaming import record
11+
12+
13+
def _read_wav(path):
14+
with wave.open(str(path), "rb") as w:
15+
return w.getnchannels(), w.getsampwidth(), w.getframerate(), w.readframes(w.getnframes())
16+
17+
18+
def test_tee_wav_yields_chunks_unchanged(tmp_path):
19+
chunks = [b"\x01\x02", b"\x03\x04\x05\x06"]
20+
out = list(record.tee_wav(iter(chunks), tmp_path / "a.wav", rate=16000))
21+
assert out == chunks # the tee must not alter what's streamed onward
22+
23+
24+
def test_tee_wav_writes_a_valid_wav_with_the_source_rate(tmp_path):
25+
path = tmp_path / "a.wav"
26+
list(record.tee_wav(iter([b"\x01\x02", b"\x03\x04"]), path, rate=44100))
27+
channels, width, rate, frames = _read_wav(path)
28+
assert channels == 1
29+
assert width == 2
30+
assert rate == 44100 # the declared source rate, not a hardcoded default
31+
assert frames == b"\x01\x02\x03\x04"
32+
33+
34+
def test_tee_wav_finalizes_a_valid_wav_on_early_close(tmp_path):
35+
# Ctrl-C closes the generator mid-stream; the partial file must still be valid WAV.
36+
path = tmp_path / "a.wav"
37+
gen = record.tee_wav(iter([b"\x01\x02", b"\x03\x04"]), path, rate=16000)
38+
assert next(gen) == b"\x01\x02" # consume only the first chunk
39+
gen.close() # raises GeneratorExit at the yield -> finally closes the WAV
40+
_channels, _width, _rate, frames = _read_wav(path)
41+
assert frames == b"\x01\x02" # only the consumed chunk landed
42+
43+
44+
def test_tee_wav_empty_stream_writes_a_zero_length_wav(tmp_path):
45+
path = tmp_path / "a.wav"
46+
assert list(record.tee_wav(iter([]), path, rate=16000)) == []
47+
_channels, _width, _rate, frames = _read_wav(path)
48+
assert frames == b""
49+
50+
51+
def test_tee_wav_unopenable_path_is_a_clean_error(tmp_path):
52+
# Pointing at a directory can't be opened for writing -> a CLIError, not a raw OSError.
53+
with pytest.raises(CLIError) as excinfo:
54+
# tee_wav opens lazily on first iteration, so the generator must be started.
55+
next(record.tee_wav(iter([b"\x01\x02"]), tmp_path, rate=16000))
56+
assert excinfo.value.error_type == "save_audio_path"
57+
58+
59+
def test_validate_target_accepts_an_existing_directory(tmp_path):
60+
record.validate_target(tmp_path / "rec.wav") # parent exists -> no raise
61+
62+
63+
def test_validate_target_rejects_a_missing_parent_directory(tmp_path):
64+
with pytest.raises(CLIError) as excinfo:
65+
record.validate_target(tmp_path / "nope" / "rec.wav")
66+
assert excinfo.value.error_type == "save_audio_path"
67+
assert excinfo.value.exit_code == 2

0 commit comments

Comments
 (0)