Cognitive-rag / 01_data_ingestion / ingest_multimodal.py
ingest_multimodal.py
Raw
#!/usr/bin/env python
"""Ingest multimodal knowledge base files into structured JSONL."""

from __future__ import annotations

import argparse
import csv
import hashlib
import json
import os
import re
import shutil
import subprocess
import tempfile
from collections import Counter
from datetime import datetime, timezone
from pathlib import Path
from typing import Dict, Iterable, List, Optional, Tuple


TEXT_EXTS = {".pdf", ".docx", ".pptx", ".txt", ".md", ".html", ".htm"}
IMAGE_EXTS = {".jpg", ".jpeg", ".png", ".tif", ".tiff", ".bmp"}
AUDIO_EXTS = {".mp3", ".wav", ".m4a", ".flac", ".aac", ".ogg"}
VIDEO_EXTS = {".mp4", ".mov", ".mkv", ".avi", ".webm"}


def normalize_text(text: str) -> str:
    text = text.replace("\u0000", "").replace("\ufffd", "")
    text = text.replace("\r\n", "\n").replace("\r", "\n")
    lines = [re.sub(r"[ \t]+", " ", line).strip() for line in text.split("\n")]
    cleaned = []
    blank_run = 0
    for line in lines:
        if not line:
            blank_run += 1
            if blank_run <= 1:
                cleaned.append("")
        else:
            blank_run = 0
            cleaned.append(line)
    return "\n".join(cleaned).strip()


def slugify(value: str) -> str:
    value = value.strip().lower()
    value = re.sub(r"[^a-z0-9]+", "-", value)
    value = value.strip("-")
    return value or "doc"


def build_doc_id(path: Path, root: Path) -> str:
    try:
        rel_path = path.relative_to(root).as_posix()
    except ValueError:
        rel_path = path.as_posix()
    slug = slugify(rel_path)
    digest = hashlib.sha1(rel_path.encode("utf-8")).hexdigest()[:8]
    return f"{slug}__{digest}"


def format_timestamp(seconds: float) -> str:
    millis = int(round(seconds * 1000))
    hours = millis // 3_600_000
    minutes = (millis % 3_600_000) // 60_000
    secs = (millis % 60_000) // 1000
    ms = millis % 1000
    return f"{hours:02d}:{minutes:02d}:{secs:02d}.{ms:03d}"


def safe_int(value: object) -> Optional[int]:
    if value is None:
        return None
    try:
        return int(value)
    except (TypeError, ValueError):
        return None


def detect_file_type(path: Path) -> Optional[str]:
    ext = path.suffix.lower()
    if ext in TEXT_EXTS:
        return ext.lstrip(".")
    if ext in IMAGE_EXTS:
        return "image"
    if ext in AUDIO_EXTS:
        return "audio"
    if ext in VIDEO_EXTS:
        return "video"
    if ext == ".csv":
        return "csv"
    return None


def iter_files(root: Path, max_files: Optional[int] = None) -> Iterable[Path]:
    count = 0
    for path in sorted(root.rglob("*")):
        if path.is_dir():
            continue
        if path.name.startswith("."):
            continue
        if "output" in path.parts:
            continue
        if max_files is not None and count >= max_files:
            break
        count += 1
        yield path


