File size: 6,058 Bytes
fd1e711 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 | 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"
}
|