# rag/ingest_net.py import uuid import os import time import traceback from concurrent.futures import ThreadPoolExecutor, as_completed, TimeoutError from typing import List, Dict, Any from core.books.fetch_url import get_pdf_bytes from core.books.storage import ( init_db, check_db_health, insert_raw_document, mark_raw_status, ) from schemas.books.sources_schema import DocRaw MAX_INGEST_WORKERS = int(os.getenv("MAX_INGEST_WORKERS", "3")) MIN_PDF_SIZE_KB = int(os.getenv("MIN_PDF_SIZE_KB", "60")) HEAD_PAGES_N = int(os.getenv("HEAD_PAGES_N", "7")) FETCH_TIMEOUT_SECONDS = int(os.getenv("FETCH_TIMEOUT_SECONDS", "120")) EMBEDDING_MODEL = os.getenv( "EMBEDDING_MODEL", "sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2", ) def ingest_from_net( user_id: str, book_id: str, sources: List[Dict[str, Any]], ): """ Two-phase ingest pipeline. Phase A (Concurrent): - Fetch PDF - Extract text / OCR - Preprocess - Insert raw metadata into DB Phase B (Sequential): - Load embedding model once - Embed pages - Upsert to Qdrant """ start_total = time.time() print("=" * 80) print("[INGEST START]") print(f"user_id={user_id}") print(f"book_id={book_id}") print(f"sources_count={len(sources)}") print("=" * 80) init_db() if not check_db_health(): raise RuntimeError("Supabase / DB is not reachable") collection_name = f"user_{user_id}__book_{book_id}" # ============================================================ # Phase A Worker # ============================================================ def fetch_and_prepare(src: Dict[str, Any]) -> Dict[str, Any]: url = src.get("url") doc_id = str(uuid.uuid4()) start = time.time() result = { "url": url, "doc_id": doc_id, "status": "rejected", "reason": None, } try: print(f"\n[FETCH START] {url}") # ---------------------------------------------------- # Fetch PDF # ---------------------------------------------------- pdf_bytes = get_pdf_bytes(url) if not pdf_bytes: result["reason"] = "no_pdf_or_blocked" return result pdf_size_mb = round(len(pdf_bytes) / (1024 * 1024), 2) print(f"[FETCH DONE] {url} | size={pdf_size_mb} MB") # ---------------------------------------------------- # Validate size # ---------------------------------------------------- if len(pdf_bytes) < MIN_PDF_SIZE_KB * 1024: result["reason"] = "pdf_too_small" return result # ---------------------------------------------------- # Delayed imports # ---------------------------------------------------- from .pdf_text import ( extract_text_pypdf2, is_text_usable, ) from .ocr import mistral_ocr_pdf # ---------------------------------------------------- # Extract text # ---------------------------------------------------- print(f"[TEXT EXTRACTION START] {url}") pages = extract_text_pypdf2(pdf_bytes) joined = "\n".join(pages) # ---------------------------------------------------- # Text usable? # ---------------------------------------------------- if is_text_usable(joined): extraction_method = "text" print(f"[TEXT EXTRACTION SUCCESS] {url}") else: print(f"[OCR REQUIRED] {url}") # prevent OCR huge PDFs if len(pdf_bytes) > 10 * 1024 * 1024: result["reason"] = "pdf_too_large_for_ocr" return result try: print(f"[OCR START] {url}") pages = mistral_ocr_pdf(pdf_bytes) extraction_method = "ocr" print(f"[OCR DONE] {url}") except Exception as e: traceback.print_exc() result["reason"] = f"ocr_failed:{str(e)}" return result # ---------------------------------------------------- # Validate extraction # ---------------------------------------------------- if not pages: result["reason"] = "no_text_extracted" return result # ---------------------------------------------------- # Preprocessing # ---------------------------------------------------- print(f"[PREPROCESS START] {url}") from .preprocess import ( normalize_arabic, drop_common_headers_footers, ) language = src.get("language", "ar") if language == "ar": pages = [normalize_arabic(p) for p in pages] pages = drop_common_headers_footers(pages) pages_head = pages[:HEAD_PAGES_N] print(f"[PREPROCESS DONE] {url}") # ---------------------------------------------------- # Create raw doc # ---------------------------------------------------- raw_doc = DocRaw( doc_id=doc_id, user_id=user_id, book_id=book_id, source_url=url, source_type=src.get("source_type", "pdf"), domain=src.get("domain", ""), language=language, pages_head=pages_head, extraction_method=extraction_method, pdf_size_bytes=len(pdf_bytes), status="pending", error_reason="", ) # ---------------------------------------------------- # Insert DB row # ---------------------------------------------------- print(f"[DB INSERT START] {url}") insert_raw_document(raw_doc) print(f"[DB INSERT DONE] {url}") result.update( { "status": "prepared", "doc_id": doc_id, "raw_doc": raw_doc, "pages": pages, "extraction": extraction_method, "duration": round(time.time() - start, 2), } ) print( f"[PREPARED SUCCESS] {url} " f"| pages={len(pages)} " f"| duration={round(time.time() - start, 2)}s" ) return result except Exception as e: traceback.print_exc() result["reason"] = f"fetch_prepare_failed:{str(e)}" print(f"[PREPARED FAILED] {url} | {e}") return result # ============================================================ # Phase A - Concurrent # ============================================================ print("\n" + "=" * 80) print("[PHASE A START]") print("=" * 80) prepared_items = [] with ThreadPoolExecutor( max_workers=MAX_INGEST_WORKERS ) as executor: futures = { executor.submit(fetch_and_prepare, src): src for src in sources } for future in as_completed(futures): src = futures[future] try: res = future.result( timeout=FETCH_TIMEOUT_SECONDS ) except TimeoutError: res = { "url": src.get("url"), "status": "rejected", "reason": "timeout", } print( f"[TIMEOUT] {src.get('url')} " f"after {FETCH_TIMEOUT_SECONDS}s" ) except Exception as e: traceback.print_exc() res = { "url": src.get("url"), "status": "rejected", "reason": f"future_failed:{str(e)}", } prepared_items.append(res) print("=" * 80) print("[PHASE A DONE]") print("=" * 80) # ============================================================ # Load RAG Engine Once # ============================================================ print("\n[LOADING EMBEDDING MODEL + QDRANT]") from .rag_engine_sources import ArabicBookRAGWithSources rag = ArabicBookRAGWithSources( user_id=user_id, book_id=book_id, embedding_model=EMBEDDING_MODEL, ) print("[RAG ENGINE READY]") # ============================================================ # Phase B - Sequential Qdrant ingest # ============================================================ print("\n" + "=" * 80) print("[PHASE B START]") print("=" * 80) items_out = [] ingested = 0 rejected = 0 for item in prepared_items: if item.get("status") != "prepared": items_out.append(item) rejected += 1 continue doc_id = item["doc_id"] pages = item["pages"] raw_doc = item["raw_doc"] try: print(f"\n[QDRANT START] {doc_id}") stats = rag.ingest_pages( pages=pages, raw_doc=raw_doc, ) print(f"[QDRANT DONE] {doc_id}") items_out.append( { "url": item["url"], "status": "ingested", "doc_id": doc_id, "source_type": raw_doc.source_type, "extraction": item["extraction"], "stats": stats, } ) ingested += 1 # free memory item["pages"] = None except Exception as e: traceback.print_exc() try: mark_raw_status( doc_id, "failed", f"qdrant_failed:{str(e)}", ) except Exception: traceback.print_exc() items_out.append( { "url": item["url"], "status": "rejected", "reason": f"qdrant_failed:{str(e)}", "doc_id": doc_id, } ) rejected += 1 # ============================================================ # Done # ============================================================ total_duration = round( time.time() - start_total, 2, ) print("\n" + "=" * 80) print("[INGEST FINISHED]") print(f"ingested={ingested}") print(f"rejected={rejected}") print(f"duration={total_duration}s") print("=" * 80) return ( items_out, ingested, rejected, collection_name, )