"""Parallel ingestion engine with event-driven progress reporting.""" from __future__ import annotations import json import threading import time from collections import Counter from concurrent.futures import ThreadPoolExecutor, as_completed from dataclasses import dataclass, field from datetime import datetime, timezone from pathlib import Path from typing import Any, Dict, List, Optional import sys, os sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..")) from ingest_multimodal import ( OcrRunner, WhisperTranscriber, add_relations, build_doc_id, detect_file_type, iter_files, load_langchain_documents, process_video_chunked, records_from_audio, records_from_csv, records_from_documents, records_from_image, _resolve_modality, ) from ingestion.events import ( EventEmitter, FileCompleted, FileEmpty, FileFailed, FileProgress, FileSkipped, FileStarted, IngestionCompleted, IngestionProgress, IngestionStarted, NullEmitter, ) @dataclass class IngestionConfig: """All configuration for an ingestion run.""" input_dir: Path output_dir: Path output_name: str = "cleaned_documents.jsonl" errors_name: str = "ingestion_errors.jsonl" pdf_strategy: str = "auto" ocr_engine: str = "tesseract" ocr_language: str = "eng" whisper_model: str = "small" whisper_device: str = "cpu" whisper_compute_type: str = "int8" max_files: Optional[int] = None chunk_duration: float = 600.0 extract_keyframes: bool = False scene_threshold: float = 0.3 workers: int = 4 @dataclass class FileResult: """Result from processing a single file.""" file_index: int path: Path filetype: str modality: str records: List[Dict[str, Any]] = field(default_factory=list) error: Optional[str] = None empty: bool = False duration_sec: float = 0.0 class CancelledError(Exception): pass class IngestionEngine: """Concurrent ingestion engine using ThreadPoolExecutor. - PDF / text / CSV / image work runs freely in parallel threads - Audio / video (Whisper) is serialized via ``_whisper_lock`` to avoid model contention (safe on both CPU and GPU) - Cooperative cancellation via ``threading.Event`` - Emits SSE events for every lifecycle transition """ def __init__( self, config: IngestionConfig, emitter: EventEmitter | None = None, ) -> None: self.config = config self.emitter: EventEmitter = emitter or NullEmitter() self._cancel = threading.Event() self._whisper_lock = threading.Lock() self._transcriber: Optional[WhisperTranscriber] = None self._ocr_runner = OcrRunner(config.ocr_engine, config.ocr_language) # Thread-safe progress counters self._lock = threading.Lock() self._files_completed = 0 self._total_records = 0 self._start_time: float = 0.0 # ------------------------------------------------------------------ # Public API # ------------------------------------------------------------------ def run(self) -> Dict[str, Any]: """Execute the ingestion pipeline. Returns stats dict.""" self.config.output_dir.mkdir(parents=True, exist_ok=True) self._start_time = time.monotonic() # Discover and classify files all_files = list(iter_files(self.config.input_dir, self.config.max_files)) supported, skipped = self._classify_files(all_files) self.emitter.emit(IngestionStarted( total_files=len(supported), skipped_files=skipped, )) # Emit individual skip events for s in skipped: self.emitter.emit(FileSkipped(file_path=s["path"], reason=s["reason"])) # Process files if self.config.workers <= 0 or len(supported) == 0: results = self._run_sequential(supported) else: results = self._run_parallel(supported) # Sort by original file index for deterministic output results.sort(key=lambda r: r.file_index) # Write output and build stats stats = self._write_output(results, skipped) self.emitter.emit(IngestionCompleted(stats=stats)) return stats def cancel(self) -> None: """Signal all workers to stop after their current file.""" self._cancel.set() @property def is_cancelled(self) -> bool: return self._cancel.is_set() # ------------------------------------------------------------------ # File classification # ------------------------------------------------------------------ def _classify_files( self, files: List[Path] ) -> tuple[List[tuple[int, Path, str, str]], List[Dict[str, str]]]: """Split files into supported (with index, path, filetype, modality) and skipped (with path + reason).""" supported: List[tuple[int, Path, str, str]] = [] skipped: List[Dict[str, str]] = [] idx = 0 for path in files: filetype = detect_file_type(path) if not filetype: rel = self._rel(path) skipped.append({ "path": rel, "reason": f"Unsupported file type: {path.suffix}", }) continue modality = _resolve_modality(filetype) supported.append((idx, path, filetype, modality)) idx += 1 return supported, skipped # ------------------------------------------------------------------ # Execution strategies # ------------------------------------------------------------------ def _run_sequential( self, supported: List[tuple[int, Path, str, str]] ) -> List[FileResult]: results: List[FileResult] = [] for idx, path, filetype, modality in supported: if self._cancel.is_set(): break result = self._process_file(idx, path, filetype, modality, len(supported)) results.append(result) self._update_progress(result, len(supported)) return results def _run_parallel( self, supported: List[tuple[int, Path, str, str]] ) -> List[FileResult]: results: List[FileResult] = [] total = len(supported) with ThreadPoolExecutor(max_workers=self.config.workers) as pool: futures = { pool.submit( self._process_file, idx, path, filetype, modality, total ): idx for idx, path, filetype, modality in supported } for future in as_completed(futures): if self._cancel.is_set(): pool.shutdown(wait=False, cancel_futures=True) break result = future.result() results.append(result) self._update_progress(result, total) return results # ------------------------------------------------------------------ # Per-file processing # ------------------------------------------------------------------ def _process_file( self, file_index: int, path: Path, filetype: str, modality: str, total_files: int, ) -> FileResult: if self._cancel.is_set(): return FileResult( file_index=file_index, path=path, filetype=filetype, modality=modality, error="Cancelled", ) rel_path = self._rel(path) self.emitter.emit(FileStarted( file_path=rel_path, file_index=file_index, total_files=total_files, modality=modality, filetype=filetype, )) t0 = time.monotonic() doc_id = build_doc_id(path, self.config.input_dir) try: records = self._extract_records(path, filetype, modality, doc_id, rel_path) except Exception as exc: duration = time.monotonic() - t0 self.emitter.emit(FileFailed( file_path=rel_path, error=str(exc), modality=modality )) return FileResult( file_index=file_index, path=path, filetype=filetype, modality=modality, error=str(exc), duration_sec=duration, ) duration = time.monotonic() - t0 if not records: self.emitter.emit(FileEmpty( file_path=rel_path, modality=modality, reason="Extraction produced no content", )) return FileResult( file_index=file_index, path=path, filetype=filetype, modality=modality, empty=True, duration_sec=duration, ) self.emitter.emit(FileCompleted( file_path=rel_path, records_emitted=len(records), duration_sec=duration, )) return FileResult( file_index=file_index, path=path, filetype=filetype, modality=modality, records=records, duration_sec=duration, ) def _extract_records( self, path: Path, filetype: str, modality: str, doc_id: str, rel_path: str, ) -> List[Dict[str, Any]]: """Route to the correct extraction function. Audio and video acquire the whisper lock; everything else runs freely. """ root = self.config.input_dir if filetype == "csv": records = records_from_csv(path, doc_id, root) add_relations(records) return records if filetype == "image": records = records_from_image(path, doc_id, root, self._ocr_runner) add_relations(records) return records if filetype == "audio": with self._whisper_lock: self._ensure_transcriber() records = records_from_audio( path, doc_id, root, self._transcriber, "audio" ) add_relations(records) return records if filetype == "video": total_chunks = self._estimate_video_chunks(path) with self._whisper_lock: self._ensure_transcriber() records = self._process_video_with_progress( path, doc_id, rel_path, total_chunks ) return records # relations already added by process_video_chunked # Text types (pdf, docx, pptx, txt, md, html) docs, loader_name = load_langchain_documents( path, filetype, self.config.pdf_strategy ) records = records_from_documents( docs, doc_id, path, root, filetype, "text", loader_name ) add_relations(records) return records def _process_video_with_progress( self, path: Path, doc_id: str, rel_path: str, total_chunks: int, ) -> List[Dict[str, Any]]: """Wrap process_video_chunked with chunk-level progress events.""" self.emitter.emit(FileProgress( file_path=rel_path, detail=f"Video: {total_chunks} chunks to process", percent=0.0, )) records = process_video_chunked( video_path=path, doc_id=doc_id, root=self.config.input_dir, transcriber=self._transcriber, output_dir=self.config.output_dir, chunk_duration=self.config.chunk_duration, extract_keyframes=self.config.extract_keyframes, ocr_runner=self._ocr_runner if self.config.extract_keyframes else None, scene_threshold=self.config.scene_threshold, ) return records def _estimate_video_chunks(self, path: Path) -> int: """Get approximate chunk count without holding the whisper lock.""" try: from ingest_multimodal import get_video_duration duration = get_video_duration(path) cd = self.config.chunk_duration return max(1, int(duration // cd) + (1 if duration % cd > 0 else 0)) except Exception: return 1 # ------------------------------------------------------------------ # Progress tracking # ------------------------------------------------------------------ def _update_progress(self, result: FileResult, total_files: int) -> None: with self._lock: self._files_completed += 1 self._total_records += len(result.records) elapsed = time.monotonic() - self._start_time completed = self._files_completed records = self._total_records # ETA: simple linear extrapolation eta = None if completed > 0 and completed < total_files: avg_per_file = elapsed / completed remaining = total_files - completed eta = avg_per_file * remaining self.emitter.emit(IngestionProgress( files_completed=completed, total_files=total_files, records_emitted=records, elapsed_sec=round(elapsed, 2), eta_sec=round(eta, 2) if eta is not None else None, )) # ------------------------------------------------------------------ # Output writing # ------------------------------------------------------------------ def _write_output( self, results: List[FileResult], skipped: List[Dict[str, str]], ) -> Dict[str, Any]: output_path = self.config.output_dir / self.config.output_name errors_path = self.config.output_dir / self.config.errors_name file_counts: Counter = Counter() record_counts: Counter = Counter() success_files: List[str] = [] failed_files: List[Dict[str, str]] = [] empty_files: List[Dict[str, str]] = [] total_records = 0 with open(output_path, "w", encoding="utf-8") as out_handle, \ open(errors_path, "w", encoding="utf-8") as error_handle: for result in results: file_counts[result.filetype] += 1 rel = self._rel(result.path) if result.error: failed_files.append({ "path": rel, "modality": result.modality, "reason": result.error, }) error_entry = { "path": rel, "stage": "ingest", "error": result.error, "modality": result.modality, "filetype": result.filetype, } error_handle.write(json.dumps(error_entry, ensure_ascii=True) + "\n") elif result.empty: empty_files.append({ "path": rel, "modality": result.modality, "reason": "Extraction produced no content", }) error_entry = { "path": rel, "stage": "empty_content", "error": "Extraction produced no content", "modality": result.modality, "filetype": result.filetype, } error_handle.write(json.dumps(error_entry, ensure_ascii=True) + "\n") else: for record in result.records: out_handle.write(json.dumps(record, ensure_ascii=True) + "\n") total_records += len(result.records) record_counts[result.filetype] += len(result.records) success_files.append(rel) elapsed = time.monotonic() - self._start_time stats: Dict[str, Any] = { "input_root": str(self.config.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": self.config.pdf_strategy, "ocr_engine": self.config.ocr_engine, "ocr_language": self.config.ocr_language, "whisper_model": self.config.whisper_model, "whisper_device": self.config.whisper_device, "whisper_compute_type": self.config.whisper_compute_type, "video_chunk_duration": self.config.chunk_duration, "extract_keyframes": self.config.extract_keyframes, "scene_threshold": ( self.config.scene_threshold if self.config.extract_keyframes else None ), "workers": self.config.workers, }, "counts": { "files_processed": len(results), "files_succeeded": len(success_files), "files_failed": len(failed_files), "files_empty": len(empty_files), "files_skipped": len(skipped), "records_emitted": total_records, "errors": len(failed_files) + len(empty_files), }, "duration_sec": round(elapsed, 2), "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, } stats_path = self.config.output_dir / "ingestion_stats.json" with open(stats_path, "w", encoding="utf-8") as f: json.dump(stats, f, indent=2, ensure_ascii=True) return stats # ------------------------------------------------------------------ # Helpers # ------------------------------------------------------------------ def _ensure_transcriber(self) -> None: """Lazily initialize the Whisper transcriber (must hold _whisper_lock).""" if self._transcriber is None: self._transcriber = WhisperTranscriber( self.config.whisper_model, self.config.whisper_device, self.config.whisper_compute_type, ) def _rel(self, path: Path) -> str: try: return path.relative_to(self.config.input_dir).as_posix() except ValueError: return path.as_posix()