src/document_processor.py

#!/usr/bin/env python3
"""
Dynamic RAG Database Updater
Processes new PDFs and updates the vector database in real-time

"""

# Disable PaddleOCR's connectivity check — models are already cached after
# first run; this check adds ~10-30s to every cold start for nothing.
import os
os.environ.setdefault("PADDLE_PDX_DISABLE_MODEL_SOURCE_CHECK", "True")

import json
import numpy as np
from pathlib import Path
from typing import List, Dict
import pickle
from datetime import datetime

# PDF processing
import fitz

# OCR (CPU optimized)
from paddleocr import PaddleOCR

# Embeddings
from sentence_transformers import SentenceTransformer

# FAISS (CPU)
import faiss


class DynamicRAGUpdater:
    """
    Handles dynamic updates to RAG database:
    1. Upload PDF
    2. OCR extraction (PaddleOCR CPU)
    3. Generate embeddings (BiomedBERT)
    4. Update FAISS index
    5. Update metadata
    """

    def __init__(
        self,
        vector_db_path: str,
        embedding_model: str = "microsoft/BiomedNLP-BiomedBERT-base-uncased-abstract-fulltext",
        upload_dir: str = "uploaded_reports",
        ocr_engine=None,
        embedder=None,
        ocr_dpi: int = 300,
    ):
        self.vector_db_path = Path(vector_db_path)
        self.upload_dir = Path(upload_dir)
        self.upload_dir.mkdir(exist_ok=True)
        self.ocr_dpi = ocr_dpi

        self.ocr_dir = self.upload_dir / "ocr_text"
        self.embeddings_dir = self.upload_dir / "embeddings"
        self.ocr_dir.mkdir(exist_ok=True)
        self.embeddings_dir.mkdir(exist_ok=True)

        # Both models are expensive to load (PaddleOCR init + BiomedBERT weights).
        # Callers (e.g. the Streamlit app) should build these once via a cached
        # resource and pass them in so repeated uploads reuse warm models instead
        # of reloading from disk every time.
        if ocr_engine is not None:
            self.ocr = ocr_engine
        else:
            # PaddleOCR 3.x API. Doc-orientation/unwarping/textline-orientation
            # sub-models are disabled: pathology scans are upright, and skipping
            # them avoids extra model downloads and speeds up CPU inference.
            self.ocr = PaddleOCR(
                lang="en",
                use_doc_orientation_classify=False,
                use_doc_unwarping=False,
                use_textline_orientation=False,
            )

        if embedder is not None:
            self.embedding_model = embedder
        else:
            self.embedding_model = SentenceTransformer(
                embedding_model,
                device="cpu"
            )

        self.embedding_dim = self.embedding_model.get_sentence_embedding_dimension()

        self.load_database()

    def load_database(self):
        index_file = self.vector_db_path / "faiss.index"
        metadata_file = self.vector_db_path / "metadata.pkl"

        self.faiss_index = faiss.read_index(str(index_file))

        with open(metadata_file, "rb") as f:
            data = pickle.load(f)
            self.chunks = data["chunks"]
            self.chunk_id_to_idx = data.get("chunk_id_to_idx", {})

    def save_database(self):
        faiss.write_index(self.faiss_index, str(self.vector_db_path / "faiss.index"))

        with open(self.vector_db_path / "metadata.pkl", "wb") as f:
            pickle.dump(
                {
                    "chunks": self.chunks,
                    "chunk_id_to_idx": self.chunk_id_to_idx,
                    "embedding_dim": self.embedding_dim,
                    "model": "microsoft/BiomedNLP-BiomedBERT-base-uncased-abstract-fulltext"
                },
                f
            )

    def extract_text_from_pdf(self, pdf_path: str, progress_callback=None) -> Dict:
        """Runs OCR page by page, keeping each line's page number and bounding
        box (converted from pixmap pixel space back to PDF point space) so
        chunks can later be traced back to a highlightable location.
        Returns {"full_text": str, "lines": [{"page", "text", "bbox"}, ...]}.
        Optional progress_callback(stage, detail) reports per-page OCR progress.
        """
        doc = fitz.open(pdf_path)
        scale = self.ocr_dpi / 72.0  # pixmap pixels -> PDF points
        num_pages = len(doc)

        full_text_parts = []
        all_lines = []

        for page_num in range(num_pages):
            if progress_callback:
                progress_callback("ocr", {"page": page_num + 1, "total": num_pages})
            page = doc.load_page(page_num)
            pix = page.get_pixmap(dpi=self.ocr_dpi, alpha=False)
            image_np = np.frombuffer(pix.samples, dtype=np.uint8).reshape(pix.h, pix.w, pix.n)

            # PaddleOCR 3.x: predict() -> [OCRResult]; result[0] is a dict with
            # parallel lists rec_texts (strings) and dt_polys (Nx2 point arrays).
            ocr_result = self.ocr.predict(image_np)

            page_text = []
            if ocr_result:
                res = ocr_result[0]
                texts = res["rec_texts"]
                polys = res["dt_polys"]
                for line_text, poly in zip(texts, polys):
                    xs = [float(pt[0]) for pt in poly]
                    ys = [float(pt[1]) for pt in poly]
                    bbox_pdf = [
                        min(xs) / scale, min(ys) / scale,
                        max(xs) / scale, max(ys) / scale,
                    ]
                    page_text.append(line_text)
                    all_lines.append({
                        "page": page_num + 1,
                        "text": line_text,
                        "bbox": bbox_pdf,
                    })

            full_text_parts.append(
                f"\n{'='*50}\nPAGE {page_num + 1}\n{'='*50}\n" +
                "\n".join(page_text)
            )

        return {
            "full_text": "\n".join(full_text_parts),
            "lines": all_lines,
        }

    def chunk_text(self, lines: List[Dict], chunk_size: int = 512) -> List[Dict]:
        """Groups OCR lines into ~chunk_size-character chunks, sentence-aware
        like before, but keeps each chunk's originating page and the bounding
        boxes of every line that fell inside it for PDF highlighting.
        Returns [{"text", "page", "line_bboxes"}, ...].
        """
        chunks = []
        current_text = []
        current_bboxes = []
        current_page = None
        length = 0

        def flush():
            if current_text:
                chunks.append({
                    "text": "".join(current_text),
                    "page": current_page,
                    "line_bboxes": list(current_bboxes),
                })

        for line in lines:
            s = line["text"].strip()
            if not s:
                continue
            s = s + ". "

            if length + len(s) > chunk_size and current_text:
                flush()
                current_text = [s]
                current_bboxes = [line["bbox"]]
                current_page = line["page"]
                length = len(s)
            else:
                if current_page is None:
                    current_page = line["page"]
                current_text.append(s)
                current_bboxes.append(line["bbox"])
                length += len(s)

        flush()
        return chunks

    def generate_embeddings(self, chunks: List[Dict]) -> np.ndarray:
        return self.embedding_model.encode(
            [c["text"] for c in chunks],
            batch_size=32,
            convert_to_numpy=True,
            normalize_embeddings=True,   # must match vectordbs.py (IndexFlatIP)
            show_progress_bar=True
        )

    def add_to_database(
        self,
        embeddings: np.ndarray,
        chunks: List[Dict],
        filename: str
    ) -> int:
        start_idx = self.faiss_index.ntotal
        self.faiss_index.add(embeddings.astype("float32"))

        for i, chunk in enumerate(chunks):
            meta = {
                "chunk_id": start_idx + i,
                "text": chunk["text"],
                "filename": filename,
                "upload_date": datetime.now().isoformat(),
                "source": "user_upload",
                "page": chunk.get("page"),
                "line_bboxes": chunk.get("line_bboxes", []),
            }
            self.chunks.append(meta)
            self.chunk_id_to_idx[f"{filename}_{i}"] = start_idx + i

        return len(embeddings)

    def process_and_add_pdf(self, pdf_path: str, progress_callback=None) -> Dict:
        start = datetime.now()
        filename = Path(pdf_path).stem

        extracted = self.extract_text_from_pdf(pdf_path, progress_callback=progress_callback)
        full_text = extracted["full_text"]
        (self.ocr_dir / f"{filename}.txt").write_text(full_text, encoding="utf-8")

        if progress_callback:
            progress_callback("embedding", {})
        chunks = self.chunk_text(extracted["lines"])
        embeddings = self.generate_embeddings(chunks)

        np.save(self.embeddings_dir / f"{filename}_embeddings.npy", embeddings)

        if progress_callback:
            progress_callback("indexing", {})
        vectors_added = self.add_to_database(embeddings, chunks, filename)
        self.save_database()

        if progress_callback:
            progress_callback("done", {"num_chunks": len(chunks), "vectors_added": vectors_added})

        return {
            "filename": filename,
            "text_length": len(full_text),
            "num_chunks": len(chunks),
            "vectors_added": vectors_added,
            "total_vectors": self.faiss_index.ntotal,
            "processing_time_seconds": (datetime.now() - start).total_seconds(),
            "timestamp": datetime.now().isoformat()
        }


def main():
    vector_db_path = "/usr/users/3d_dimension_est/selva_sur/RAG/output/biomedbert_vector_db"

    updater = DynamicRAGUpdater(
        vector_db_path=vector_db_path,
        embedding_model="microsoft/BiomedNLP-BiomedBERT-base-uncased-abstract-fulltext",
        upload_dir="uploaded_reports"
    )

    test_pdf = "path/to/new_report.pdf"

    if Path(test_pdf).exists():
        stats = updater.process_and_add_pdf(test_pdf)
        print(json.dumps(stats, indent=2))
    else:
        print("Test PDF not found. Update the path in main().")


if __name__ == "__main__":
    main()