def load_langchain_documents(path: Path, filetype: str, pdf_strategy: str) -> Tuple[List[object], str]:
    try:
        from langchain_community.document_loaders import (
            TextLoader,
            UnstructuredPDFLoader,
            UnstructuredWordDocumentLoader,
            UnstructuredPowerPointLoader,
            UnstructuredHTMLLoader,
            UnstructuredMarkdownLoader,
        )
    except Exception as exc:  # pragma: no cover - import guard
        raise RuntimeError(
            "Missing langchain-community loaders. "
            "Install with `pip install -r 01_data_ingestion/requirements.txt`."
        ) from exc

    loader = None
    loader_name = ""

    if filetype == "pdf":
        loader = UnstructuredPDFLoader(str(path), mode="elements", strategy=pdf_strategy)
        loader_name = "UnstructuredPDFLoader"
    elif filetype == "docx":
        loader = UnstructuredWordDocumentLoader(str(path), mode="elements")
        loader_name = "UnstructuredWordDocumentLoader"
    elif filetype == "pptx":
        loader = UnstructuredPowerPointLoader(str(path), mode="elements")
        loader_name = "UnstructuredPowerPointLoader"
    elif filetype in {"html", "htm"}:
        loader = UnstructuredHTMLLoader(str(path), mode="elements")
        loader_name = "UnstructuredHTMLLoader"
    elif filetype == "md":
        loader = UnstructuredMarkdownLoader(str(path), mode="elements")
        loader_name = "UnstructuredMarkdownLoader"
    elif filetype == "txt":
        loader = TextLoader(str(path), encoding="utf-8")
        loader_name = "TextLoader"
    else:
        raise ValueError(f"Unsupported text filetype: {filetype}")

    return loader.load(), loader_name


def extract_page(metadata: Dict[str, object]) -> Tuple[Optional[int], Optional[str]]:
    for key in ("page_number", "page", "page_num", "page_index", "page_no"):
        if key in metadata:
            page_val = safe_int(metadata.get(key))
            if page_val is not None:
                return page_val, "page"
    for key in ("slide_number", "slide", "slide_index"):
        if key in metadata:
            slide_val = safe_int(metadata.get(key))
            if slide_val is not None:
                return slide_val, "slide"
    return None, None


def build_source_base(
    path: Path,
    root: Path,
    filetype: str,
    modality: str,
    loader_name: str,
) -> Dict[str, object]:
    try:
        rel_path = path.relative_to(root).as_posix()
    except ValueError:
        rel_path = path.as_posix()
    return {
        "path": rel_path,
        "filename": path.name,
        "filetype": filetype,
        "modality": modality,
        "loader": loader_name,
    }


def add_relations(records: List[Dict[str, object]]) -> None:
    total = len(records)
    for idx, record in enumerate(records):
        record_id = f"{record['doc_id']}:{idx:06d}"
        source = record.setdefault("source", {})
        source["record_id"] = record_id
        source["prev_record_id"] = f"{record['doc_id']}:{idx - 1:06d}" if idx > 0 else None
        source["next_record_id"] = (
            f"{record['doc_id']}:{idx + 1:06d}" if idx + 1 < total else None
        )


def records_from_documents(
    docs: List[object],
    doc_id: str,
    path: Path,
    root: Path,
    filetype: str,
    modality: str,
    loader_name: str,
) -> List[Dict[str, object]]:
    records: List[Dict[str, object]] = []
    for doc in docs:
        content = normalize_text(getattr(doc, "page_content", "") or "")
        if not content:
            continue
        metadata = getattr(doc, "metadata", {}) or {}
        page_num, page_type = extract_page(metadata)
        element_type = metadata.get("category") or metadata.get("element_type")
        element_id = metadata.get("element_id") or metadata.get("id")
        section_label = None
        if page_num is not None:
            prefix = "slide" if page_type == "slide" else "page"
            section_label = f"{prefix}_{page_num}"
        else:
            section_label = f"element_{len(records)}"

        source = build_source_base(path, root, filetype, modality, loader_name)
        if page_num is not None:
            source["page_number"] = page_num
            source["page_type"] = page_type
        if element_type:
            source["element_type"] = element_type
        if element_id:
            source["element_id"] = element_id

        records.append(
            {
                "doc_id": doc_id,
                "page": page_num if page_type in {"page", "slide"} else None,
                "section": section_label,
                "section_index": len(records),
                "content": content,
                "source": source,
            }
        )
    return records


