Spaces:
Runtime error
Runtime error
| """Upload pipeline: parse -> chunk -> embed -> store, run in the background with progress.""" | |
| import hashlib | |
| import io | |
| import logging | |
| import uuid | |
| from pathlib import Path | |
| from pypdf import PdfReader | |
| from . import store | |
| from .bm25 import bm25 | |
| from .chunking import chunk_text | |
| from .config import settings | |
| from .registry import get_registry | |
| log = logging.getLogger("rag.ingest") | |
| ALLOWED = (".pdf", ".txt", ".md") | |
| MIN_TEXT_CHARS = 20 # a page with less text than this counts as "empty" (blank or scanned image) | |
| def parse_file(filename: str, data: bytes, on_progress=None) -> list[tuple[int, str]]: | |
| """Return [(page_number, text), ...]. Page numbers are 1-based PDF page positions.""" | |
| name = filename.lower() | |
| if name.endswith(".pdf"): | |
| try: | |
| reader = PdfReader(io.BytesIO(data)) | |
| if reader.is_encrypted: | |
| reader.decrypt("") | |
| total = len(reader.pages) | |
| except Exception as e: | |
| raise ValueError(f"Could not open this PDF (password-protected or damaged): {e}") | |
| pages = [] | |
| for i, page in enumerate(reader.pages): | |
| try: | |
| text = page.extract_text() or "" | |
| except Exception: | |
| text = "" | |
| pages.append((i + 1, text)) | |
| if on_progress and (i % 10 == 0 or i == total - 1): | |
| on_progress(i + 1, total) | |
| return pages | |
| if name.endswith((".txt", ".md")): | |
| text = data.decode("utf-8", errors="ignore") | |
| return [(i + 1, t) for i, t in enumerate(text.split("\f"))] # form feed = page break | |
| raise ValueError("Unsupported file type. Use PDF, TXT or MD.") | |
| def submit_document(filename: str, data: bytes) -> tuple[dict, bool]: | |
| """Register an upload. Returns (entry, is_duplicate). Call process_document() afterwards.""" | |
| if not filename.lower().endswith(ALLOWED): | |
| raise ValueError("Unsupported file type. Use PDF, TXT or MD.") | |
| if not data: | |
| raise ValueError("The file is empty.") | |
| reg = get_registry() | |
| sha = hashlib.sha256(data).hexdigest() | |
| existing = reg.find_by_hash(sha) | |
| if existing: | |
| return existing, True | |
| entry = reg.create( | |
| { | |
| "doc_id": str(uuid.uuid4()), | |
| "filename": Path(filename).name, | |
| "sha256": sha, | |
| "status": "processing", | |
| "stage": "Queued", | |
| "progress": 0, | |
| "pages": 0, | |
| "chunks": 0, | |
| "empty_pages": [], | |
| "warning": None, | |
| "error": None, | |
| } | |
| ) | |
| return entry, False | |
| def process_document(doc_id: str, filename: str, data: bytes) -> None: | |
| reg = get_registry() | |
| try: | |
| reg.update(doc_id, stage="Reading pages", progress=1) | |
| def parse_progress(done, total): | |
| reg.update(doc_id, progress=1 + int(14 * done / max(total, 1))) | |
| pages = parse_file(filename, data, parse_progress) | |
| empty = [p for p, t in pages if len(t.strip()) < MIN_TEXT_CHARS] | |
| records = [] # (page, chunk_index, text) | |
| for page_no, text in pages: | |
| for chunk in chunk_text(text, settings.chunk_size, settings.chunk_overlap): | |
| records.append((page_no, len(records), chunk)) | |
| if not records: | |
| raise ValueError("No extractable text found. This looks like a scanned PDF (images), which needs OCR.") | |
| warning = None | |
| if len(empty) > 0.3 * len(pages): | |
| warning = ( | |
| f"{len(empty)} of {len(pages)} pages have no readable text (scanned images?). " | |
| "Questions about those pages cannot be answered without OCR." | |
| ) | |
| reg.update(doc_id, pages=len(pages), empty_pages=empty, warning=warning, stage="Embedding", progress=15) | |
| stem = Path(filename).stem | |
| total = len(records) | |
| for i in range(0, total, settings.embed_batch): | |
| batch = records[i : i + settings.embed_batch] | |
| # the header gives the embedding model the document name and page as context | |
| vectors = store.embed([f"{stem} | page {p}\n{t}" for p, _, t in batch]) | |
| points = [ | |
| store.PointStruct( | |
| id=str(uuid.uuid4()), | |
| vector=vec, | |
| payload={"doc_id": doc_id, "filename": filename, "page": p, "chunk_index": idx, "text": t}, | |
| ) | |
| for (p, idx, t), vec in zip(batch, vectors) | |
| ] | |
| store.upsert(points) | |
| done = min(i + len(batch), total) | |
| reg.update(doc_id, progress=15 + int(84 * done / total), chunks=done) | |
| bm25.mark_dirty() | |
| reg.update(doc_id, status="ready", stage="Ready", progress=100, chunks=total) | |
| log.info("Ingested %s: %d pages, %d chunks", filename, len(pages), total) | |
| except Exception as e: | |
| log.exception("Ingest failed for %s", filename) | |
| try: | |
| store.delete_doc(doc_id) | |
| except Exception: | |
| pass | |
| bm25.mark_dirty() | |
| reg.update(doc_id, status="failed", stage="Failed", error=str(e)) | |
| def delete_document(doc_id: str) -> bool: | |
| reg = get_registry() | |
| if not reg.get(doc_id): | |
| return False | |
| store.delete_doc(doc_id) | |
| reg.delete(doc_id) | |
| bm25.mark_dirty() | |
| return True | |
| def recover_interrupted() -> None: | |
| """If the server stopped mid-upload, clean up the half-finished documents.""" | |
| reg = get_registry() | |
| for d in reg.all(): | |
| if d["status"] == "processing": | |
| store.delete_doc(d["doc_id"]) | |
| reg.update(d["doc_id"], status="failed", stage="Failed", error="Interrupted by a server restart. Delete it and upload again.") | |
| bm25.mark_dirty() | |