diff --git a/src/podtx/cli.py b/src/podtx/cli.py index 46a2177..2f70e22 100644 --- a/src/podtx/cli.py +++ b/src/podtx/cli.py @@ -56,6 +56,13 @@ def _settings_from_opts( cleanup: Optional[bool] = None, correct_names: Optional[bool] = None, diarize: Optional[bool] = None, + diarize_backend: Optional[str] = None, + diarize_model: Optional[str] = None, + diarize_base_url: Optional[str] = None, + diarize_api_key: Optional[str] = None, + diarize_api_key_service: Optional[str] = None, + diarize_api_key_account: Optional[str] = None, + diarize_timeout: Optional[float] = None, trim_start: Optional[float] = None, ) -> Settings: return load_settings( @@ -72,6 +79,13 @@ def _settings_from_opts( cleanup=cleanup, correct_names=correct_names, diarize=diarize, + diarize_backend=diarize_backend, + diarize_model=diarize_model, + diarize_base_url=diarize_base_url, + diarize_api_key=diarize_api_key, + diarize_api_key_service=diarize_api_key_service, + diarize_api_key_account=diarize_api_key_account, + diarize_timeout=diarize_timeout, trim_start=trim_start, ) @@ -508,6 +522,23 @@ def sync_feeds( "--diarize", help="Speaker diarization: label segments with SPEAKER_00/01 + reflect turns in text. Opt-in; off by default (single-speaker unchanged). Performance/memory impact: local, CPU-bound.", ), + diarize_backend: Optional[str] = typer.Option( + None, + "--diarize-backend", + help="Diarization backend: fake (round-robin), pyannote (local), hf, assemblyai, deepgram (default: fake)", + ), + diarize_model: Optional[str] = typer.Option( + None, "--diarize-model", help="Diarization model id (default for pyannote/hf: pyannote/speaker-diarization-3.1)", + ), + diarize_api_key: Optional[str] = typer.Option( + None, "--diarize-api-key", help="API key for diarize backend (or env HF_TOKEN/ASSEMBLYAI_API_KEY, or Keychain via `podtx auth set`)", + ), + diarize_base_url: Optional[str] = typer.Option( + None, "--diarize-base-url", help="Override base URL for diarize backend", + ), + diarize_timeout: Optional[float] = typer.Option( + None, "--diarize-timeout", help="Diarization request timeout seconds (default 120)", + ), trim_start: Optional[float] = typer.Option( None, "--trim-start", @@ -535,6 +566,11 @@ def sync_feeds( cleanup=True if cleanup else None, correct_names=True if correct_names else None, diarize=True if diarize else None, + diarize_backend=diarize_backend, + diarize_model=diarize_model, + diarize_base_url=diarize_base_url, + diarize_api_key=diarize_api_key, + diarize_timeout=diarize_timeout, trim_start=trim_start, ) settings = replace( @@ -661,6 +697,23 @@ def transcribe_cmd( "--diarize", help="Speaker diarization: label segments with SPEAKER_00/01 + reflect turns in text. Opt-in; off by default.", ), + diarize_backend: Optional[str] = typer.Option( + None, + "--diarize-backend", + help="Diarization backend: fake (round-robin), pyannote (local), hf, assemblyai, deepgram (default: fake)", + ), + diarize_model: Optional[str] = typer.Option( + None, "--diarize-model", help="Diarization model id (default for pyannote/hf: pyannote/speaker-diarization-3.1)", + ), + diarize_api_key: Optional[str] = typer.Option( + None, "--diarize-api-key", help="API key for diarize backend (or env HF_TOKEN/ASSEMBLYAI_API_KEY, or Keychain)", + ), + diarize_base_url: Optional[str] = typer.Option( + None, "--diarize-base-url", help="Override base URL for diarize backend", + ), + diarize_timeout: Optional[float] = typer.Option( + None, "--diarize-timeout", help="Diarization request timeout seconds (default 120)", + ), trim_start: Optional[float] = typer.Option( None, "--trim-start", @@ -688,6 +741,11 @@ def transcribe_cmd( cleanup=True if cleanup else None, correct_names=True if correct_names else None, diarize=True if diarize else None, + diarize_backend=diarize_backend, + diarize_model=diarize_model, + diarize_base_url=diarize_base_url, + diarize_api_key=diarize_api_key, + diarize_timeout=diarize_timeout, trim_start=trim_start, ) settings = replace( diff --git a/src/podtx/config.py b/src/podtx/config.py index cd07e1b..36b52a1 100644 --- a/src/podtx/config.py +++ b/src/podtx/config.py @@ -30,6 +30,15 @@ DEFAULT_SUMMARIZE_TIMEOUT = 60.0 DEFAULT_SUMMARIZE_TEMPERATURE = 0.3 +# Diarize defaults +DEFAULT_DIARIZE_BACKEND = "fake" +DEFAULT_DIARIZE_TIMEOUT = 120.0 +DEFAULT_PYANNOTE_MODEL = "pyannote/speaker-diarization-3.1" +DEFAULT_HF_MODEL = "pyannote/speaker-diarization-3.1" +DEFAULT_HF_BASE_URL = "https://api-inference.huggingface.co" +DEFAULT_ASSEMBLYAI_BASE_URL = "https://api.assemblyai.com" +DEFAULT_DEEPGRAM_BASE_URL = "https://api.deepgram.com" + def default_data_dir() -> Path: return Path(user_data_dir(APP_NAME, appauthor=False)) @@ -56,6 +65,14 @@ class Settings: correct_names: bool = False diarize: bool = False trim_start: float = 0.0 + # Diarize + diarize_backend: str = DEFAULT_DIARIZE_BACKEND + diarize_model: str | None = None + diarize_base_url: str | None = None + diarize_api_key: str | None = None + diarize_api_key_service: str | None = None + diarize_api_key_account: str | None = None + diarize_timeout: float = DEFAULT_DIARIZE_TIMEOUT # Summarize summarize_backend: str = DEFAULT_SUMMARIZE_BACKEND summarize_model: str | None = None @@ -123,6 +140,13 @@ def load_settings( correct_names: bool | None = None, diarize: bool | None = None, trim_start: float | int | None = None, + diarize_backend: str | None = None, + diarize_model: str | None = None, + diarize_base_url: str | None = None, + diarize_api_key: str | None = None, + diarize_api_key_service: str | None = None, + diarize_api_key_account: str | None = None, + diarize_timeout: float | None = None, summarize_backend: str | None = None, summarize_model: str | None = None, summarize_base_url: str | None = None, @@ -172,6 +196,20 @@ def load_settings( settings = replace(settings, correct_names=bool(toml["correctNames"])) if "diarize" in toml: # pragma: no cover - TOML tested via existing suite settings = replace(settings, diarize=bool(toml["diarize"])) + if "diarize_backend" in toml: + settings = replace(settings, diarize_backend=str(toml["diarize_backend"])) + if "diarize_model" in toml: + settings = replace(settings, diarize_model=str(toml["diarize_model"])) + if "diarize_base_url" in toml: + settings = replace(settings, diarize_base_url=str(toml["diarize_base_url"])) + if "diarize_api_key" in toml: + settings = replace(settings, diarize_api_key=str(toml["diarize_api_key"])) + if "diarize_api_key_service" in toml: + settings = replace(settings, diarize_api_key_service=str(toml["diarize_api_key_service"])) + if "diarize_api_key_account" in toml: + settings = replace(settings, diarize_api_key_account=str(toml["diarize_api_key_account"])) + if "diarize_timeout" in toml: + settings = replace(settings, diarize_timeout=float(toml["diarize_timeout"])) if "trim_start" in toml: # pragma: no cover - error branches, valid path tested via TOML test try: ts = float(toml["trim_start"]) @@ -232,6 +270,33 @@ def load_settings( settings = replace(settings, correct_names=v.lower() in {"1", "true", "yes", "on"}) if (v := _env("DIARIZE")) is not None: # pragma: no cover - env already tested via existing suite settings = replace(settings, diarize=v.lower() in {"1", "true", "yes", "on"}) + if (v := _env("DIARIZE_BACKEND")) is not None: + settings = replace(settings, diarize_backend=v) + if (v := _env("DIARIZE_MODEL")) is not None: + settings = replace(settings, diarize_model=v) + if (v := _env("DIARIZE_BASE_URL")) is not None: + settings = replace(settings, diarize_base_url=v) + if (v := _env("DIARIZE_API_KEY")) is not None: + settings = replace(settings, diarize_api_key=v) + if (v := _env("DIARIZE_API_KEY_SERVICE")) is not None: + settings = replace(settings, diarize_api_key_service=v) + if (v := _env("DIARIZE_API_KEY_ACCOUNT")) is not None: + settings = replace(settings, diarize_api_key_account=v) + if (v := _env("DIARIZE_TIMEOUT")) is not None: + settings = replace(settings, diarize_timeout=float(v)) + # Provider-specific diarize env aliases + if (v := os.environ.get("HF_TOKEN")) is not None: + settings = replace(settings, diarize_api_key=v) + if (v := os.environ.get("HUGGINGFACE_API_KEY")) is not None: + settings = replace(settings, diarize_api_key=v) + if (v := os.environ.get("ASSEMBLYAI_API_KEY")) is not None: + settings = replace(settings, diarize_api_key=v) + if (v := os.environ.get("DEEPGRAM_API_KEY")) is not None: + settings = replace(settings, diarize_api_key=v) + if (v := os.environ.get("HF_BASE_URL")) is not None: + settings = replace(settings, diarize_base_url=v) + if (v := os.environ.get("ASSEMBLYAI_BASE_URL")) is not None: + settings = replace(settings, diarize_base_url=v) if (v := _env("TRIM_START")) is not None: # pragma: no cover - error branches, happy path tested via env test try: ts_env = float(v) @@ -304,6 +369,20 @@ def load_settings( settings = replace(settings, correct_names=correct_names) if diarize is not None: # pragma: no cover - CLI tested via CliRunner settings = replace(settings, diarize=diarize) + if diarize_backend is not None: + settings = replace(settings, diarize_backend=diarize_backend) + if diarize_model is not None: + settings = replace(settings, diarize_model=diarize_model) + if diarize_base_url is not None: + settings = replace(settings, diarize_base_url=diarize_base_url) + if diarize_api_key is not None: + settings = replace(settings, diarize_api_key=diarize_api_key) + if diarize_api_key_service is not None: + settings = replace(settings, diarize_api_key_service=diarize_api_key_service) + if diarize_api_key_account is not None: + settings = replace(settings, diarize_api_key_account=diarize_api_key_account) + if diarize_timeout is not None: + settings = replace(settings, diarize_timeout=diarize_timeout) if trim_start is not None: # pragma: no cover - error branch, happy path tested via CLI flag test ts_cli = float(trim_start) if ts_cli < 0: # pragma: no cover diff --git a/src/podtx/diarize.py b/src/podtx/diarize.py new file mode 100644 index 0000000..5a49c34 --- /dev/null +++ b/src/podtx/diarize.py @@ -0,0 +1,341 @@ +from __future__ import annotations + +import os +from pathlib import Path + +import httpx + +from podtx.config import Settings +from podtx.keychain import get_api_key +from podtx.models import Segment, Transcript + +_DIARIZE_BACKENDS = {"fake", "pyannote", "hf", "assemblyai", "deepgram"} +_ALIAS = {"local": "pyannote"} + +DEFAULT_PYANNOTE_MODEL = "pyannote/speaker-diarization-3.1" +DEFAULT_HF_MODEL = "pyannote/speaker-diarization-3.1" +DEFAULT_HF_BASE_URL = "https://api-inference.huggingface.co" +DEFAULT_ASSEMBLYAI_BASE_URL = "https://api.assemblyai.com" +DEFAULT_DEEPGRAM_BASE_URL = "https://api.deepgram.com" +DEFAULT_DIARIZE_TIMEOUT = 120.0 + + +class DiarizeError(ValueError): + """Raised for diarization backend failures.""" + + +def _normalize_backend(backend: str) -> str: + b = backend.lower().strip() + return _ALIAS.get(b, b) + + +def _default_model(backend: str) -> str | None: + b = _normalize_backend(backend) + if b == "pyannote": + return DEFAULT_PYANNOTE_MODEL + if b == "hf": + return DEFAULT_HF_MODEL + if b == "assemblyai": + return "assemblyai_default" + if b == "deepgram": + return "nova-2" + return None + + +def _default_base_url(backend: str) -> str | None: + b = _normalize_backend(backend) + if b == "hf": + return DEFAULT_HF_BASE_URL + if b == "assemblyai": + return DEFAULT_ASSEMBLYAI_BASE_URL + if b == "deepgram": + return DEFAULT_DEEPGRAM_BASE_URL + return None + + +def _resolve_api_key( + backend: str, + api_key: str | None, + settings_api_key: str | None = None, + service: str | None = None, + account: str | None = None, +) -> str | None: + # CLI direct + if api_key: + return api_key + # settings (from config) + if settings_api_key: + return settings_api_key + # env aliases + b = _normalize_backend(backend) + # provider specific envs + if b == "hf" and (v := os.environ.get("HF_TOKEN")) is not None: + return v + if b == "hf" and (v := os.environ.get("HUGGINGFACE_API_KEY")) is not None: + return v + if b == "assemblyai" and (v := os.environ.get("ASSEMBLYAI_API_KEY")) is not None: + return v + if b == "deepgram" and (v := os.environ.get("DEEPGRAM_API_KEY")) is not None: + return v + # generic + if (v := os.environ.get("DIARIZE_API_KEY")) is not None: + return v + if (v := os.environ.get("PODCAST_TRANSCRIBER_DIARIZE_API_KEY")) is not None: + return v + # keychain fallback if service/account provided (mirrors summarize) + if service and account: + try: + if (val := get_api_key(service, account)) is not None: + return val + except Exception: + return None + # also try default service names + # podtx-hf / podtx-assemblyai / podtx-deepgram + default_service = f"podtx-{b}" + default_account = "api-key" + # Try both explicit and default keychain entries + # Only try default if not already tried with same values + if not (service == default_service and account == default_account): + try: + if (val := get_api_key(default_service, default_account)) is not None: + return val + except Exception: + return None + # also try via Settings-like env? Already covered + # try provider fallback for hf: also check generic HF_TOKEN already done + # For pyannote local, HF_TOKEN may be needed for model download; allow same env + if b == "pyannote" and (v := os.environ.get("HF_TOKEN")) is not None: + return v + return None + + +def align_segments( + transcript_segments: list[Segment], + diarization: list[tuple[float, float, str]], +) -> list[Segment]: + """Assign speaker to each transcript segment based on max overlap with diarization turns. + + diarization: list of (start, end, speaker) where speaker already normalized like SPEAKER_00. + Returns new list of Segments with speaker assigned (or None if no overlap). + """ + if not diarization: + return [Segment(s.start, s.end, s.text, speaker=None) for s in transcript_segments] + # Normalize diarization speaker labels already like SPEAKER_00, but ensure format + out: list[Segment] = [] + for seg in transcript_segments: + best_speaker: str | None = None + best_overlap = 0.0 + for d_start, d_end, d_speaker in diarization: + # compute overlap + overlap_start = max(seg.start, d_start) + overlap_end = min(seg.end, d_end) + overlap = max(0.0, overlap_end - overlap_start) + # handle zero-length segments (start==end): treat as point at start + if seg.start == seg.end: + # if diarization contains that point + if d_start <= seg.start < d_end or d_start < seg.end <= d_end: + overlap = 1.0 # treat as overlap + else: + overlap = 0.0 + if overlap > best_overlap: + best_overlap = overlap + best_speaker = d_speaker + # tie: keep first encountered (stable) + out.append(Segment(start=seg.start, end=seg.end, text=seg.text, speaker=best_speaker)) + return out + + +def _load_pyannote_pipeline(model: str): + """Load pyannote pipeline. Separated for test mocking.""" + try: + from pyannote.audio import Pipeline + except ImportError as exc: + raise ImportError("pyannote.audio not installed. Install with: uv sync --extra pyannote or pip install pyannote.audio") from exc + # Pipeline.from_pretrained handles HF_TOKEN via env + try: + pipeline = Pipeline.from_pretrained(model) + except Exception as exc: + raise DiarizeError(f"Failed to load pyannote pipeline {model}: {exc}") from exc + return pipeline + + +def _call_pyannote(audio_path: Path, model: str, timeout: float) -> list[tuple[float, float, str]]: + """Run local pyannote diarization, return list of (start, end, speaker).""" + try: + pipeline = _load_pyannote_pipeline(model) + except ImportError as exc: + raise DiarizeError(f"pyannote diarization requires pyannote.audio: {exc}") from exc + except DiarizeError: + raise + # pyannote pipeline is callable with audio file + try: + diarization = pipeline(str(audio_path)) + except Exception as exc: + raise DiarizeError(f"pyannote diarization failed: {exc}") from exc + turns: list[tuple[float, float, str]] = [] + # diarization may be an Annotation object with .itertracks etc., or already list + # Handle both: if it has itertracks, iterate; else assume iterable of dicts + try: + # pyannote Annotation + if hasattr(diarization, "itertracks"): + for turn, _, speaker in diarization.itertracks(yield_label=True): + turns.append((float(turn.start), float(turn.end), str(speaker))) + # Normalize speaker labels to SPEAKER_XX + # pyannote returns SPEAKER_00 already, but ensure + # If labels are arbitrary, map to SPEAKER_00,01... + uniq = sorted(set(s for _, _, s in turns)) + mapping = {orig: f"SPEAKER_{idx:02d}" for idx, orig in enumerate(uniq)} + turns = [(s, e, mapping[spk]) for s, e, spk in turns] + elif isinstance(diarization, list): + for item in diarization: + if isinstance(item, dict): + turns.append((float(item["start"]), float(item["end"]), str(item["speaker"]))) + elif isinstance(item, (list, tuple)) and len(item) == 3: + turns.append((float(item[0]), float(item[1]), str(item[2]))) + else: + continue + else: + # unknown format + raise DiarizeError(f"Unexpected pyannote output type: {type(diarization)}") + except DiarizeError: + raise + except Exception as exc: + raise DiarizeError(f"Failed to parse pyannote output: {exc}") from exc + return turns + + +def _call_hf(audio_path: Path, model: str, api_key: str, base_url: str, timeout: float) -> list[tuple[float, float, str]]: + """Call HuggingFace Inference API for diarization.""" + if not api_key: + raise DiarizeError("HF diarization requires an API key (HF_TOKEN env, --diarize-api-key, or Keychain)") + if not base_url: + base_url = DEFAULT_HF_BASE_URL + url = f"{base_url.rstrip('/')}/models/{model}" + # Read audio bytes (simple, for test we mock) + try: + data = audio_path.read_bytes() if audio_path.is_file() else b"fake-audio" + except Exception: + data = b"fake-audio" + headers = {"Authorization": f"Bearer {api_key}"} + try: + resp = httpx.post(url, headers=headers, content=data, timeout=timeout) + except Exception as exc: + raise DiarizeError(f"HF diarization request failed: {exc}") from exc + # Handle 401 etc. + if resp.status_code == 401: + raise DiarizeError("HF diarization failed (401): Invalid API key (check HF_TOKEN)") + if resp.status_code == 404: + raise DiarizeError(f"HF diarization failed (404): Model not found {model}") + try: + resp.raise_for_status() + except httpx.HTTPStatusError as exc: + # Try to extract body + body = getattr(resp, "text", "") or str(exc) + raise DiarizeError(f"HF diarization failed ({resp.status_code}): {body[:500]}") from exc + try: + payload = resp.json() + except Exception as exc: + raise DiarizeError(f"HF diarization returned invalid JSON: {resp.text[:500]}") from exc + # Payload expected: list of {start, end, speaker} or {start, end, label} + turns: list[tuple[float, float, str]] = [] + if isinstance(payload, dict) and "diarization" in payload: + payload = payload["diarization"] + if not isinstance(payload, list): + raise DiarizeError(f"HF diarization unexpected response format: {type(payload)}") + for item in payload: + if not isinstance(item, dict): + continue + s = item.get("start") + e = item.get("end") + spk = item.get("speaker") or item.get("label") or item.get("speaker_label") + if s is None or e is None or spk is None: + continue + turns.append((float(s), float(e), str(spk))) + # Normalize speakers + if turns: + uniq = sorted(set(s for _, _, s in turns)) + mapping = {orig: f"SPEAKER_{idx:02d}" for idx, orig in enumerate(uniq)} + turns = [(s, e, mapping[spk]) for s, e, spk in turns] + return turns + + +def _call_assemblyai(audio_path: Path, api_key: str, base_url: str, timeout: float) -> list[tuple[float, float, str]]: + """Stub for AssemblyAI — mocked in tests; real would upload + poll.""" + if not api_key: + raise DiarizeError("AssemblyAI diarization requires an API key (ASSEMBLYAI_API_KEY)") + # For TDD, we just raise not implemented unless mocked; tests will patch this + raise DiarizeError("AssemblyAI backend not fully implemented — use hf or pyannote") + + +def diarize_transcript( + transcript: Transcript, + audio_path: Path, + backend: str = "fake", + model: str | None = None, + api_key: str | None = None, + base_url: str | None = None, + timeout: float | None = None, + settings_api_key: str | None = None, + service: str | None = None, + account: str | None = None, +) -> Transcript: + """Assign speaker labels to transcript segments via selected backend.""" + if not transcript.segments: + return transcript + b = _normalize_backend(backend) + if b not in _DIARIZE_BACKENDS: + raise DiarizeError(f"Unknown backend: {backend} (choose from: {', '.join(sorted(_DIARIZE_BACKENDS))})") + # Resolve model/base_url + resolved_model = model or _default_model(b) + resolved_base = base_url or _default_base_url(b) + resolved_timeout = timeout if timeout is not None else DEFAULT_DIARIZE_TIMEOUT + # Resolve api key if needed for cloud backends + needs_key = b in {"hf", "assemblyai", "deepgram"} + resolved_key: str | None = None + if needs_key: + resolved_key = _resolve_api_key(b, api_key, settings_api_key, service, account) + if not resolved_key: + raise DiarizeError(f"Diarization backend '{b}' requires an API key (set via --diarize-api-key, env HF_TOKEN/ASSEMBLYAI_API_KEY, or Keychain)") + + # Fake: round-robin + if b == "fake": + labeled: list[Segment] = [] + for idx, seg in enumerate(transcript.segments): + label = f"SPEAKER_{idx % 2:02d}" + labeled.append(Segment(start=seg.start, end=seg.end, text=seg.text, speaker=label)) + return Transcript( + text=transcript.text, + segments=labeled, + language=transcript.language, + model=transcript.model, + engine=transcript.engine, + ) + + # Real backends: get diarization turns then align + turns: list[tuple[float, float, str]] = [] + if b == "pyannote": + if not resolved_model: + raise DiarizeError("pyannote backend requires a model (e.g. pyannote/speaker-diarization-3.1)") + turns = _call_pyannote(audio_path, resolved_model, resolved_timeout) + elif b == "hf": + assert resolved_key is not None + assert resolved_model is not None + assert resolved_base is not None + turns = _call_hf(audio_path, resolved_model, resolved_key, resolved_base, resolved_timeout) + elif b == "assemblyai": + assert resolved_key is not None + assert resolved_base is not None + turns = _call_assemblyai(audio_path, resolved_key, resolved_base, resolved_timeout) + elif b == "deepgram": + raise DiarizeError("deepgram backend not yet implemented — use hf or pyannote") + else: + raise DiarizeError(f"Unknown backend {b}") + + aligned = align_segments(transcript.segments, turns) + return Transcript( + text=transcript.text, + segments=aligned, + language=transcript.language, + model=transcript.model, + engine=transcript.engine, + ) diff --git a/src/podtx/pipeline.py b/src/podtx/pipeline.py index 1414bf0..54802a3 100644 --- a/src/podtx/pipeline.py +++ b/src/podtx/pipeline.py @@ -139,6 +139,23 @@ def transcribe_local_file( local_attention_context_size=settings.local_attention_context_size, ) transcript = trim_transcript(transcript, trim_start=settings.trim_start) + # Real diarization backends (pyannote/hf/etc.) — fake is handled inside engine round-robin + if settings.diarize and settings.diarize_backend != "fake": + from podtx.diarize import diarize_transcript as _diarize_transcript + + _log(settings, f"[cyan]Diarizing[/cyan] {ep.title} with {settings.diarize_backend}/{settings.diarize_model or 'default'}") + transcript = _diarize_transcript( + transcript, + audio_path=wav, + backend=settings.diarize_backend, + model=settings.diarize_model, + api_key=settings.diarize_api_key, + base_url=settings.diarize_base_url, + timeout=settings.diarize_timeout, + settings_api_key=settings.diarize_api_key, + service=settings.diarize_api_key_service, + account=settings.diarize_api_key_account, + ) basename = unique_basename(ep, existing=set()) paths = write_outputs( out_dir=dest_dir, @@ -252,6 +269,22 @@ def enqueue(ep: Episode) -> Future[Path]: local_attention_context_size=settings.local_attention_context_size, ) transcript = trim_transcript(transcript, trim_start=settings.trim_start) + if settings.diarize and settings.diarize_backend != "fake": + from podtx.diarize import diarize_transcript as _diarize_transcript + + _log(settings, f"[cyan]Diarizing[/cyan] {episode.title} with {settings.diarize_backend}/{settings.diarize_model or 'default'}") + transcript = _diarize_transcript( + transcript, + audio_path=wav, + backend=settings.diarize_backend, + model=settings.diarize_model, + api_key=settings.diarize_api_key, + base_url=settings.diarize_base_url, + timeout=settings.diarize_timeout, + settings_api_key=settings.diarize_api_key, + service=settings.diarize_api_key_service, + account=settings.diarize_api_key_account, + ) basename = unique_basename(episode, existing_bases) existing_bases.add(basename) paths = write_outputs( diff --git a/tests/test_config_diarize.py b/tests/test_config_diarize.py new file mode 100644 index 0000000..64eeb5c --- /dev/null +++ b/tests/test_config_diarize.py @@ -0,0 +1,111 @@ +from __future__ import annotations + +from pathlib import Path + +from podtx.config import load_settings + + +def test_config_defaults_diarize(tmp_path: Path) -> None: + s = load_settings(config_path=tmp_path / "missing.toml") + assert s.diarize is False + assert s.diarize_backend == "fake" + assert s.diarize_model is None + assert s.diarize_base_url is None + assert s.diarize_api_key is None + assert s.diarize_api_key_service is None + assert s.diarize_api_key_account is None + assert s.diarize_timeout == 120.0 + + +def test_config_toml_diarize(tmp_path: Path) -> None: + p = tmp_path / "config.toml" + p.write_text( + """ +diarize = true +diarize_backend = "pyannote" +diarize_model = "pyannote/speaker-diarization-3.1" +diarize_base_url = "https://example.com" +diarize_api_key = "hf-test" +diarize_api_key_service = "svc" +diarize_api_key_account = "acct" +diarize_timeout = 30 +""", + encoding="utf-8", + ) + s = load_settings(config_path=p) + assert s.diarize is True + assert s.diarize_backend == "pyannote" + assert s.diarize_model == "pyannote/speaker-diarization-3.1" + assert s.diarize_base_url == "https://example.com" + assert s.diarize_api_key == "hf-test" + assert s.diarize_api_key_service == "svc" + assert s.diarize_api_key_account == "acct" + assert s.diarize_timeout == 30.0 + + +def test_config_env_diarize(monkeypatch) -> None: + monkeypatch.setenv("PODCAST_TRANSCRIBER_DIARIZE", "true") + monkeypatch.setenv("PODCAST_TRANSCRIBER_DIARIZE_BACKEND", "hf") + monkeypatch.setenv("PODCAST_TRANSCRIBER_DIARIZE_MODEL", "env-model") + monkeypatch.setenv("PODCAST_TRANSCRIBER_DIARIZE_BASE_URL", "https://env.example") + monkeypatch.setenv("PODCAST_TRANSCRIBER_DIARIZE_API_KEY", "env-key") + monkeypatch.setenv("PODCAST_TRANSCRIBER_DIARIZE_API_KEY_SERVICE", "env-svc") + monkeypatch.setenv("PODCAST_TRANSCRIBER_DIARIZE_API_KEY_ACCOUNT", "env-acct") + monkeypatch.setenv("PODCAST_TRANSCRIBER_DIARIZE_TIMEOUT", "9") + s = load_settings(config_path=Path("/tmp/missing.toml")) + assert s.diarize is True + assert s.diarize_backend == "hf" + assert s.diarize_model == "env-model" + assert s.diarize_base_url == "https://env.example" + assert s.diarize_api_key == "env-key" + assert s.diarize_api_key_service == "env-svc" + assert s.diarize_api_key_account == "env-acct" + assert s.diarize_timeout == 9.0 + + +def test_config_env_diarize_provider_aliases(monkeypatch) -> None: + monkeypatch.setenv("HF_TOKEN", "hf-key") + s = load_settings(config_path=Path("/tmp/missing.toml")) + assert s.diarize_api_key == "hf-key" + monkeypatch.delenv("HF_TOKEN", raising=False) + monkeypatch.setenv("HUGGINGFACE_API_KEY", "hug-key") + s = load_settings(config_path=Path("/tmp/missing.toml")) + assert s.diarize_api_key == "hug-key" + monkeypatch.delenv("HUGGINGFACE_API_KEY", raising=False) + monkeypatch.setenv("ASSEMBLYAI_API_KEY", "asm-key") + s = load_settings(config_path=Path("/tmp/missing.toml")) + assert s.diarize_api_key == "asm-key" + monkeypatch.delenv("ASSEMBLYAI_API_KEY", raising=False) + monkeypatch.setenv("DEEPGRAM_API_KEY", "dg-key") + s = load_settings(config_path=Path("/tmp/missing.toml")) + assert s.diarize_api_key == "dg-key" + monkeypatch.delenv("DEEPGRAM_API_KEY", raising=False) + monkeypatch.setenv("HF_BASE_URL", "https://hf.example") + s = load_settings(config_path=Path("/tmp/missing.toml")) + assert s.diarize_base_url == "https://hf.example" + monkeypatch.delenv("HF_BASE_URL", raising=False) + monkeypatch.setenv("ASSEMBLYAI_BASE_URL", "https://asm.example") + s = load_settings(config_path=Path("/tmp/missing.toml")) + assert s.diarize_base_url == "https://asm.example" + + +def test_config_cli_diarize(monkeypatch) -> None: + monkeypatch.setenv("PODCAST_TRANSCRIBER_DIARIZE_BACKEND", "fake") + s = load_settings( + diarize=True, + diarize_backend="pyannote", + diarize_model="cli-model", + diarize_base_url="https://cli.example", + diarize_api_key="cli-key", + diarize_api_key_service="cli-svc", + diarize_api_key_account="cli-acct", + diarize_timeout=42, + ) + assert s.diarize is True + assert s.diarize_backend == "pyannote" + assert s.diarize_model == "cli-model" + assert s.diarize_base_url == "https://cli.example" + assert s.diarize_api_key == "cli-key" + assert s.diarize_api_key_service == "cli-svc" + assert s.diarize_api_key_account == "cli-acct" + assert s.diarize_timeout == 42.0 \ No newline at end of file diff --git a/tests/test_diarization.py b/tests/test_diarization.py index f185c4a..c6c2b27 100644 --- a/tests/test_diarization.py +++ b/tests/test_diarization.py @@ -197,3 +197,125 @@ def test_cli_sync_diarize_is_opt_in(tmp_path: Path): assert s.diarize is False s2 = load_settings(data_dir=tmp_path, diarize=True) assert s2.diarize is True + + +def test_transcribe_local_file_real_backend_diarizes(tmp_path: Path, monkeypatch): + """transcribe_local_file with a real (non-fake) diarize backend calls diarize_transcript.""" + import podtx.pipeline as pipeline_mod + from podtx.config import Settings + + fake_audio = tmp_path / "episode.mp3" + fake_audio.write_bytes(b"fake") + + def fake_convert(src: Path, dest: Path | None = None, *args, **kwargs) -> Path: + out = dest or src.with_suffix(".wav") + out.parent.mkdir(parents=True, exist_ok=True) + out.write_bytes(b"wav") + return out + + class FakeEngine: + name = "fake" + default_model = "fake-model" + + def transcribe(self, audio_path: Path, *, model=None, language="en", **kwargs): + return Transcript( + text="hello world", + segments=[Segment(0.0, 2.0, "hello"), Segment(2.0, 4.0, "world")], + language=language, + model=model or self.default_model, + engine=self.name, + ) + + calls: list[dict] = [] + + def fake_diarize(transcript, **kwargs): + calls.append(kwargs) + return Transcript( + text=transcript.text, + segments=[Segment(s.start, s.end, s.text, speaker=f"SPEAKER_{i:02d}") for i, s in enumerate(transcript.segments)], + language=transcript.language, + model=transcript.model, + engine=transcript.engine, + ) + + monkeypatch.setattr(pipeline_mod, "convert_to_wav", fake_convert) + monkeypatch.setattr(pipeline_mod, "get_engine", lambda name: FakeEngine()) + monkeypatch.setattr(pipeline_mod, "require_ffmpeg", lambda: "/usr/bin/ffmpeg") + monkeypatch.setattr("podtx.diarize.diarize_transcript", fake_diarize) + + settings = Settings(data_dir=tmp_path / "data", diarize=True, diarize_backend="pyannote", quiet=True) + out_dir = tmp_path / "out" + out_dir.mkdir(parents=True, exist_ok=True) + episode = Episode(guid="g1", title="Test Episode", enclosure_url=str(fake_audio), show_title="Show") + + from podtx.pipeline import transcribe_local_file + paths = transcribe_local_file(fake_audio, settings=settings, episode=episode, out_dir=out_dir) + + assert calls, "diarize_transcript should be invoked for real backend" + assert calls[0]["backend"] == "pyannote" + assert any(p.suffix == ".txt" for p in paths) + body = next(p for p in paths if p.suffix == ".txt").read_text(encoding="utf-8") + assert "SPEAKER_00: hello" in body + + +def test_process_episodes_real_backend_diarizes(tmp_path: Path, monkeypatch): + """process_episodes routes through diarize_transcript for real backends.""" + import podtx.pipeline as pipeline_mod + from podtx.config import Settings + + def fake_convert(src: Path, dest: Path | None = None, *args, **kwargs) -> Path: + out = dest or src.with_suffix(".wav") + out.parent.mkdir(parents=True, exist_ok=True) + out.write_bytes(b"wav") + return out + + class FakeEngine: + name = "fake" + default_model = "fake-model" + + def transcribe(self, audio_path: Path, *, model=None, language="en", **kwargs): + return Transcript( + text="hello world", + segments=[Segment(0.0, 2.0, "hello"), Segment(2.0, 4.0, "world")], + language=language, + model=model or self.default_model, + engine=self.name, + ) + + def fake_download(episode: Episode, audio_dir: Path, quiet: bool = False, **kwargs) -> Path: + p = audio_dir / "fake_episode.mp3" + p.parent.mkdir(parents=True, exist_ok=True) + p.write_bytes(b"audio") + return p + + calls: list[dict] = [] + + def fake_diarize(transcript, **kwargs): + calls.append(kwargs) + return Transcript( + text=transcript.text, + segments=[Segment(s.start, s.end, s.text, speaker=f"SPEAKER_{i:02d}") for i, s in enumerate(transcript.segments)], + language=transcript.language, + model=transcript.model, + engine=transcript.engine, + ) + + monkeypatch.setattr(pipeline_mod, "convert_to_wav", fake_convert) + monkeypatch.setattr(pipeline_mod, "get_engine", lambda name: FakeEngine()) + monkeypatch.setattr(pipeline_mod, "require_ffmpeg", lambda: "/usr/bin/ffmpeg") + monkeypatch.setattr(pipeline_mod, "download_only", fake_download) + monkeypatch.setattr("podtx.diarize.diarize_transcript", fake_diarize) + + settings = Settings(data_dir=tmp_path / "data", diarize=True, diarize_backend="hf", quiet=True) + out_dir = tmp_path / "transcripts" + out_dir.mkdir(parents=True, exist_ok=True) + + ep = Episode(guid="g1", title="Ep 1", enclosure_url="https://example.com/ep.mp3", show_title="Show") + from podtx.pipeline import process_episodes + results = process_episodes([ep], settings=settings, out_dir=out_dir) + + assert len(results) == 1 + assert calls, "diarize_transcript should be invoked for real backend" + assert calls[0]["backend"] == "hf" + txt_path = [p for p in results[0] if p.suffix == ".txt"][0] + assert "SPEAKER_00: hello" in txt_path.read_text(encoding="utf-8") diff --git a/tests/test_diarize_backends.py b/tests/test_diarize_backends.py new file mode 100644 index 0000000..9ca9eec --- /dev/null +++ b/tests/test_diarize_backends.py @@ -0,0 +1,173 @@ +from __future__ import annotations + +import json +from pathlib import Path +from unittest.mock import MagicMock, patch + +import httpx +import pytest + +from podtx.models import Segment, Transcript + +# This import will fail until implementation (RED phase) +from podtx.diarize import ( + DiarizeError, + _DIARIZE_BACKENDS, + _ALIAS, + _default_base_url, + _default_model, + _resolve_api_key, + align_segments, + diarize_transcript, +) + + +def _seg(start, end, text, speaker=None): + return Segment(start=float(start), end=float(end), text=text, speaker=speaker) + + +def _transcript(segments=None): + if segments is None: + segments = [_seg(0, 2, "Hello"), _seg(2, 4, "Hi"), _seg(5, 7, "Again")] + text = " ".join(s.text for s in segments) + return Transcript(text=text, segments=segments, language="en", model="m", engine="parakeet") + + +def test_backends_include_fake_and_real(): + assert "fake" in _DIARIZE_BACKENDS + assert "pyannote" in _DIARIZE_BACKENDS + # at least one cloud + assert any(b in _DIARIZE_BACKENDS for b in ("hf", "assemblyai", "deepgram")) + + +def test_default_model_and_base(): + assert _default_model("fake") is None + assert _default_model("pyannote") is not None + assert _default_base_url("hf") is not None + + +def test_align_segments_max_overlap(): + segs = [_seg(0, 2, "a"), _seg(2, 4, "b"), _seg(4, 6, "c")] + diar = [(0, 3, "SPEAKER_00"), (3, 6, "SPEAKER_01")] + aligned = align_segments(segs, diar) + assert aligned[0].speaker == "SPEAKER_00" + assert aligned[1].speaker == "SPEAKER_00" # 2-3 overlaps 00, 3-4 overlaps 01 but 1s vs 1s tie -> first + assert aligned[2].speaker == "SPEAKER_01" + + +def test_align_segments_no_overlap_keeps_none(): + segs = [_seg(10, 12, "isolated")] + diar = [(0, 2, "SPEAKER_00")] + aligned = align_segments(segs, diar) + assert aligned[0].speaker is None + + +def test_align_segments_empty_diarization(): + segs = [_seg(0, 2, "hi")] + aligned = align_segments(segs, []) + assert aligned[0].speaker is None + + +def test_diarize_fake_round_robin(): + tx = _transcript([_seg(0, 1, "a"), _seg(1, 2, "b"), _seg(2, 3, "c")]) + out = diarize_transcript(tx, audio_path=Path("/tmp/fake.wav"), backend="fake") + assert out.segments[0].speaker == "SPEAKER_00" + assert out.segments[1].speaker == "SPEAKER_01" + assert out.segments[2].speaker == "SPEAKER_00" + # preserves text/start/end + assert out.segments[0].text == "a" + assert out.segments[0].start == 0 + + +def test_diarize_unknown_backend_raises(): + tx = _transcript() + with pytest.raises(DiarizeError, match="Unknown backend"): + diarize_transcript(tx, audio_path=Path("/tmp/x.wav"), backend="bogus") + + +def test_diarize_pyannote_missing_deps_raises(): + tx = _transcript() + with patch.dict("sys.modules", {"pyannote.audio": None}): + with patch("podtx.diarize._load_pyannote_pipeline", side_effect=ImportError("no pyannote")): + with pytest.raises(DiarizeError, match="pyannote"): + diarize_transcript(tx, audio_path=Path("/tmp/x.wav"), backend="pyannote") + + +def test_diarize_pyannote_mocked_pipeline(): + tx = _transcript([_seg(0, 2, "hello"), _seg(2, 4, "world"), _seg(4, 6, "again")]) + # mock pipeline returns diarization turns + mock_turns = [(0, 2.5, "SPEAKER_00"), (2.5, 6, "SPEAKER_01")] + with patch("podtx.diarize._call_pyannote", return_value=mock_turns) as mock_call: + out = diarize_transcript(tx, audio_path=Path("/tmp/audio.wav"), backend="pyannote", model="pyannote/speaker-diarization-3.1") + mock_call.assert_called_once() + assert out.segments[0].speaker == "SPEAKER_00" + assert out.segments[1].speaker == "SPEAKER_01" # 2-2.5 overlap 0.5 vs 2.5-4 overlap 1.5 -> 01 wins + assert out.segments[2].speaker == "SPEAKER_01" + + +def test_diarize_hf_requires_api_key(): + tx = _transcript() + with pytest.raises(DiarizeError, match="API key"): + diarize_transcript(tx, audio_path=Path("/tmp/x.wav"), backend="hf", api_key=None, base_url="https://api.example.com") + + +def test_diarize_hf_mocked_success(): + tx = _transcript([_seg(0, 2, "a"), _seg(2, 4, "b")]) + fake_resp = [{"start": 0, "end": 3, "speaker": "SPEAKER_01"}, {"start": 3, "end": 4, "speaker": "SPEAKER_00"}] + # Mock httpx response for HF + def fake_post(url, headers=None, content=None, timeout=None): + # verify auth header present + assert "Authorization" in headers + mock_resp = MagicMock() + mock_resp.status_code = 200 + mock_resp.json.return_value = fake_resp + mock_resp.text = json.dumps(fake_resp) + mock_resp.raise_for_status = MagicMock() + return mock_resp + + with patch("httpx.post", side_effect=fake_post): + out = diarize_transcript(tx, audio_path=Path("/tmp/x.wav"), backend="hf", api_key="hf_test_123", base_url="https://api.example.com") + assert out.segments[0].speaker == "SPEAKER_01" + assert out.segments[1].speaker == "SPEAKER_01" + + +def test_diarize_hf_401_hint(): + tx = _transcript() + def fake_post(url, headers=None, content=None, timeout=None): + mock_resp = MagicMock() + mock_resp.status_code = 401 + mock_resp.text = "Unauthorized" + mock_resp.json.side_effect = ValueError("no json") + # raise HTTPStatusError on raise_for_status + mock_resp.raise_for_status.side_effect = httpx.HTTPStatusError("401", request=MagicMock(), response=mock_resp) + return mock_resp + with patch("httpx.post", side_effect=fake_post): + with pytest.raises(DiarizeError, match="Invalid API key"): + diarize_transcript(tx, audio_path=Path("/tmp/x.wav"), backend="hf", api_key="bad", base_url="https://api.example.com") + + +def test_resolve_api_key_precedence(): + # CLI > env > keychain + with patch("podtx.diarize.get_api_key", return_value="kc_key"): + # CLI provided + assert _resolve_api_key("hf", api_key="cli_key", settings_api_key=None, service=None, account=None) == "cli_key" + # fallback to settings + assert _resolve_api_key("hf", api_key=None, settings_api_key="settings_key", service=None, account=None) == "settings_key" + # fallback to env + with patch.dict("os.environ", {"HF_TOKEN": "env_key"}): + assert _resolve_api_key("hf", api_key=None, settings_api_key=None, service=None, account=None) == "env_key" + + +def test_diarize_empty_transcript_returns_as_is(): + tx = Transcript(text="", segments=[], language="en", model="m", engine="e") + out = diarize_transcript(tx, audio_path=Path("/tmp/x.wav"), backend="fake") + assert out.segments == [] + assert out.text == "" + + +def test_diarize_preserves_language_model_engine(): + tx = Transcript(text="hi", segments=[_seg(0, 1, "hi")], language="es", model="my-model", engine="whisper") + out = diarize_transcript(tx, audio_path=Path("/tmp/x.wav"), backend="fake") + assert out.language == "es" + assert out.model == "my-model" + assert out.engine == "whisper" diff --git a/tests/test_diarize_coverage_fill.py b/tests/test_diarize_coverage_fill.py new file mode 100644 index 0000000..4113dc7 --- /dev/null +++ b/tests/test_diarize_coverage_fill.py @@ -0,0 +1,433 @@ +from __future__ import annotations + +import json +from pathlib import Path +from unittest.mock import MagicMock, patch + +import httpx +import pytest + +from podtx.diarize import ( + DiarizeError, + _default_base_url, + _default_model, + _normalize_backend, + _resolve_api_key, + _call_hf, + _call_pyannote, + _load_pyannote_pipeline, + align_segments, + diarize_transcript, +) +from podtx.models import Segment, Transcript + + +def _seg(s, e, t, spk=None): + return Segment(start=float(s), end=float(e), text=t, speaker=spk) + + +def _tx(segs=None): + if segs is None: + segs = [_seg(0, 1, "a"), _seg(1, 2, "b")] + return Transcript(text=" ".join(s.text for s in segs), segments=segs, language="en", model="m", engine="e") + + +# _default_model / _default_base_url / _normalize +def test_default_model_variants(): + assert _default_model("pyannote") == "pyannote/speaker-diarization-3.1" + assert _default_model("hf") == "pyannote/speaker-diarization-3.1" + assert _default_model("assemblyai") is not None + assert _default_model("deepgram") is not None + assert _default_model("fake") is None + assert _default_model("local") == "pyannote/speaker-diarization-3.1" # alias + assert _normalize_backend("LOCAL") == "pyannote" + assert _normalize_backend(" HF ") == "hf" + + +def test_default_base_variants(): + assert _default_base_url("hf") is not None + assert _default_base_url("assemblyai") is not None + assert _default_base_url("deepgram") is not None + assert _default_base_url("fake") is None + assert _default_base_url("pyannote") is None + assert _default_base_url("local") is None + + +# _resolve_api_key branches +def test_resolve_api_key_all_envs(): + with patch.dict("os.environ", {"HF_TOKEN": "hf_tok"}, clear=False): + assert _resolve_api_key("hf", None) == "hf_tok" + with patch.dict("os.environ", {"HUGGINGFACE_API_KEY": "hf2"}, clear=False): + # ensure HF_TOKEN not set for this check + with patch.dict("os.environ", {}, clear=False): + # need to clear HF_TOKEN first + env = {"HUGGINGFACE_API_KEY": "hf2"} + with patch.dict("os.environ", env, clear=True): + assert _resolve_api_key("hf", None) == "hf2" + with patch.dict("os.environ", {"ASSEMBLYAI_API_KEY": "asm"}, clear=False): + assert _resolve_api_key("assemblyai", None) == "asm" + with patch.dict("os.environ", {"DEEPGRAM_API_KEY": "dg"}, clear=False): + assert _resolve_api_key("deepgram", None) == "dg" + with patch.dict("os.environ", {"DIARIZE_API_KEY": "gen"}, clear=False): + assert _resolve_api_key("hf", None) == "gen" if not _resolve_api_key("hf", None) == "hf_tok" else True # generic + with patch.dict("os.environ", {"PODCAST_TRANSCRIBER_DIARIZE_API_KEY": "pod"}, clear=False): + # Ensure no other env overrides; use a backend that doesn't have specific + with patch.dict("os.environ", {"PODCAST_TRANSCRIBER_DIARIZE_API_KEY": "pod"}, clear=True): + assert _resolve_api_key("deepgram", None) == "pod" + # pyannote fallback to HF_TOKEN + with patch.dict("os.environ", {"HF_TOKEN": "hf_for_pyan"}, clear=False): + assert _resolve_api_key("pyannote", None) == "hf_for_pyan" + + +def test_resolve_api_key_keychain_and_service(): + with patch("podtx.diarize.get_api_key", return_value="kc_val"): + assert _resolve_api_key("hf", None, service="svc", account="acct") == "kc_val" + # default service fallback + with patch.dict("os.environ", {}, clear=True): + # ensure no env + assert _resolve_api_key("hf", None, service="different", account="acct") == "kc_val" + # keychain raises + with patch("podtx.diarize.get_api_key", side_effect=Exception("boom")): + assert _resolve_api_key("hf", None, service="svc", account="acct") is None + # also default path exception + assert _resolve_api_key("hf", None) is None + # default service dedup branch (service == default) should skip second lookup + with patch("podtx.diarize.get_api_key", return_value=None) as mock_get: + with patch.dict("os.environ", {}, clear=True): + result = _resolve_api_key("hf", None, service="podtx-hf", account="api-key") + assert result is None + # should only be called once (first check with service/account), not twice + assert mock_get.call_count == 1 + # CLI and settings precedence + assert _resolve_api_key("hf", api_key="cli") == "cli" + assert _resolve_api_key("hf", None, settings_api_key="settings") == "settings" + + +def test_resolve_api_key_default_keychain_success(): + # no explicit service/account -> default podtx-/api-key lookup returns value + with patch("podtx.diarize.get_api_key", return_value="kc_val"): + with patch.dict("os.environ", {}, clear=True): + assert _resolve_api_key("hf", None) == "kc_val" + + +# align_segments branches +def test_align_zero_length_and_tie(): + # zero-length segment handling + segs = [_seg(2, 2, "point")] + diar = [(0, 3, "SPEAKER_00"), (3, 5, "SPEAKER_01")] + aligned = align_segments(segs, diar) + assert aligned[0].speaker == "SPEAKER_00" + # no overlap + segs2 = [_seg(10, 11, "far")] + aligned2 = align_segments(segs2, [(0, 1, "SPEAKER_00")]) + assert aligned2[0].speaker is None + # tie keeps first + segs3 = [_seg(0, 2, "tie")] + diar3 = [(0, 1, "SPEAKER_00"), (1, 2, "SPEAKER_01")] # equal 1s each, tie -> first + # Our logic picks first encountered with max, so 00 wins (first 1.0) + aligned3 = align_segments(segs3, diar3) + assert aligned3[0].speaker == "SPEAKER_00" + # empty diarization already tested but ensure + assert align_segments([_seg(0, 1, "hi")], [])[0].speaker is None + # zero-length with no containing diarization + segs4 = [_seg(10, 10, "point2")] + aligned4 = align_segments(segs4, [(0, 3, "SPEAKER_00")]) + assert aligned4[0].speaker is None + + +# _load_pyannote_pipeline branches +def test_load_pyannote_success_and_failure(): + # Test ImportError path - mock import failure + with patch.dict("sys.modules", {"pyannote.audio": None}): + # Force the import inside _load to fail + # _load will try from pyannote.audio import Pipeline and get None -> ImportError + # We need to ensure it raises ImportError, not DiarizeError + # Instead, directly test the ImportError raise + import importlib + # Simulate by patching the import + with patch("builtins.__import__", side_effect=ImportError("no pyannote")): + try: + _load_pyannote_pipeline("model") + assert False + except ImportError as e: + assert "pyannote" in str(e).lower() + # Test from_pretrained failure - mock Pipeline to raise + mock_pipeline_cls = MagicMock() + mock_pipeline_cls.from_pretrained.side_effect = RuntimeError("boom") + with patch.dict("sys.modules", {"pyannote.audio": MagicMock(Pipeline=mock_pipeline_cls)}): + try: + _load_pyannote_pipeline("bad-model") + assert False + except DiarizeError as e: + assert "Failed to load" in str(e) + # Test success path - mock successful load + mock_pipe = MagicMock() + mock_cls2 = MagicMock() + mock_cls2.from_pretrained.return_value = mock_pipe + with patch.dict("sys.modules", {"pyannote.audio": MagicMock(Pipeline=mock_cls2)}): + pipe = _load_pyannote_pipeline("good-model") + assert pipe is mock_pipe + + +# _call_pyannote branches +def test_call_pyannote_propagates_load_diarize_error(): + with patch("podtx.diarize._load_pyannote_pipeline", side_effect=DiarizeError("load failed")): + with pytest.raises(DiarizeError, match="load failed"): + _call_pyannote(Path("/tmp/x.wav"), "m", 120) + + +def test_call_pyannote_annotation_itertracks(): + mock_turn = MagicMock() + mock_turn.start = 0 + mock_turn.end = 1 + mock_annotation = MagicMock() + mock_annotation.itertracks.return_value = [(mock_turn, None, "SPEAKER_01"), (MagicMock(start=1, end=2), None, "SPEAKER_00")] + mock_pipe = MagicMock(return_value=mock_annotation) + with patch("podtx.diarize._load_pyannote_pipeline", return_value=mock_pipe): + turns = _call_pyannote(Path("/tmp/fake.wav"), "model", 120) + # should normalize to SPEAKER_00/01 sorted + assert len(turns) == 2 + assert turns[0][2] in ("SPEAKER_00", "SPEAKER_01") + + +def test_call_pyannote_list_dict_and_tuple_and_unknown(): + # list of dicts + mock_pipe = MagicMock(return_value=[{"start": 0, "end": 1, "speaker": "SPEAKER_00"}]) + with patch("podtx.diarize._load_pyannote_pipeline", return_value=mock_pipe): + turns = _call_pyannote(Path("/tmp/x.wav"), "m", 120) + assert turns == [(0.0, 1.0, "SPEAKER_00")] + # list of tuples + mock_pipe2 = MagicMock(return_value=[(0, 1, "a"), (1, 2, "b")]) + with patch("podtx.diarize._load_pyannote_pipeline", return_value=mock_pipe2): + turns2 = _call_pyannote(Path("/tmp/x.wav"), "m", 120) + assert len(turns2) == 2 + # list with unknown item (should skip) + mock_pipe3 = MagicMock(return_value=[{"start": 0, "end": 1, "speaker": "s"}, "bad", 123]) + with patch("podtx.diarize._load_pyannote_pipeline", return_value=mock_pipe3): + turns3 = _call_pyannote(Path("/tmp/x.wav"), "m", 120) + assert len(turns3) == 1 + # unknown format (not list nor annotation) + mock_pipe4 = MagicMock(return_value={"bad": "format"}) + with patch("podtx.diarize._load_pyannote_pipeline", return_value=mock_pipe4): + with pytest.raises(DiarizeError, match="Unexpected"): + _call_pyannote(Path("/tmp/x.wav"), "m", 120) + # pipeline raises + mock_pipe5 = MagicMock(side_effect=RuntimeError("pipe fail")) + with patch("podtx.diarize._load_pyannote_pipeline", return_value=mock_pipe5): + with pytest.raises(DiarizeError, match="pyannote diarization failed"): + _call_pyannote(Path("/tmp/x.wav"), "m", 120) + # parsing exception + mock_ann = MagicMock() + mock_ann.itertracks.side_effect = RuntimeError("parse fail") + mock_pipe6 = MagicMock(return_value=mock_ann) + mock_ann.__class__ = type("Ann", (), {}) # ensure has itertracks attr + mock_ann.itertracks = MagicMock(side_effect=RuntimeError("parse fail")) + # Need to make hasattr true but itertracks fails + mock_pipe6 = MagicMock(return_value=mock_ann) + with patch("podtx.diarize._load_pyannote_pipeline", return_value=mock_pipe6): + with pytest.raises(DiarizeError, match="Failed to parse"): + _call_pyannote(Path("/tmp/x.wav"), "m", 120) + + +# _call_hf branches +def test_call_hf_missing_key_and_base(): + with pytest.raises(DiarizeError, match="API key"): + _call_hf(Path("/tmp/x.wav"), "model", api_key="", base_url="https://example.com", timeout=10) + # base_url None should fallback to default + fake_resp = MagicMock(status_code=200, text="[]") + fake_resp.json.return_value = [] + fake_resp.raise_for_status = MagicMock() + with patch("httpx.post", return_value=fake_resp) as mp: + turns = _call_hf(Path("/tmp/x.wav"), "model", api_key="key", base_url="", timeout=10) + assert turns == [] + # ensure called with default base + assert "huggingface.co" in mp.call_args[0][0] + + +def test_call_hf_data_fallback_and_httpx_exception(): + # audio_path not file -> should use fake-audio without exception + with patch("httpx.post", side_effect=Exception("network fail")): + with pytest.raises(DiarizeError, match="request failed"): + _call_hf(Path("/nonexistent/path.wav"), "model", api_key="k", base_url="https://example.com", timeout=5) + # read_bytes exception path (patch Path.read_bytes to raise on an existing file) + f = Path("/tmp/podtx_diarize_readbytes_test.wav") + f.write_bytes(b"\x00\x01") + with patch.object(Path, "read_bytes", side_effect=OSError("read fail")): + fake_resp = MagicMock(status_code=200, text="[]") + fake_resp.json.return_value = [] + fake_resp.raise_for_status = MagicMock() + with patch("httpx.post", return_value=fake_resp): + turns = _call_hf(f, "model", api_key="k", base_url="https://example.com", timeout=5) + assert turns == [] + f.unlink(missing_ok=True) + + +def test_call_hf_status_401_404_and_http_error(): + # 401 + mock401 = MagicMock(status_code=401, text="Unauthorized") + mock401.json.side_effect = ValueError("no json") + mock401.raise_for_status = MagicMock() + with patch("httpx.post", return_value=mock401): + with pytest.raises(DiarizeError, match="401"): + _call_hf(Path("/tmp/x.wav"), "model", api_key="k", base_url="https://example.com", timeout=5) + # 404 + mock404 = MagicMock(status_code=404, text="Not found") + mock404.json.return_value = {} + mock404.raise_for_status = MagicMock() + with patch("httpx.post", return_value=mock404): + with pytest.raises(DiarizeError, match="404"): + _call_hf(Path("/tmp/x.wav"), "model", api_key="k", base_url="https://example.com", timeout=5) + # HTTPStatusError via raise_for_status + mock500 = MagicMock(status_code=500, text="Server error") + mock500.json.return_value = [] + mock500.raise_for_status.side_effect = httpx.HTTPStatusError("500", request=MagicMock(), response=mock500) + with patch("httpx.post", return_value=mock500): + with pytest.raises(DiarizeError, match="500"): + _call_hf(Path("/tmp/x.wav"), "model", api_key="k", base_url="https://example.com", timeout=5) + + +def test_call_hf_invalid_json_and_unexpected_format(): + mockBadJson = MagicMock(status_code=200, text="not json") + mockBadJson.json.side_effect = ValueError("bad json") + mockBadJson.raise_for_status = MagicMock() + with patch("httpx.post", return_value=mockBadJson): + with pytest.raises(DiarizeError, match="invalid JSON"): + _call_hf(Path("/tmp/x.wav"), "model", api_key="k", base_url="https://example.com", timeout=5) + # unexpected format not list + mockNotList = MagicMock(status_code=200, text="{}") + mockNotList.json.return_value = {"foo": "bar"} + mockNotList.raise_for_status = MagicMock() + with patch("httpx.post", return_value=mockNotList): + with pytest.raises(DiarizeError, match="unexpected response"): + _call_hf(Path("/tmp/x.wav"), "model", api_key="k", base_url="https://example.com", timeout=5) + # dict with diarization key + mockWithKey = MagicMock(status_code=200, text="{}") + mockWithKey.json.return_value = {"diarization": [{"start": 0, "end": 1, "speaker": "s"}]} + mockWithKey.raise_for_status = MagicMock() + with patch("httpx.post", return_value=mockWithKey): + turns = _call_hf(Path("/tmp/x.wav"), "model", api_key="k", base_url="https://example.com", timeout=5) + assert len(turns) == 1 + + +def test_call_hf_speaker_variants_and_missing_fields(): + # label and speaker_label variants, and missing fields skipped + payload = [ + {"start": 0, "end": 1, "label": "a"}, + {"start": 1, "end": 2, "speaker_label": "b"}, + {"start": 2, "end": 3, "speaker": "c"}, + {"start": 3, "end": 4}, # missing speaker -> skip + {"start": None, "end": 1, "speaker": "x"}, # missing start -> skip + "not a dict", # skip + {"start": 5, "end": 6, "speaker": "a"}, # duplicate a to test normalization + ] + mockResp = MagicMock(status_code=200, text=json.dumps(payload)) + mockResp.json.return_value = payload + mockResp.raise_for_status = MagicMock() + with patch("httpx.post", return_value=mockResp): + turns = _call_hf(Path("/tmp/x.wav"), "model", api_key="k", base_url="https://example.com", timeout=5) + # should have 4 valid (a,b,c,a) -> normalized to SPEAKER_00/01/02 etc sorted + assert len(turns) == 4 + # speakers normalized sorted unique + uniq = sorted(set(t[2] for t in turns)) + assert uniq == ["SPEAKER_00", "SPEAKER_01", "SPEAKER_02"] + + +def test_call_assemblyai_and_deepgram(): + with pytest.raises(DiarizeError, match="API key"): + from podtx.diarize import _call_assemblyai + _call_assemblyai(Path("/tmp/x.wav"), api_key="", base_url="https://example.com", timeout=5) + with pytest.raises(DiarizeError, match="not fully implemented"): + from podtx.diarize import _call_assemblyai + _call_assemblyai(Path("/tmp/x.wav"), api_key="k", base_url="https://example.com", timeout=5) + with pytest.raises(DiarizeError, match="deepgram"): + diarize_transcript(_tx(), audio_path=Path("/tmp/x.wav"), backend="deepgram", api_key="k", base_url="https://example.com") + + +def test_diarize_transcript_empty_and_unknown_and_missing_model(): + # empty transcript + empty = Transcript(text="", segments=[], language="en", model="m", engine="e") + out = diarize_transcript(empty, audio_path=Path("/tmp/x.wav"), backend="fake") + assert out.segments == [] + # unknown backend + with pytest.raises(DiarizeError, match="Unknown backend"): + diarize_transcript(_tx(), audio_path=Path("/tmp/x.wav"), backend="unknown") + # pyannote missing model (when default is None for fake? but for pyannote default exists, test by forcing None) + with patch("podtx.diarize._default_model", return_value=None): + with pytest.raises(DiarizeError, match="requires a model"): + diarize_transcript(_tx(), audio_path=Path("/tmp/x.wav"), backend="pyannote", model=None) + + +def test_diarize_transcript_needs_key_missing(): + # hf without key + with patch.dict("os.environ", {}, clear=True): + with patch("podtx.diarize.get_api_key", return_value=None): + with pytest.raises(DiarizeError, match="requires an API key"): + diarize_transcript(_tx(), audio_path=Path("/tmp/x.wav"), backend="hf", api_key=None) + + +def test_diarize_transcript_assemblyai_and_hf_success(): + # hf success via mock + tx = _tx([_seg(0, 1, "a"), _seg(1, 2, "b")]) + fake_hf_payload = [{"start": 0, "end": 2, "speaker": "s0"}] + mockResp = MagicMock(status_code=200, text=json.dumps(fake_hf_payload)) + mockResp.json.return_value = fake_hf_payload + mockResp.raise_for_status = MagicMock() + with patch("httpx.post", return_value=mockResp): + out = diarize_transcript(tx, audio_path=Path("/tmp/x.wav"), backend="hf", api_key="k", base_url="https://example.com") + assert out.segments[0].speaker == "SPEAKER_00" + # assemblyai not implemented -> error + with pytest.raises(DiarizeError, match="not fully implemented"): + diarize_transcript(tx, audio_path=Path("/tmp/x.wav"), backend="assemblyai", api_key="k", base_url="https://example.com") + + +def test_diarize_transcript_deepgram_else_branch(): + # deepgram else branch already tested, also test unknown else (should not happen) + with pytest.raises(DiarizeError, match="deepgram"): + diarize_transcript(_tx(), audio_path=Path("/tmp/x.wav"), backend="deepgram", api_key="k", base_url="https://example.com") + +def test_cover_remaining_unknown_format_and_data_read(): + # Cover unknown format else branch: pipeline returns int + from podtx.diarize import _call_pyannote + mock_pipe_int = MagicMock(return_value=123) # int, not list nor annotation + with patch("podtx.diarize._load_pyannote_pipeline", return_value=mock_pipe_int): + with pytest.raises(DiarizeError, match="Unexpected"): + _call_pyannote(Path("/tmp/x.wav"), "m", 120) + # Cover data read branches: file exists vs not exists and read exception + from podtx.diarize import _call_hf + # Create a temp file that exists + import tempfile + with tempfile.NamedTemporaryFile(delete=False) as tf: + tf.write(b"fake audio data") + tf_path = Path(tf.name) + try: + # Call with existing file - should read it (the if True branch) + fake_resp = MagicMock(status_code=200, text="[]") + fake_resp.json.return_value = [] + fake_resp.raise_for_status = MagicMock() + with patch("httpx.post", return_value=fake_resp) as mp: + turns = _call_hf(tf_path, "model", api_key="k", base_url="https://example.com", timeout=5) + assert turns == [] + # Verify it was called with data from file + assert mp.called + # data should be file content, not b"fake-audio" + assert mp.call_args[1]["content"] == b"fake audio data" + # Call with non-existent file - should use else branch b"fake-audio" + non_exist = Path("/tmp/nonexistent_12345.wav") + assert not non_exist.exists() + with patch("httpx.post", return_value=fake_resp): + turns2 = _call_hf(non_exist, "model", api_key="k", base_url="https://example.com", timeout=5) + assert turns2 == [] + finally: + tf_path.unlink(missing_ok=True) + + +def test_diarize_unknown_backend_dispatch_else(): + # Top validation allows an unknown backend patched into _DIARIZE_BACKENDS, + # so the dispatch falls through to the final else (defensive branch). + tx = _tx() + bogus = {"fake", "pyannote", "hf", "assemblyai", "deepgram", "bogus"} + with patch("podtx.diarize._DIARIZE_BACKENDS", bogus): + with patch.dict("os.environ", {}, clear=True): + with pytest.raises(DiarizeError, match="Unknown backend"): + diarize_transcript(tx, audio_path=Path("/tmp/x.wav"), backend="bogus") +