textflowkit 0.1.3__py3-none-any.whl

textflowkit/core/pipeline.py

@@ -0,0 +1,368 @@
+"""The shared pipeline.
+
+Every platform, every caller, every adapter runs this same path:
+
+    resolve source -> acquire media -> extract audio -> transcribe -> render
+
+Platform differences live entirely in the source layer.
+"""
+
+from __future__ import annotations
+
+import shutil
+import tempfile
+import time
+import uuid
+from collections.abc import Callable
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any
+
+from textflowkit.core.checkpoint import (
+    CheckpointRecord,
+    local_source_identity,
+    parse_checkpoint,
+    validate_local_resume,
+)
+from textflowkit.core.diarize import DiarizationError, assign_speakers, get_diarizer
+from textflowkit.core.engine import get_engine
+from textflowkit.core.model import Transcript
+from textflowkit.core.paths import (
+    UnsafeInputPathError,
+    default_input_root,
+    resolve_input_path,
+)
+from textflowkit.core.service import enforce_media_limits, enforce_predecode_limits
+from textflowkit.core.translate import (
+    TranslationError,
+    get_translator,
+    translate_segments,
+)
+from textflowkit.render import SUPPORTED_FORMATS, write_all
+from textflowkit.sources.acquire import (
+    AcquisitionError,
+    extract_audio,
+    fetch_media,
+    require_tool,
+    stage_confined_local_media,
+)
+from textflowkit.sources.detect import resolve_source
+
+
+class PipelineError(RuntimeError):
+    """Raised when any stage of the pipeline fails."""
+
+
+def _cleanup_scratch(path: Path) -> None:
+    """Remove a completed attempt, tolerating brief Windows decoder file locks.
+
+    A failed cleanup must not be silently ignored: the scratch tree may contain
+    downloaded media, so report a persistent failure to the caller.
+    """
+    for attempt in range(10):
+        try:
+            shutil.rmtree(path)
+            return
+        except FileNotFoundError:
+            return
+        except OSError as exc:
+            if attempt == 9:
+                raise PipelineError(f"scratch cleanup failed: {exc}") from exc
+            time.sleep(0.1)
+
+
+@dataclass(slots=True)
+class TranscribeResult:
+    transcript: Transcript
+    outputs: list[Path]
+
+
+def _existing_path(raw: str | None) -> Path | None:
+    if not raw:
+        return None
+    path = Path(raw)
+    return path if path.exists() else None
+
+
+def _resume_transcript(
+    resumed: CheckpointRecord | None,
+    *,
+    source: str,
+    ref,
+) -> Transcript | None:
+    if resumed is None or resumed.transcript is None:
+        return None
+    try:
+        transcript = Transcript.from_dict(resumed.transcript)
+    except (KeyError, TypeError, ValueError):
+        return None
+    transcript.source = source
+    transcript.platform = ref.platform
+    return transcript
+
+
+def _checkpoint_paths_usable(resumed: CheckpointRecord | None) -> bool:
+    if resumed is None or resumed.transcript is None:
+        return False
+    return "transcribe" in resumed.finished_stages
+
+
+def transcribe(
+    source: str,
+    *,
+    language: str | None = None,
+    formats: list[str] | None = None,
+    output_dir: str | Path | None = None,
+    model: str = "small",
+    engine: str = "whisper",
+    device: str | None = None,
+    cookies_from_browser: str | None = None,
+    keep_media: bool = False,
+    work_dir: str | Path | None = None,
+    check_cancel: Callable[[], None] | None = None,
+    input_root: str | Path | None = None,
+    diarize: bool = False,
+    diarizer_backend: str = "pyannote",
+    translate_to: str | None = None,
+    translator_backend: str = "ollama",
+    resume_checkpoint: dict[str, Any] | None = None,
+    on_checkpoint: Callable[[dict[str, Any]], None] | None = None,
+    output_id: str | None = None,
+) -> TranscribeResult:
+    """Run the full pipeline for a URL or local file.
+
+    `check_cancel` is an optional callback invoked at stage boundaries and
+    during yt-dlp progress and ffmpeg decode. It is expected to raise to abort
+    the run. Model inference and optional postprocessors still only stop at
+    their next stage boundary.
+
+    `resume_checkpoint` carries a validated snapshot from an earlier run.
+    Completed transcript work is reused; acquisition and engine work are skipped.
+    `on_checkpoint` receives an atomic snapshot after each completed stage.
+    """
+    resumed = parse_checkpoint(resume_checkpoint)
+    finished_stages = list(resumed.finished_stages) if resumed else []
+    media: Path | None = None
+    audio: Path | None = None
+    transcript: Transcript | None = None
+    local_identity: dict[str, Any] | None = resumed.local_identity if resumed else None
+
+    def _checkpoint(stage: str | None = None) -> CheckpointRecord | None:
+        if check_cancel is not None:
+            check_cancel()
+        if stage is None or on_checkpoint is None:
+            return None
+        if stage not in finished_stages:
+            finished_stages.append(stage)
+        snapshot = CheckpointRecord(
+            source=source,
+            model=model,
+            language=language,
+            engine=engine,
+            device=resumed.device if resumed else device,
+            options={
+                "formats": list(formats or []),
+                "diarize": diarize,
+                "diarizer_backend": diarizer_backend,
+                "translate_to": translate_to,
+                "translator_backend": translator_backend,
+            },
+            finished_stages=list(finished_stages),
+            transcript=transcript.to_dict() if isinstance(transcript, Transcript) else None,
+            media_path=_recordable(media),
+            audio_path=_recordable(audio),
+            local_identity=local_identity,
+        )
+        on_checkpoint(snapshot.to_dict())
+        return snapshot
+
+    def _recordable(path: Path | None) -> str | None:
+        """A path we can honestly promise to a later run.
+
+        The scratch directory is deleted at the end of a run, so recording a
+        path inside it produces a checkpoint that looks valid but can never be
+        used - resume finds the file gone and silently redoes the work. Only
+        paths outside scratch are worth storing.
+        """
+        if path is None:
+            return None
+        try:
+            resolved = Path(path).resolve()
+        except OSError:
+            return None
+        if scratch is not None and scratch.resolve() in resolved.parents:
+            return None
+        return str(path)
+
+    _checkpoint()
+
+    formats = formats or ["json", "srt", "txt"]
+    for fmt in formats:
+        if fmt.lower().lstrip(".") not in SUPPORTED_FORMATS:
+            raise PipelineError(
+                f"unsupported format '{fmt}'; choose from {', '.join(SUPPORTED_FORMATS)}"
+            )
+
+    # A local path may be confined; a URL is guarded separately by the SSRF
+    # check inside resolve_source. `input_root=None` means "use the configured
+    # root if one is set", which keeps the CLI unconfined by default.
+    root = input_root if input_root is not None else default_input_root()
+    resolved_source = source
+    if not source.startswith(("http://", "https://")):
+        try:
+            resolved_source = str(resolve_input_path(source, root=root))
+        except (FileNotFoundError, ValueError, UnsafeInputPathError) as exc:
+            raise PipelineError(str(exc)) from exc
+
+    try:
+        ref = resolve_source(resolved_source)
+    except (FileNotFoundError, ValueError) as exc:
+        raise PipelineError(str(exc)) from exc
+
+    _checkpoint("source")
+
+    if resumed is not None and ref.kind == "file":
+        try:
+            validate_local_resume(resumed, resolved_source, input_root=root)
+        except ValueError as exc:
+            raise PipelineError(str(exc)) from exc
+
+    if work_dir is not None:
+        Path(work_dir).mkdir(parents=True, exist_ok=True)
+    scratch = Path(tempfile.mkdtemp(prefix="textflowkit-", dir=work_dir))
+
+    try:
+        transcript = _resume_transcript(resumed, source=source, ref=ref)
+        if resumed is not None:
+            media = _existing_path(resumed.media_path)
+            audio = _existing_path(resumed.audio_path)
+
+        # Resuming means "the transcript already exists, do not transcribe again".
+        # The media files are a separate question: they live in scratch and are
+        # deleted after every run, so requiring them would make resume impossible.
+        # If a later stage (diarization) genuinely needs the audio and it is gone,
+        # re-acquire it - that is far cheaper than re-running Whisper.
+        can_resume = transcript is not None and _checkpoint_paths_usable(resumed)
+        if not can_resume:
+            transcript = None
+
+        try:
+            if not can_resume:
+                require_tool("ffmpeg")
+                if ref.kind == "file" and root is not None:
+                    media = stage_confined_local_media(
+                        ref.location, work_dir=scratch, input_root=root
+                    )
+                else:
+                    media = fetch_media(
+                        ref,
+                        work_dir=scratch,
+                        cookies_from_browser=cookies_from_browser,
+                        check_cancel=check_cancel,
+                    )
+                if ref.kind == "file":
+                    local_identity = local_source_identity(
+                        resolved_source, input_root=root, content_path=media,
+                    )
+                _checkpoint("fetch")
+                enforce_predecode_limits(media)
+                audio = extract_audio(media, work_dir=scratch, check_cancel=check_cancel)
+                enforce_media_limits(media, audio)
+                _checkpoint("extract")
+
+                eng = get_engine(engine, model=model, device=device)
+                try:
+                    transcript = eng.transcribe(audio, language=language)
+                except Exception as exc:  # engine failures are user-facing
+                    raise PipelineError(f"transcription failed: {exc}") from exc
+                if ref.kind == "file":
+                    try:
+                        current = local_source_identity(resolved_source, input_root=root)
+                    except (FileNotFoundError, OSError, ValueError) as exc:
+                        raise PipelineError("local source changed during transcription") from exc
+                    if current != local_identity:
+                        raise PipelineError("local source changed during transcription")
+                _checkpoint("transcribe")
+            elif diarize and audio is None:
+                # A finished transcript is the expensive checkpoint. Reacquire
+                # only the audio required by pyannote; never rerun Whisper.
+                require_tool("ffmpeg")
+                if media is None:
+                    if ref.kind == "file" and root is not None:
+                        media = stage_confined_local_media(
+                            ref.location, work_dir=scratch, input_root=root
+                        )
+                    else:
+                        media = fetch_media(
+                            ref, work_dir=scratch,
+                            cookies_from_browser=cookies_from_browser,
+                            check_cancel=check_cancel,
+                        )
+                    _checkpoint("fetch")
+                enforce_predecode_limits(media)
+                audio = extract_audio(media, work_dir=scratch, check_cancel=check_cancel)
+                enforce_media_limits(media, audio)
+                _checkpoint("extract")
+        except (AcquisitionError, UnsafeInputPathError) as exc:
+            raise PipelineError(str(exc)) from exc
+
+        if diarize:
+            # Refuse loudly rather than returning a transcript with empty speakers.
+            # A silent no-op here is exactly the defect that was removed from
+            # --speaker-labels, and it must not come back through this door.
+            if audio is None:
+                raise PipelineError("diarization requested but no audio is available for resume")
+            try:
+                diarizer = get_diarizer(diarizer_backend)
+                turns = diarizer.diarize(audio)
+            except DiarizationError as exc:
+                raise PipelineError(f"diarization requested but unavailable: {exc}") from exc
+            except Exception as exc:
+                raise PipelineError(f"diarization failed: {exc}") from exc
+            labelled = assign_speakers(transcript.segments, turns)
+            transcript.metadata["diarization"] = {
+                "backend": getattr(diarizer, "name", diarizer_backend),
+                "speakers": sorted({t.speaker for t in turns}),
+                "turns": len(turns),
+                "segments_labelled": labelled,
+            }
+
+        if translate_to:
+            # Refuse loudly: never present source text as though it were translated.
+            try:
+                translator = get_translator(translator_backend)
+                translated = translate_segments(
+                    transcript.segments, translate_to, translator=translator
+                )
+            except TranslationError as exc:
+                raise PipelineError(f"translation requested but unavailable: {exc}") from exc
+            except ValueError as exc:
+                raise PipelineError(f"translation failed: {exc}") from exc
+            transcript.metadata["translation"] = {
+                "backend": getattr(translator, "name", translator_backend),
+                "route": getattr(translator, "route", "unknown"),
+                "target": translate_to,
+                "segments_translated": translated,
+            }
+
+        _checkpoint("postprocess")
+
+        transcript.source = source
+        transcript.platform = ref.platform
+
+        outputs: list[Path] = []
+        if output_dir is not None:
+            stem = Path(ref.location).stem if ref.kind != "url" else "transcript"
+            stem = stem.replace("textflowkit-", "") or "transcript"
+            stem = f"{stem}-{output_id or uuid.uuid4().hex[:16]}"
+            outputs = write_all(transcript, formats=formats, output_dir=output_dir, stem=stem)
+        _checkpoint("render")
+
+        assert transcript is not None
+        result = TranscribeResult(transcript=transcript, outputs=outputs)
+    except BaseException:
+        _cleanup_scratch(scratch)
+        raise
+    if not keep_media:
+        _cleanup_scratch(scratch)
+    return result

