import shutil import tempfile import uuid from pathlib import Path from fastapi import APIRouter, UploadFile, File, HTTPException, BackgroundTasks from pydantic import BaseModel from backend.ingestion.pdf_loader import ingest_pdf from backend.ingestion.url_loader import ingest_url from backend.ingestion.youtube_loader import ingest_youtube from backend.database.connection import get_connection router = APIRouter() # ── Request models ──────────────────────────────────────────────────────────── class UrlRequest(BaseModel): url: str language: str = "en" class YoutubeRequest(BaseModel): url: str language: str = "en" # ── Background Task Wrappers ────────────────────────────────────────────────── def run_ingest_pdf_bg(tmp_path: Path, filename: str, source_id: str): """Safely runs PDF ingestion and always cleans up the temporary file.""" try: ingest_pdf(tmp_path, filename, source_id) except Exception as e: print(f"[BG Ingest PDF] Critical failure: {e}") finally: if tmp_path.exists(): tmp_path.unlink() print(f"[BG Ingest PDF] Cleaned up temporary file: {tmp_path}") # ── Endpoints ───────────────────────────────────────────────────────────────── @router.post("/upload-pdf") async def upload_pdf(background_tasks: BackgroundTasks, file: UploadFile = File(...)): """ POST /upload-pdf Accepts a PDF file upload, saves it temporarily, starts ingestion in a background task, and immediately returns source ID. """ if not file.filename.endswith(".pdf"): raise HTTPException(status_code=400, detail="Only PDF files are accepted.") # Save uploaded file to a temp location first with tempfile.NamedTemporaryFile(delete=False, suffix=".pdf") as tmp: shutil.copyfileobj(file.file, tmp) tmp_path = Path(tmp.name) source_id = str(uuid.uuid4()) dest_path = tmp_path # temporary placeholder path until moved in background task try: # Create initial pending source entry in MySQL with get_connection() as conn: cursor = conn.cursor() cursor.execute(""" INSERT INTO sources (id, type, title, origin, language, status, progress_percentage, chunk_count) VALUES (%s, %s, %s, %s, %s, %s, %s, %s) """, (source_id, "pdf", file.filename, str(dest_path), "en", "UPLOADING", 10, 0)) conn.commit() # Dispatch background task background_tasks.add_task(run_ingest_pdf_bg, tmp_path, file.filename, source_id) except Exception as e: if tmp_path.exists(): tmp_path.unlink() raise HTTPException(status_code=500, detail=f"Database initialization failed: {str(e)}") return { "message" : "PDF upload initiated", "source_id" : source_id, "title" : file.filename, "chunk_count": 0, "status" : "UPLOADING" } @router.post("/add-url") def add_url(req: UrlRequest, background_tasks: BackgroundTasks): """ POST /add-url Accepts a website URL, scrapes and ingests it in the background. """ if not req.url.startswith("http"): raise HTTPException(status_code=400, detail="Invalid URL. Must start with http:// or https://") source_id = str(uuid.uuid4()) try: # Create initial pending source entry in MySQL with get_connection() as conn: cursor = conn.cursor() cursor.execute(""" INSERT INTO sources (id, type, title, origin, language, status, progress_percentage, chunk_count) VALUES (%s, %s, %s, %s, %s, %s, %s, %s) """, (source_id, "url", req.url, req.url, "en", "EXTRACTING", 10, 0)) conn.commit() # Dispatch background task background_tasks.add_task(ingest_url, req.url, source_id) except Exception as e: raise HTTPException(status_code=500, detail=f"Database initialization failed: {str(e)}") return { "message" : "URL ingestion initiated", "source_id" : source_id, "title" : req.url, "chunk_count": 0, "status" : "EXTRACTING" } @router.post("/add-youtube") def add_youtube(req: YoutubeRequest, background_tasks: BackgroundTasks): """ POST /add-youtube Accepts a YouTube URL, fetches transcript and ingests it in the background. """ if "youtube.com" not in req.url and "youtu.be" not in req.url: raise HTTPException(status_code=400, detail="Invalid YouTube URL.") source_id = str(uuid.uuid4()) try: # Create initial pending source entry in MySQL with get_connection() as conn: cursor = conn.cursor() cursor.execute(""" INSERT INTO sources (id, type, title, origin, language, status, progress_percentage, chunk_count) VALUES (%s, %s, %s, %s, %s, %s, %s, %s) """, (source_id, "youtube", req.url, req.url, "en", "EXTRACTING", 10, 0)) conn.commit() # Dispatch background task background_tasks.add_task(ingest_youtube, req.url, source_id) except Exception as e: raise HTTPException(status_code=500, detail=f"Database initialization failed: {str(e)}") return { "message" : "YouTube video ingestion initiated", "source_id" : source_id, "title" : req.url, "chunk_count": 0, "language" : req.language, "status" : "EXTRACTING" }