class OcrRunner:
    def __init__(self, engine: str, language: str):
        self.engine = engine
        self.language = language
        self._paddle = None

    def extract_text(self, path: Path) -> str:
        if self.engine == "tesseract":
            try:
                from PIL import Image
                import pytesseract
            except Exception as exc:  # pragma: no cover - import guard
                raise RuntimeError(
                    "Missing OCR deps. Install with `pip install -r 01_data_ingestion/requirements.txt` "
                    "and ensure tesseract is installed."
                ) from exc
            image = Image.open(path)
            return pytesseract.image_to_string(image, lang=self.language)
        if self.engine == "paddle":
            try:
                from paddleocr import PaddleOCR
            except Exception as exc:  # pragma: no cover - import guard
                raise RuntimeError(
                    "Missing paddleocr. Install with `pip install paddleocr`."
                ) from exc
            if self._paddle is None:
                lang = "en" if self.language in {"eng", "en"} else self.language
                self._paddle = PaddleOCR(use_angle_cls=True, lang=lang, show_log=False)
            result = self._paddle.ocr(str(path), cls=True)
            lines = []
            for page in result:
                for item in page:
                    text = item[1][0]
                    if text:
                        lines.append(text)
            return "\n".join(lines)
        raise ValueError(f"Unsupported OCR engine: {self.engine}")


class WhisperTranscriber:
    def __init__(self, model_name: str, device: str, compute_type: str):
        try:
            from faster_whisper import WhisperModel
        except Exception as exc:  # pragma: no cover - import guard
            raise RuntimeError(
                "Missing faster-whisper. Install with `pip install faster-whisper`."
            ) from exc
        self.model = WhisperModel(model_name, device=device, compute_type=compute_type)

    def transcribe(self, audio_path: Path):
        return self.model.transcribe(
            str(audio_path),
            beam_size=5,
            vad_filter=True,
        )