textflowkit/core/retrieval.py

@@ -0,0 +1,148 @@
+"""Transcript retrieval: paging, time ranges, and search.
+
+A long transcript returned whole is a context problem. A 19-minute video is
+already ~51 KB of JSON (~214 segments); a 3-hour podcast is several hundred KB
+dropped into a model's context in a single tool result, most of it irrelevant to
+the question being asked.
+
+This module does the slicing in one place so the MCP and HTTP adapters cannot
+disagree about what `offset` or `start` mean - the same reasoning that put the
+pipeline in core rather than in each adapter.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+
+from textflowkit.core.model import Segment, Transcript
+
+
+@dataclass(slots=True)
+class Page:
+    """A window onto a transcript's segments."""
+
+    segments: list[Segment]
+    total: int
+    offset: int
+    returned: int
+    has_more: bool
+    start: float
+    end: float
+
+    def as_dict(self) -> dict:
+        return {
+            "total_segments": self.total,
+            "offset": self.offset,
+            "returned": self.returned,
+            "has_more": self.has_more,
+            "start": self.start,
+            "end": self.end,
+        }
+
+
+@dataclass(slots=True)
+class Match:
+    """A search hit, plus optional surrounding context."""
+
+    index: int
+    segment: Segment
+    context_before: list[Segment] = field(default_factory=list)
+    context_after: list[Segment] = field(default_factory=list)
+
+
+def page_segments(
+    transcript: Transcript,
+    *,
+    offset: int = 0,
+    limit: int | None = None,
+    start: float | None = None,
+    end: float | None = None,
+) -> Page:
+    """Slice a transcript by time range and/or page offset.
+
+    Filtering order is time first, then offset/limit within the filtered set, so
+    `offset` counts from the start of the requested window rather than from the
+    beginning of the transcript.
+
+    Raises `ValueError` for a negative offset or limit.
+    """
+    if offset < 0:
+        raise ValueError("offset must be >= 0")
+    if limit is not None and limit < 0:
+        raise ValueError("limit must be >= 0")
+
+    segments = [s for s in transcript.segments if not s.hidden]
+
+    if start is not None:
+        segments = [s for s in segments if s.end >= start]
+    if end is not None:
+        segments = [s for s in segments if s.start <= end]
+
+    total = len(segments)
+    window = segments[offset:] if offset else segments
+    if limit is not None:
+        window = window[:limit]
+
+    first = window[0].start if window else 0.0
+    last = window[-1].end if window else 0.0
+    consumed = offset + len(window)
+
+    return Page(
+        segments=window,
+        total=total,
+        offset=offset,
+        returned=len(window),
+        has_more=consumed < total,
+        start=first,
+        end=last,
+    )
+
+
+def search_segments(
+    transcript: Transcript,
+    query: str,
+    *,
+    limit: int = 20,
+    context: int = 0,
+    case_sensitive: bool = False,
+) -> list[Match]:
+    """Find segments whose text contains `query`.
+
+    `context` includes that many neighbouring segments either side, which is
+    usually what makes a hit readable. Matching is substring, not fuzzy - the
+    caller can ask again with a different phrase.
+
+    Raises `ValueError` for an empty query or negative limit/context.
+    """
+    if not query:
+        raise ValueError("query must not be empty")
+    if limit < 0:
+        raise ValueError("limit must be >= 0")
+    if context < 0:
+        raise ValueError("context must be >= 0")
+    if limit == 0:
+        return []
+
+    needle = query if case_sensitive else query.lower()
+    visible = [s for s in transcript.segments if not s.hidden]
+
+    matches: list[Match] = []
+    for idx, segment in enumerate(visible):
+        hay = segment.text if case_sensitive else segment.text.lower()
+        translated = segment.translated_text
+        if translated:
+            hay_translated = translated if case_sensitive else translated.lower()
+        else:
+            hay_translated = ""
+        if needle in hay or (hay_translated and needle in hay_translated):
+            matches.append(
+                Match(
+                    index=idx,
+                    segment=segment,
+                    context_before=visible[max(0, idx - context): idx] if context else [],
+                    context_after=visible[idx + 1: idx + 1 + context] if context else [],
+                )
+            )
+            if len(matches) >= limit:
+                break
+    return matches

