Add aai transcribe --download-sections (yt-dlp passthrough)

aai_cli/code_gen/__init__.py

@@ -33,9 +33,12 @@ def transcribe(
     source: str,
     *,
     llm_gateway: dict[str, object] | None = None,
+    download_sections: list[str] | None = None,
 ) -> str:
     """Generate runnable Python that reproduces this transcribe invocation."""
-    return _transcribe.render(merged, source, llm_gateway=llm_gateway)
+    return _transcribe.render(
+        merged, source, llm_gateway=llm_gateway, download_sections=download_sections
+    )
 
 
 def stream(

aai_cli/code_gen/transcribe.py

@@ -26,6 +26,7 @@ def render(
     *,
     llm_gateway: dict[str, object] | None = None,
     output: str | None = None,
+    download_sections: list[str] | None = None,
 ) -> str:
     """Generate a runnable transcribe script reproducing this CLI invocation.
 
@@ -37,32 +38,82 @@ def render(
     When `output` (a ``-o/--output`` field name) is given, the script prints that one
     field instead — and, as in the real command, it takes precedence over the LLM chain
     and the analysis sections.
+
+    When `download_sections` (yt-dlp ``--download-sections`` specs) is given for a
+    downloadable URL, the generated yt-dlp call fetches only those parts of the source.
     """
     if output is not None:
         llm_gateway = None  # `-o` returns before the chain runs in the real command
     needs_download = youtube.is_downloadable_url(source)
+    # Sections only apply to the download path; ignore them for a local file.
+    sections = list(download_sections) if (needs_download and download_sections) else []
+    ranges_expr, needs_re = _download_ranges(sections)
     parts = (
-        _header_block(llm_gateway, output, needs_download=needs_download)
-        + _transcribe_block(merged, source, needs_download=needs_download)
+        _header_block(
+            llm_gateway,
+            output,
+            needs_download=needs_download,
+            needs_re=needs_re,
+            has_sections=ranges_expr is not None,
+        )
+        + _transcribe_block(merged, source, needs_download=needs_download, ranges_expr=ranges_expr)
         + _result_block(merged, llm_gateway, output)
     )
     parts.append("")
     return "\n".join(parts)
 
 
+def _render_seconds(value: float) -> str:
+    """A Python literal for a section bound (``inf`` has no bare literal form)."""
+    if value == float("inf"):
+        return "float('inf')"
+    if value == float("-inf"):
+        return "float('-inf')"
+    return repr(value)
+
+
+def _download_ranges(sections: list[str]) -> tuple[str | None, bool]:
+    """Render a ``download_range_func(...)`` expression for `sections`.
+
+    Returns ``(expression_or_None, needs_re_import)`` — ``None`` when there are no
+    sections, and the flag is true when a chapter-regex spec means the generated
+    script needs ``import re``.
+    """
+    if not sections:
+        return None, False
+    chapters, ranges, from_url = youtube.parse_download_sections(sections)
+    chapters_src = "[" + ", ".join(f"re.compile({c!r})" for c in chapters) + "]"
+    ranges_src = (
+        "[" + ", ".join(f"({_render_seconds(s)}, {_render_seconds(e)})" for s, e in ranges) + "]"
+    )
+    return f"download_range_func({chapters_src}, {ranges_src}, {from_url})", bool(chapters)
+
+
 def _header_block(
-    llm_gateway: dict[str, object] | None, output: str | None, *, needs_download: bool
+    llm_gateway: dict[str, object] | None,
+    output: str | None,
+    *,
+    needs_download: bool,
+    needs_re: bool,
+    has_sections: bool,
 ) -> list[str]:
     """Imports plus the api-key (and non-default environment) settings lines."""
     stdlib_imports = ["import os"]
     if needs_download:
         # The download path fetches audio to a temp dir before uploading.
-        stdlib_imports += ["import tempfile"]
+        stdlib_imports.append("import tempfile")
+    if needs_re:
+        # Chapter-regex sections compile to re.compile(...) in the generated script.
+        stdlib_imports.append("import re")
     if output == "json":
-        stdlib_imports.insert(0, "import json")
+        stdlib_imports.append("import json")
+    stdlib_imports.sort()
     imports = ["import assemblyai as aai"]
     if needs_download:
         imports.append("import yt_dlp")
+    if has_sections:
+        # --download-sections builds yt-dlp's download_ranges via this helper.
+        imports.append("from yt_dlp.utils import download_range_func")
     if llm_gateway:
         imports.append("from openai import OpenAI")
     parts = [
@@ -81,7 +132,24 @@ def _header_block(
     return parts
 
 
-def _transcribe_block(merged: dict[str, object], source: str, *, needs_download: bool) -> list[str]:
+def _ydl_options_lines(ranges_expr: str | None) -> list[str]:
+    """The yt-dlp options dict for the download — one line, or multi-line when sections
+    add a ``download_ranges``/``force_keyframes_at_cuts`` pair."""
+    if ranges_expr is None:
+        return ['        {"format": "bestaudio/best", "outtmpl": f"{_tmp}/%(id)s.%(ext)s"}']
+    return [
+        "        {",
+        '            "format": "bestaudio/best",',
+        '            "outtmpl": f"{_tmp}/%(id)s.%(ext)s",',
+        f'            "download_ranges": {ranges_expr},',
+        '            "force_keyframes_at_cuts": True,',
+        "        }",
+    ]
+
+
+def _transcribe_block(
+    merged: dict[str, object], source: str, *, needs_download: bool, ranges_expr: str | None
+) -> list[str]:
     """The transcriber setup, optional config, the transcribe call, and error check."""
     parts = ["", "transcriber = aai.Transcriber()"]
     config_arg = ""
@@ -97,7 +165,7 @@ def _transcribe_block(merged: dict[str, object], source: str, *, needs_download:
             "# AssemblyAI can't fetch this page URL itself; download the audio first.",
             "with tempfile.TemporaryDirectory() as _tmp:",
             "    with yt_dlp.YoutubeDL(",
-            '        {"format": "bestaudio/best", "outtmpl": f"{_tmp}/%(id)s.%(ext)s"}',
+            *_ydl_options_lines(ranges_expr),
             "    ) as _ydl:",
             f"        _info = _ydl.extract_info({source!r}, download=True)",
             "        _audio = _ydl.prepare_filename(_info)",

aai_cli/commands/transcribe.py

@@ -29,35 +29,6 @@
 
 app = typer.Typer()
 
-# The PII policy strings the SDK accepts, validated client-side so a typo'd
-# --redact-pii-policy fails before any upload — mirroring how an unknown --config
-# key is rejected with the valid field list.
-_PII_POLICY_VALUES = frozenset(policy.value for policy in aai.PIIRedactionPolicy)
-
-
-def _validate_pii_policies(policies: list[str] | None) -> None:
-    unknown = [p for p in policies or [] if p not in _PII_POLICY_VALUES]
-    if unknown:
-        valid = ", ".join(sorted(_PII_POLICY_VALUES))
-        raise UsageError(f"Unknown PII policy(s) {unknown}. Valid policies: {valid}.")
-
-
-def _validate_language_flags(language_code: str | None, language_detection: bool | None) -> None:
-    if language_code and language_detection:
-        raise UsageError(
-            "--language-code and --language-detection can't be combined.",
-            suggestion="Force a language or auto-detect it, not both.",
-        )
-
-
-def _validate_speakers_expected(merged: dict[str, object]) -> None:
-    # Checked on the merged dict so `--config speaker_labels=true` also counts.
-    if merged.get("speakers_expected") and not merged.get("speaker_labels"):
-        raise UsageError(
-            "--speakers-expected only applies when diarization is enabled.",
-            suggestion="Add --speaker-labels.",
-        )
-
 
 @app.command(
     rich_help_panel=help_panels.TRANSCRIPTION,
@@ -278,6 +249,14 @@ def transcribe(
     audio_end: int | None = typer.Option(
         None, "--audio-end", help="End offset in ms.", rich_help_panel=help_panels.OPT_CUSTOMIZATION
     ),
+    download_sections: list[str] | None = typer.Option(
+        None,
+        "--download-sections",
+        help="For a YouTube/podcast URL, download only part of the source (yt-dlp "
+        '"--download-sections" syntax, e.g. "*0:00-5:00" for the first five minutes; '
+        "repeatable).",
+        rich_help_panel=help_panels.OPT_CUSTOMIZATION,
+    ),
     # webhooks
     webhook_url: str | None = typer.Option(
         None,
@@ -370,9 +349,11 @@ def transcribe(
     """
 
     def body(state: AppState, json_mode: bool) -> None:
-        _validate_language_flags(language_code, language_detection)
+        transcribe_exec.validate_language_flags(
+            language_code, language_detection=language_detection
+        )
         pii_policies = config_builder.split_csv(redact_pii_policy)
-        _validate_pii_policies(pii_policies)
+        transcribe_exec.validate_pii_policies(pii_policies)
         flags: dict[str, object] = {
             "speech_model": config_builder.enum_value(speech_model),
             "language_code": language_code,
@@ -428,7 +409,7 @@ def body(state: AppState, json_mode: bool) -> None:
             flags=flags, overrides=config_kv, config_file=config_file
         )
 
-        _validate_speakers_expected(merged)
+        transcribe_exec.validate_speakers_expected(merged)
 
         if show_code:
             # Print-only: build the equivalent script and exit without transcribing or
@@ -441,7 +422,13 @@ def body(state: AppState, json_mode: bool) -> None:
             )
             gateway = code_gen.gateway_options(list(llm_prompt or []), model, max_tokens)
             output.print_code(
-                render_transcribe_code(merged, audio, llm_gateway=gateway, output=output_field)
+                render_transcribe_code(
+                    merged,
+                    audio,
+                    llm_gateway=gateway,
+                    output=output_field,
+                    download_sections=list(download_sections or []),
+                )
             )
             return
 
@@ -453,7 +440,11 @@ def body(state: AppState, json_mode: bool) -> None:
         api_key = config.resolve_api_key(profile=state.profile)
         with output.status("Transcribing…", json_mode=json_mode, quiet=state.quiet):
             transcript = transcribe_exec.run_transcription(
-                api_key, source, sample=sample, transcription_config=tc
+                api_key,
+                source,
+                sample=sample,
+                transcription_config=tc,
+                download_sections=list(download_sections or []),
             )
 
         if out is not None:

aai_cli/transcribe_exec.py

@@ -17,6 +17,35 @@
 from aai_cli import choices, client, stdio, youtube
 from aai_cli.errors import UsageError
 
+# The PII policy strings the SDK accepts, validated client-side so a typo'd
+# --redact-pii-policy fails before any upload — mirroring how an unknown --config
+# key is rejected with the valid field list.
+PII_POLICY_VALUES = frozenset(policy.value for policy in aai.PIIRedactionPolicy)
+
+
+def validate_pii_policies(policies: list[str] | None) -> None:
+    unknown = [p for p in policies or [] if p not in PII_POLICY_VALUES]
+    if unknown:
+        valid = ", ".join(sorted(PII_POLICY_VALUES))
+        raise UsageError(f"Unknown PII policy(s) {unknown}. Valid policies: {valid}.")
+
+
+def validate_language_flags(language_code: str | None, *, language_detection: bool | None) -> None:
+    if language_code and language_detection:
+        raise UsageError(
+            "--language-code and --language-detection can't be combined.",
+            suggestion="Force a language or auto-detect it, not both.",
+        )
+
+
+def validate_speakers_expected(merged: dict[str, object]) -> None:
+    # Checked on the merged dict so `--config speaker_labels=true` also counts.
+    if merged.get("speakers_expected") and not merged.get("speaker_labels"):
+        raise UsageError(
+            "--speakers-expected only applies when diarization is enabled.",
+            suggestion="Add --speaker-labels.",
+        )
+
 
 def render_transform_steps(d: dict[str, Any]) -> str:
     """Human view of chained LLM-Gateway steps: the lone output, or each step labeled."""
@@ -56,6 +85,7 @@ def run_transcription(
     *,
     sample: bool,
     transcription_config: aai.TranscriptionConfig,
+    download_sections: list[str] | None = None,
 ) -> aai.Transcript:
     if source == "-":
         # Audio piped on stdin (e.g. `ffmpeg -i v.mp4 -f wav - | aai transcribe -`).
@@ -72,6 +102,6 @@ def run_transcription(
     if youtube.is_downloadable_url(audio):
         # Fetch first; AssemblyAI can't read a YouTube/podcast page URL itself.
         with tempfile.TemporaryDirectory(prefix="aai-yt-") as td:
-            local = youtube.download_audio(audio, Path(td))
+            local = youtube.download_audio(audio, Path(td), download_sections=download_sections)
             return client.transcribe(api_key, str(local), config=transcription_config)
     return client.transcribe(api_key, audio, config=transcription_config)