def extract_audio_from_video(video_path: Path, tmp_dir: Path) -> Path:
    if shutil.which("ffmpeg") is None:
        raise RuntimeError("ffmpeg is required to process video files.")
    output_path = tmp_dir / f"{video_path.stem}_audio.wav"
    cmd = [
        "ffmpeg",
        "-y",
        "-i",
        str(video_path),
        "-vn",
        "-acodec",
        "pcm_s16le",
        "-ar",
        "16000",
        "-ac",
        "1",
        str(output_path),
    ]
    subprocess.run(cmd, check=True, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
    return output_path


# ---------------------------------------------------------------------------
# Chunked video processing helpers
# ---------------------------------------------------------------------------

def get_video_duration(video_path: Path) -> float:
    """Probe video duration in seconds using ffprobe."""
    if shutil.which("ffprobe") is None:
        raise RuntimeError("ffprobe is required (usually bundled with ffmpeg).")
    cmd = [
        "ffprobe", "-v", "error",
        "-show_entries", "format=duration",
        "-of", "csv=p=0",
        str(video_path),
    ]
    result = subprocess.run(cmd, capture_output=True, text=True, check=True)
    return float(result.stdout.strip())


def extract_audio_chunk(
    video_path: Path,
    tmp_dir: Path,
    start_sec: float,
    duration_sec: float,
    chunk_index: int,
) -> Path:
    """Extract a time-bounded audio chunk from a video using FFmpeg."""
    if shutil.which("ffmpeg") is None:
        raise RuntimeError("ffmpeg is required to process video files.")
    output_path = tmp_dir / f"{video_path.stem}_chunk_{chunk_index:04d}.wav"
    cmd = [
        "ffmpeg", "-y",
        "-ss", str(start_sec),
        "-i", str(video_path),
        "-t", str(duration_sec),
        "-vn", "-acodec", "pcm_s16le", "-ar", "16000", "-ac", "1",
        str(output_path),
    ]
    subprocess.run(cmd, check=True, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
    return output_path


def load_checkpoint(checkpoint_path: Path) -> Dict:
    """Load a video processing checkpoint, or return empty state."""
    if checkpoint_path.exists():
        with open(checkpoint_path, "r", encoding="utf-8") as f:
            return json.load(f)
    return {"completed_chunks": [], "total_segments": 0}


def save_checkpoint(checkpoint_path: Path, state: Dict) -> None:
    """Persist checkpoint state to disk."""
    with open(checkpoint_path, "w", encoding="utf-8") as f:
        json.dump(state, f, ensure_ascii=True)


# ---------------------------------------------------------------------------
# Keyframe extraction
# ---------------------------------------------------------------------------

def extract_keyframes_ffmpeg(
    video_path: Path,
    frames_dir: Path,
    scene_threshold: float = 0.3,
) -> List[Tuple[Path, float]]:
    """Extract scene-change keyframes from a video using FFmpeg.

    Returns list of (frame_path, timestamp_seconds) tuples.
    """
    if shutil.which("ffmpeg") is None:
        raise RuntimeError("ffmpeg is required for keyframe extraction.")
    frames_dir.mkdir(parents=True, exist_ok=True)
    cmd = [
        "ffmpeg", "-y",
        "-i", str(video_path),
        "-vf", f"select='gt(scene,{scene_threshold})',showinfo",
        "-vsync", "vfr",
        str(frames_dir / f"{video_path.stem}_frame_%04d.png"),
    ]
    result = subprocess.run(cmd, capture_output=True, text=True)

    # Parse pts_time from showinfo output on stderr
    frame_files = sorted(frames_dir.glob(f"{video_path.stem}_frame_*.png"))
    timestamps: List[float] = []
    for line in result.stderr.split("\n"):
        if "pts_time:" in line:
            match = re.search(r"pts_time:\s*([\d.]+)", line)
            if match:
                timestamps.append(float(match.group(1)))

    results: List[Tuple[Path, float]] = []
    for i, frame_path in enumerate(frame_files):
        ts = timestamps[i] if i < len(timestamps) else 0.0
        results.append((frame_path, ts))
    return results


def extract_keyframes_with_ocr(
    video_path: Path,
    doc_id: str,
    root: Path,
    output_dir: Path,
    ocr_runner: OcrRunner,
    scene_threshold: float = 0.3,
) -> List[Dict[str, object]]:
    """Extract keyframes from a video, run OCR, and produce frame records.

    Frame images are saved to <output_dir>/frames/<doc_id>/.
    Only frames with OCR text produce JSONL records.
    """
    frames_output_dir = output_dir / "frames" / doc_id
    keyframes = extract_keyframes_ffmpeg(video_path, frames_output_dir, scene_threshold)
    if not keyframes:
        return []

    try:
        rel_video_path = video_path.relative_to(root).as_posix()
    except ValueError:
        rel_video_path = video_path.as_posix()

    records: List[Dict[str, object]] = []
    for idx, (frame_path, timestamp_sec) in enumerate(keyframes):
        ocr_text = normalize_text(ocr_runner.extract_text(frame_path))
        if not ocr_text:
            continue

        try:
            rel_frame_path = frame_path.relative_to(output_dir).as_posix()
        except ValueError:
            rel_frame_path = frame_path.as_posix()

        source: Dict[str, object] = {
            "path": rel_video_path,
            "filename": video_path.name,
            "filetype": "video",
            "modality": "video_frame",
            "loader": f"ffmpeg+ocr:{ocr_runner.engine}",
            "timestamp_start": timestamp_sec,
            "timestamp_end": timestamp_sec,
            "timestamp": format_timestamp(timestamp_sec),
            "frame_path": rel_frame_path,
            "frame_index": idx,
            "scene_threshold": scene_threshold,
            "ocr_engine": ocr_runner.engine,
            "ocr_language": ocr_runner.language,
        }
        records.append({
            "doc_id": doc_id,
            "page": None,
            "section": f"frame_{len(records):04d}",
            "section_index": len(records),
            "content": ocr_text,
            "source": source,
        })
    return records


# ---------------------------------------------------------------------------
# Chunked video orchestrator
# ---------------------------------------------------------------------------

def process_video_chunked(
    video_path: Path,
    doc_id: str,
    root: Path,
    transcriber: WhisperTranscriber,
    output_dir: Path,
    chunk_duration: float = 600.0,
    extract_keyframes: bool = False,
    ocr_runner: Optional[OcrRunner] = None,
    scene_threshold: float = 0.3,
) -> List[Dict[str, object]]:
    """Process a video in time-based chunks with checkpointing.

    Returns all transcript records (and optionally keyframe records)
    with globally correct timestamps and segment numbering.
    """
    duration = get_video_duration(video_path)
    num_chunks = max(1, int(duration // chunk_duration) + (1 if duration % chunk_duration > 0 else 0))

    checkpoint_path = output_dir / f".checkpoint_{doc_id}.json"
    state = load_checkpoint(checkpoint_path)
    completed = set(state["completed_chunks"])
    segment_offset = state["total_segments"]

    all_transcript_records: List[Dict[str, object]] = []

    for chunk_idx in range(num_chunks):
        if chunk_idx in completed:
            continue

        start_sec = chunk_idx * chunk_duration
        this_duration = min(chunk_duration, duration - start_sec)

        with tempfile.TemporaryDirectory() as tmp_dir:
            audio_path = extract_audio_chunk(
                video_path, Path(tmp_dir), start_sec, this_duration, chunk_idx
            )
            chunk_records = records_from_audio(
                audio_path, doc_id, root, transcriber, "video",
                time_offset=start_sec, segment_offset=segment_offset,
            )

        # Patch source metadata to point at original video
        for rec in chunk_records:
            try:
                rec["source"]["path"] = video_path.relative_to(root).as_posix()
            except ValueError:
                rec["source"]["path"] = video_path.as_posix()
            rec["source"]["filename"] = video_path.name
            rec["source"]["filetype"] = "video"

        all_transcript_records.extend(chunk_records)
        segment_offset += len(chunk_records)

        completed.add(chunk_idx)
        state["completed_chunks"] = sorted(completed)
        state["total_segments"] = segment_offset
        save_checkpoint(checkpoint_path, state)

    # Keyframe extraction (optional)
    frame_records: List[Dict[str, object]] = []
    if extract_keyframes and ocr_runner is not None:
        frame_records = extract_keyframes_with_ocr(
            video_path, doc_id, root, output_dir, ocr_runner,
            scene_threshold=scene_threshold,
        )

    # Link transcript and frame records independently
    if all_transcript_records:
        add_relations(all_transcript_records)
    if frame_records:
        add_relations(frame_records)

    # Clean up checkpoint on success
    if checkpoint_path.exists():
        checkpoint_path.unlink()

    return all_transcript_records + frame_records


def records_from_csv(path: Path, doc_id: str, root: Path) -> List[Dict[str, object]]:
    records: List[Dict[str, object]] = []
    with open(path, newline="", encoding="utf-8") as handle:
        reader = csv.DictReader(handle)
        for row_idx, row in enumerate(reader, start=1):
            parts = []
            for key, value in row.items():
                if value is None:
                    continue
                field = (key or "").strip()
                if not field:
                    field = "column"
                cleaned = str(value).strip()
                if not cleaned:
                    continue
                parts.append(f"{field}: {cleaned}")
            content = normalize_text(" | ".join(parts))
            if not content:
                continue

            source = build_source_base(path, root, "csv", "table", "csv")
            source["row_index"] = row_idx
            source["column_count"] = len(reader.fieldnames or [])

            records.append(
                {
                    "doc_id": doc_id,
                    "page": None,
                    "section": f"row_{row_idx}",
                    "section_index": len(records),
                    "content": content,
                    "source": source,
                }
            )
    return records


def records_from_image(
    path: Path,
    doc_id: str,
    root: Path,
    ocr_runner: OcrRunner,
) -> List[Dict[str, object]]:
    content = normalize_text(ocr_runner.extract_text(path))
    if not content:
        return []
    source = build_source_base(path, root, "image", "image", f"ocr:{ocr_runner.engine}")
    source["ocr_engine"] = ocr_runner.engine
    source["ocr_language"] = ocr_runner.language
    return [
        {
            "doc_id": doc_id,
            "page": None,
            "section": "image_1",
            "section_index": 0,
            "content": content,
            "source": source,
        }
    ]


def records_from_audio(
    path: Path,
    doc_id: str,
    root: Path,
    transcriber: WhisperTranscriber,
    modality: str,
    time_offset: float = 0.0,
    segment_offset: int = 0,
) -> List[Dict[str, object]]:
    segments, info = transcriber.transcribe(path)
    language = getattr(info, "language", None)
    language_prob = getattr(info, "language_probability", None)

    records: List[Dict[str, object]] = []
    for idx, segment in enumerate(segments):
        text = normalize_text(getattr(segment, "text", "") or "")
        if not text:
            continue
        start = float(getattr(segment, "start", 0.0)) + time_offset
        end = float(getattr(segment, "end", 0.0)) + time_offset
        global_idx = segment_offset + len(records)
        source = build_source_base(path, root, modality, modality, "faster-whisper")
        source["timestamp_start"] = start
        source["timestamp_end"] = end
        source["timestamp"] = format_timestamp(start)
        if language:
            source["language"] = language
        if language_prob is not None:
            source["language_probability"] = round(float(language_prob), 4)

        records.append(
            {
                "doc_id": doc_id,
                "page": None,
                "section": f"segment_{global_idx:04d}",
                "section_index": global_idx,
                "content": text,
                "source": source,
            }
        )
    return records


def _resolve_modality(filetype: str) -> str:
    """Map a filetype string to its modality category."""
    if filetype == "image":
        return "image"
    if filetype == "audio":
        return "audio"
    if filetype == "video":
        return "video"
    if filetype == "csv":
        return "table"
    return "text"


def ingest(
    input_dir: Path,
    output_dir: Path,
    output_name: str,
    errors_name: str,
    pdf_strategy: str,
    ocr_engine: str,
    ocr_language: str,
    whisper_model: str,
    whisper_device: str,
    whisper_compute_type: str,
    max_files: Optional[int],
    chunk_duration: float = 600.0,
    extract_keyframes: bool = False,
    scene_threshold: float = 0.3,
) -> Dict[str, object]:
    output_dir.mkdir(parents=True, exist_ok=True)
    output_path = output_dir / output_name
    errors_path = output_dir / errors_name

    file_counts = Counter()
    record_counts = Counter()
    error_count = 0
    total_files = 0
    total_records = 0

    success_files: List[str] = []
    failed_files: List[Dict[str, str]] = []
    empty_files: List[Dict[str, str]] = []
    skipped_files: List[Dict[str, str]] = []

    ocr_runner = OcrRunner(ocr_engine, ocr_language)
    transcriber = None

    def _rel(path: Path) -> str:
        try:
            return path.relative_to(input_dir).as_posix()
        except ValueError:
            return path.as_posix()

    def log_error(path: Path, stage: str, err: Exception,
                  modality: str = "", filetype: str = "") -> None:
        nonlocal error_count
        error_count += 1
        error_entry: Dict[str, object] = {
            "path": str(path),
            "stage": stage,
            "error": str(err),
        }
        if modality:
            error_entry["modality"] = modality
        if filetype:
            error_entry["filetype"] = filetype
        error_handle.write(json.dumps(error_entry, ensure_ascii=True) + "\n")

    with open(output_path, "w", encoding="utf-8") as out_handle, open(
        errors_path, "w", encoding="utf-8"
    ) as error_handle:
        for path in iter_files(input_dir, max_files=max_files):
            filetype = detect_file_type(path)
            if not filetype:
                skipped_files.append({
                    "path": _rel(path),
                    "reason": f"Unsupported file type: {path.suffix}",
                })
                continue
            total_files += 1
            file_counts[filetype] += 1
            modality = _resolve_modality(filetype)
            doc_id = build_doc_id(path, input_dir)

            try:
                records: List[Dict[str, object]] = []
                if filetype == "csv":
                    records = records_from_csv(path, doc_id, input_dir)
                elif filetype == "image":
                    records = records_from_image(path, doc_id, input_dir, ocr_runner)
                elif filetype == "audio":
                    if transcriber is None:
                        transcriber = WhisperTranscriber(
                            whisper_model, whisper_device, whisper_compute_type
                        )
                    records = records_from_audio(
                        path, doc_id, input_dir, transcriber, "audio"
                    )
                elif filetype == "video":
                    if transcriber is None:
                        transcriber = WhisperTranscriber(
                            whisper_model, whisper_device, whisper_compute_type
                        )
                    records = process_video_chunked(
                        video_path=path,
                        doc_id=doc_id,
                        root=input_dir,
                        transcriber=transcriber,
                        output_dir=output_dir,
                        chunk_duration=chunk_duration,
                        extract_keyframes=extract_keyframes,
                        ocr_runner=ocr_runner if extract_keyframes else None,
                        scene_threshold=scene_threshold,
                    )
                else:
                    docs, loader_name = load_langchain_documents(path, filetype, pdf_strategy)
                    records = records_from_documents(
                        docs, doc_id, path, input_dir, filetype, "text", loader_name
                    )

                if records:
                    # Video records already have relations from process_video_chunked
                    if filetype != "video":
                        add_relations(records)
                    for record in records:
                        out_handle.write(json.dumps(record, ensure_ascii=True) + "\n")
                    total_records += len(records)
                    record_counts[filetype] += len(records)
                    success_files.append(_rel(path))
                else:
                    empty_files.append({
                        "path": _rel(path),
                        "modality": modality,
                        "reason": "Extraction produced no content",
                    })
                    log_error(path, "empty_content",
                              Exception("Extraction produced no content"),
                              modality=modality, filetype=filetype)
            except Exception as exc:
                log_error(path, "ingest", exc,
                          modality=modality, filetype=filetype)
                failed_files.append({
                    "path": _rel(path),
                    "modality": modality,
                    "reason": str(exc),
                })

    stats = {
        "input_root": str(input_dir),
        "output_file": str(output_path),
        "errors_file": str(errors_path),
        "generated_at": datetime.now(timezone.utc).isoformat().replace("+00:00", "Z"),
        "config": {
            "pdf_strategy": pdf_strategy,
            "ocr_engine": ocr_engine,
            "ocr_language": ocr_language,
            "whisper_model": whisper_model,
            "whisper_device": whisper_device,
            "whisper_compute_type": whisper_compute_type,
            "video_chunk_duration": chunk_duration,
            "extract_keyframes": extract_keyframes,
            "scene_threshold": scene_threshold if extract_keyframes else None,
        },
        "counts": {
            "files_processed": total_files,
            "files_succeeded": len(success_files),
            "files_failed": len(failed_files),
            "files_empty": len(empty_files),
            "files_skipped": len(skipped_files),
            "records_emitted": total_records,
            "errors": error_count,
        },
        "files_by_type": dict(file_counts),
        "records_by_type": dict(record_counts),
        "success_files": success_files,
        "failed_files": failed_files,
        "empty_files": empty_files,
        "skipped_files": skipped_files,
    }

    stats_path = output_dir / "ingestion_stats.json"
    with open(stats_path, "w", encoding="utf-8") as stats_handle:
        json.dump(stats, stats_handle, indent=2, ensure_ascii=True)

    return stats


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Ingest multimodal knowledge base files.")
    parser.add_argument(
        "--input",
        default=os.path.join("00_Data"),
        help="Root directory containing files to ingest.",
    )
    parser.add_argument(
        "--output-dir",
        default=os.path.join("01_data_ingestion", "output"),
        help="Output directory for JSONL and stats.",
    )
    parser.add_argument(
        "--output-name",
        default="cleaned_documents.jsonl",
        help="Output JSONL filename.",
    )
    parser.add_argument(
        "--errors-name",
        default="ingestion_errors.jsonl",
        help="Errors JSONL filename.",
    )
    parser.add_argument(
        "--pdf-strategy",
        default="auto",
        help="Unstructured PDF strategy (auto, hi_res, fast, ocr_only).",
    )
    parser.add_argument(
        "--ocr-engine",
        default="tesseract",
        choices=["tesseract", "paddle"],
        help="OCR engine for images.",
    )
    parser.add_argument(
        "--ocr-language",
        default="eng",
        help="OCR language code (tesseract or paddle).",
    )
    parser.add_argument(
        "--whisper-model",
        default="small",
        help="faster-whisper model size (tiny, base, small, medium, large-v3).",
    )
    parser.add_argument(
        "--whisper-device",
        default="cpu",
        help="faster-whisper device (cpu, cuda, auto).",
    )
    parser.add_argument(
        "--whisper-compute-type",
        default="int8",
        help="faster-whisper compute type (int8, float16, int8_float16).",
    )
    parser.add_argument(
        "--max-files",
        type=int,
        default=None,
        help="Optional cap on number of files processed.",
    )
    parser.add_argument(
        "--chunk-duration",
        type=float,
        default=600.0,
        help="Duration in seconds for video audio chunks (default: 600 = 10 min).",
    )
    parser.add_argument(
        "--extract-keyframes",
        action="store_true",
        default=False,
        help="Extract keyframes from videos and run OCR on them.",
    )
    parser.add_argument(
        "--scene-threshold",
        type=float,
        default=0.3,
        help="FFmpeg scene change threshold 0.0-1.0 (default: 0.3).",
    )
    parser.add_argument(
        "--workers",
        type=int,
        default=4,
        help="Number of parallel worker threads (0 = sequential legacy mode, default: 4).",
    )
    return parser.parse_args()


def main() -> None:
    args = parse_args()

    if args.workers > 0:
        from ingestion.engine import IngestionConfig, IngestionEngine
        from ingestion.events import LogEmitter

        config = IngestionConfig(
            input_dir=Path(args.input),
            output_dir=Path(args.output_dir),
            output_name=args.output_name,
            errors_name=args.errors_name,
            pdf_strategy=args.pdf_strategy,
            ocr_engine=args.ocr_engine,
            ocr_language=args.ocr_language,
            whisper_model=args.whisper_model,
            whisper_device=args.whisper_device,
            whisper_compute_type=args.whisper_compute_type,
            max_files=args.max_files,
            chunk_duration=args.chunk_duration,
            extract_keyframes=args.extract_keyframes,
            scene_threshold=args.scene_threshold,
            workers=args.workers,
        )
        engine = IngestionEngine(config, emitter=LogEmitter())
        engine.run()
    else:
        ingest(
            input_dir=Path(args.input),
            output_dir=Path(args.output_dir),
            output_name=args.output_name,
            errors_name=args.errors_name,
            pdf_strategy=args.pdf_strategy,
            ocr_engine=args.ocr_engine,
            ocr_language=args.ocr_language,
            whisper_model=args.whisper_model,
            whisper_device=args.whisper_device,
            whisper_compute_type=args.whisper_compute_type,
            max_files=args.max_files,
            chunk_duration=args.chunk_duration,
            extract_keyframes=args.extract_keyframes,
            scene_threshold=args.scene_threshold,
        )


if __name__ == "__main__":
    main()