textflowkit/core/runner.py

@@ -0,0 +1,146 @@
+"""Job execution.
+
+`run_job` executes one job synchronously and owns its state transitions.
+`submit` is the entry point callers use; it delegates to the process-wide
+`JobExecutor`, so jobs get bounded concurrency and can be cancelled.
+
+Cancellation contract:
+
+- `JobCancelled` is an orderly stop, not a failure. The job ends CANCELLED.
+- A job that finishes while a cancellation is in flight must not overwrite the
+  cancelled state with DONE, so the final transition re-reads the job first.
+- Explicit resume of an interrupted job is the one path allowed past the
+  terminal-state guard. Callers use `prepare_resume` first, which resets the job
+  to PENDING while preserving its checkpoint.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Callable
+from pathlib import Path
+from typing import Any
+
+from textflowkit.core.checkpoint import write_checkpoint
+from textflowkit.core.executor import JobCancelled, get_default_executor
+from textflowkit.core.jobs import Job, JobState, JobStore
+from textflowkit.core.model import Transcript
+from textflowkit.core.pipeline import PipelineError, transcribe
+
+
+def run_job(
+    job: Job,
+    store: JobStore,
+    *,
+    source: str,
+    language: str | None = None,
+    formats: list[str] | None = None,
+    output_dir: str | Path | None = None,
+    model: str = "small",
+    engine: str = "whisper",
+    device: str | None = None,
+    cookies_from_browser: str | None = None,
+    keep_media: bool = False,
+    work_dir: str | Path | None = None,
+    check_cancel: Callable[[], None] | None = None,
+    input_root: str | Path | None = None,
+    diarize: bool = False,
+    diarizer_backend: str = "pyannote",
+    translate_to: str | None = None,
+    translator_backend: str = "ollama",
+    resume_checkpoint: dict[str, Any] | None = None,
+) -> None:
+    """Execute a job, recording its terminal state. Callers decide the thread."""
+    # Do not start work that has already been cancelled or otherwise finished.
+    # `submit` only hands us fresh jobs; explicit resume prepares the row first.
+    current = store.get(job.id)
+    if current is not None and current.is_terminal:
+        return
+
+    store.update(job.id, state=JobState.RUNNING, progress="starting")
+
+    try:
+        result = transcribe(
+            source,
+            language=language,
+            formats=formats,
+            output_dir=output_dir,
+            model=model,
+            engine=engine,
+            device=device,
+            cookies_from_browser=cookies_from_browser,
+            keep_media=keep_media,
+            work_dir=work_dir,
+            check_cancel=check_cancel,
+            input_root=input_root,
+            diarize=diarize,
+            diarizer_backend=diarizer_backend,
+            translate_to=translate_to,
+            translator_backend=translator_backend,
+            resume_checkpoint=resume_checkpoint,
+            on_checkpoint=lambda record: write_checkpoint(store, job.id, record),
+            output_id=job.id,
+        )
+    except JobCancelled:
+        store.update(job.id, state=JobState.CANCELLED, progress="cancelled")
+        return
+    except PipelineError as exc:
+        store.update(job.id, state=JobState.ERROR, error=str(exc), progress="failed")
+        return
+    # Last-resort guard: a job must never be left stuck in RUNNING because of an
+    # unexpected exception type. The error is recorded on the job, not swallowed.
+    # Narrower catches above handle the expected failure modes.
+    except Exception as exc:  # noqa: BLE001
+        store.update(
+            job.id,
+            state=JobState.ERROR,
+            error=f"{type(exc).__name__}: {exc}",
+            progress="failed",
+        )
+        return
+
+    # A cancellation that arrived while the last stage ran must win over DONE.
+    latest = store.get(job.id)
+    if latest is not None and latest.state is JobState.CANCELLED:
+        return
+
+    store.update(
+        job.id,
+        state=JobState.DONE,
+        progress="complete",
+        transcript=result.transcript.to_dict(),
+        outputs=[str(p) for p in result.outputs],
+    )
+
+
+def submit(
+    store: JobStore,
+    *,
+    source: str,
+    background: bool = True,
+    **kwargs,
+) -> Job:
+    """Create a job and schedule it.
+
+    `background=False` runs inline (used by tests and by callers that want a
+    blocking call). Otherwise the job goes to the process-wide executor, which
+    bounds concurrency and can cancel it.
+    """
+    if not background:
+        job = store.create(source)
+        run_job(job, store, source=source, **kwargs)
+        return job
+
+    executor = get_default_executor()
+    if executor.store is not store:
+        # A caller passed a specific store; run it directly rather than silently
+        # routing the job to a different store than the one they hold.
+        job = store.create(source)
+        run_job(job, store, source=source, **kwargs)
+        return job
+    return executor.submit(source=source, **kwargs)
+
+
+def transcript_for(job: Job) -> Transcript | None:
+    if job.transcript is None:
+        return None
+    return Transcript.from_dict(job.transcript)