rag/ingest_net.py

# 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,
    )