"""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,
)