Update rag/ingest_net.py

rag/ingest_net.py

@@ -1,11 +1,9 @@
 # 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 concurrent.futures import ThreadPoolExecutor, as_completed
+from typing import List, Dict, Any, Tuple
 
 from core.books.fetch_url import get_pdf_bytes
 from core.books.storage import (
@@ -16,198 +14,81 @@ from core.books.storage import (
 )
 from schemas.books.sources_schema import DocRaw
 
-MAX_INGEST_WORKERS = int(os.getenv("MAX_INGEST_WORKERS", "3"))
-
+MAX_INGEST_WORKERS = int(os.getenv("MAX_INGEST_WORKERS", "6"))
 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",
 )
 
+HEAD_PAGES_N = int(os.getenv("HEAD_PAGES_N", "7"))
 
-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
+def ingest_from_net(user_id: str, book_id: str, sources: List[Dict[str, Any]]):
+    """
+    Two-phase ingest:
+      Phase A (concurrent): fetch PDFs, extract text/OCR, preprocess, insert raw doc to DB.
+      Phase B (sequential): load embedder once and ingest pages (encode + upsert) to Qdrant.
+    This avoids loading the embedding model many times (huge slow-down).
     """
-
-    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
-    # ============================================================
-
+    # Worker for phase A: fetch + extract + preprocess + insert_raw_document
     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,
-        }
+        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,
-            )
-
+            # delayed imports for optional OCR libs
+            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,
-            )
+            # preprocess
+            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,
@@ -223,16 +104,9 @@ def ingest_from_net(
                 error_reason="",
             )
 
-            # ----------------------------------------------------
-            # Insert DB row
-            # ----------------------------------------------------
-
-            print(f"[DB INSERT START] {url}")
-
+            # insert raw doc (phase 1)
             insert_raw_document(raw_doc)
 
-            print(f"[DB INSERT DONE] {url}")
-
             result.update(
                 {
                     "status": "prepared",
@@ -240,143 +114,45 @@ def ingest_from_net(
                     "raw_doc": raw_doc,
                     "pages": pages,
                     "extraction": extraction_method,
-                    "duration": round(time.time() - start, 2),
+                    "duration": time.time() - start,
                 }
             )
-
-            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)
-
+    # Phase A: concurrent fetch + prepare
     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)}",
-                }
-
+    with ThreadPoolExecutor(max_workers=MAX_INGEST_WORKERS) as ex:
+        futures = {ex.submit(fetch_and_prepare, src): src for src in sources}
+        for fut in as_completed(futures):
+            res = fut.result()
             prepared_items.append(res)
 
-    print("=" * 80)
-    print("[PHASE A DONE]")
-    print("=" * 80)
-
-    # ============================================================
-    # Load RAG Engine Once
-    # ============================================================
-
-    print("\n[LOADING EMBEDDING MODEL + QDRANT]")
-
+    # Phase B: sequential embedding & qdrant upsert with single embedder instance
     from .rag_engine_sources import ArabicBookRAGWithSources
 
     rag = ArabicBookRAGWithSources(
-        user_id=user_id,
-        book_id=book_id,
-        embedding_model=EMBEDDING_MODEL,
+        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}")
-
+            stats = rag.ingest_pages(pages=pages, raw_doc=raw_doc)
             items_out.append(
                 {
                     "url": item["url"],
@@ -387,27 +163,13 @@ def ingest_from_net(
                     "stats": stats,
                 }
             )
-
             ingested += 1
-
-            # free memory
-            item["pages"] = None
-
         except Exception as e:
-
-            traceback.print_exc()
-
+            # mark failed on db
             try:
-
-                mark_raw_status(
-                    doc_id,
-                    "failed",
-                    f"qdrant_failed:{str(e)}",
-                )
-
+                mark_raw_status(doc_id, "failed", f"qdrant_failed:{str(e)}")
             except Exception:
-                traceback.print_exc()
-
+                pass
             items_out.append(
                 {
                     "url": item["url"],
@@ -416,28 +178,6 @@ def ingest_from_net(
                     "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,
-    )
+    return items_out, ingested, rejected, rag.collection