Spaces:
Runtime error
Runtime error
File size: 5,703 Bytes
0a19322 | 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 | """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()
|