"""FastAPI router for ingestion: start, upload, SSE stream, cancel, status.""" from __future__ import annotations import shutil import tempfile import threading import time import uuid import zipfile from dataclasses import dataclass, field from enum import Enum from pathlib import Path from typing import Any, Dict, List, Optional from fastapi import APIRouter, HTTPException, UploadFile, File, Form from fastapi.responses import StreamingResponse from pydantic import BaseModel from ingestion.engine import IngestionConfig, IngestionEngine from ingestion.events import QueueEmitter, IngestionCompleted router = APIRouter(prefix="/ingestion", tags=["ingestion"]) # --------------------------------------------------------------------------- # Job management # --------------------------------------------------------------------------- class JobStatus(str, Enum): PENDING = "pending" RUNNING = "running" COMPLETED = "completed" FAILED = "failed" CANCELLING = "cancelling" CANCELLED = "cancelled" @dataclass class Job: job_id: str status: JobStatus engine: IngestionEngine emitter: QueueEmitter thread: Optional[threading.Thread] = None stats: Optional[Dict[str, Any]] = None error: Optional[str] = None temp_dir: Optional[str] = None # temp dir to clean up after completion created_at: float = field(default_factory=time.time) # Snapshot of progress for the polling endpoint files_completed: int = 0 total_files: int = 0 records_emitted: int = 0 # In-memory job store (single-server deployment) _jobs: Dict[str, Job] = {} # Auto-expire jobs older than this (seconds) _JOB_TTL = 3600 def _cleanup_expired_jobs() -> None: now = time.time() expired = [ jid for jid, job in _jobs.items() if now - job.created_at > _JOB_TTL and job.status in (JobStatus.COMPLETED, JobStatus.FAILED, JobStatus.CANCELLED) ] for jid in expired: job = _jobs.pop(jid, None) if job and job.temp_dir: shutil.rmtree(job.temp_dir, ignore_errors=True) def _get_job(job_id: str) -> Job: _cleanup_expired_jobs() job = _jobs.get(job_id) if job is None: raise HTTPException(status_code=404, detail=f"Job {job_id} not found") return job def _run_job(job: Job) -> None: """Background thread target: runs the engine and updates job state.""" try: job.status = JobStatus.RUNNING stats = job.engine.run() job.stats = stats if job.engine.is_cancelled: job.status = JobStatus.CANCELLED else: job.status = JobStatus.COMPLETED except Exception as exc: job.error = str(exc) job.status = JobStatus.FAILED finally: job.emitter.close() if job.temp_dir: shutil.rmtree(job.temp_dir, ignore_errors=True) job.temp_dir = None def _start_job(config: IngestionConfig, temp_dir: Optional[str] = None) -> Job: """Create an engine + emitter, start in a background thread, return the job.""" job_id = uuid.uuid4().hex[:12] emitter = QueueEmitter() engine = IngestionEngine(config, emitter=emitter) job = Job( job_id=job_id, status=JobStatus.PENDING, engine=engine, emitter=emitter, temp_dir=temp_dir, ) _jobs[job_id] = job t = threading.Thread(target=_run_job, args=(job,), daemon=True) job.thread = t t.start() return job # --------------------------------------------------------------------------- # Request / response models # --------------------------------------------------------------------------- class StartRequest(BaseModel): input_dir: str output_dir: str = "01_data_ingestion/output" 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 class JobResponse(BaseModel): job_id: str status: str class UploadResponse(BaseModel): job_id: str status: str uploaded_files: int class StatusResponse(BaseModel): job_id: str status: str files_completed: int total_files: int records_emitted: int stats: Optional[Dict[str, Any]] = None error: Optional[str] = None # --------------------------------------------------------------------------- # Endpoints # --------------------------------------------------------------------------- @router.post("/start", response_model=JobResponse) def start_ingestion(req: StartRequest): """Start ingestion from a server-side directory path.""" input_path = Path(req.input_dir) if not input_path.is_dir(): raise HTTPException(status_code=400, detail=f"Directory not found: {req.input_dir}") config = IngestionConfig( input_dir=input_path, output_dir=Path(req.output_dir), output_name=req.output_name, errors_name=req.errors_name, pdf_strategy=req.pdf_strategy, ocr_engine=req.ocr_engine, ocr_language=req.ocr_language, whisper_model=req.whisper_model, whisper_device=req.whisper_device, whisper_compute_type=req.whisper_compute_type, max_files=req.max_files, chunk_duration=req.chunk_duration, extract_keyframes=req.extract_keyframes, scene_threshold=req.scene_threshold, workers=req.workers, ) job = _start_job(config) return JobResponse(job_id=job.job_id, status=job.status.value) @router.post("/upload", response_model=UploadResponse) async def upload_and_ingest( files: List[UploadFile] = File(...), output_dir: str = Form(default="01_data_ingestion/output"), pdf_strategy: str = Form(default="auto"), ocr_engine: str = Form(default="tesseract"), ocr_language: str = Form(default="eng"), whisper_model: str = Form(default="small"), whisper_device: str = Form(default="cpu"), whisper_compute_type: str = Form(default="int8"), workers: int = Form(default=4), ): """Upload files (or a .zip) and start ingestion.""" temp_dir = tempfile.mkdtemp(prefix="cograg_upload_") upload_dir = Path(temp_dir) file_count = 0 for upload in files: filename = upload.filename or f"file_{file_count}" dest = upload_dir / filename dest.parent.mkdir(parents=True, exist_ok=True) content = await upload.read() if filename.lower().endswith(".zip"): # Extract zip contents into the temp directory zip_path = upload_dir / filename with open(zip_path, "wb") as f: f.write(content) with zipfile.ZipFile(zip_path, "r") as zf: zf.extractall(upload_dir) zip_path.unlink() # Count extracted files (excluding directories) file_count += sum(1 for p in upload_dir.rglob("*") if p.is_file()) else: with open(dest, "wb") as f: f.write(content) file_count += 1 config = IngestionConfig( input_dir=upload_dir, output_dir=Path(output_dir), 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, workers=workers, ) job = _start_job(config, temp_dir=temp_dir) return UploadResponse( job_id=job.job_id, status=job.status.value, uploaded_files=file_count ) @router.get("/{job_id}/stream") async def stream_events(job_id: str): """SSE endpoint that streams real-time ingestion events.""" job = _get_job(job_id) async def event_generator(): async for sse_text in job.emitter.stream(): yield sse_text return StreamingResponse( event_generator(), media_type="text/event-stream", headers={ "Cache-Control": "no-cache", "Connection": "keep-alive", "X-Accel-Buffering": "no", }, ) @router.post("/{job_id}/cancel", response_model=JobResponse) def cancel_ingestion(job_id: str): """Cancel a running ingestion job.""" job = _get_job(job_id) if job.status not in (JobStatus.PENDING, JobStatus.RUNNING): raise HTTPException( status_code=400, detail=f"Cannot cancel job in state: {job.status.value}", ) job.engine.cancel() job.status = JobStatus.CANCELLING return JobResponse(job_id=job.job_id, status=job.status.value) @router.get("/{job_id}/status", response_model=StatusResponse) def get_status(job_id: str): """Polling fallback: get current job progress.""" job = _get_job(job_id) # Pull latest counts from engine internals engine = job.engine with engine._lock: files_completed = engine._files_completed records_emitted = engine._total_records return StatusResponse( job_id=job.job_id, status=job.status.value, files_completed=files_completed, total_files=job.total_files, records_emitted=records_emitted, stats=job.stats, error=job.error, )