pdf_processor.py

"""
PDF ingestion with optional OCR fallback.

Phase 4 improvement: OCR is now triggered not only by text length but also by
a word-quality heuristic.  A quality score is stored in metadata so low-quality
OCR chunks can be identified and down-weighted during retrieval if needed.
"""
import os
import re
import logging

from langchain_community.document_loaders import PyPDFLoader
from langchain_core.documents import Document

from config.settings import (
    RAW_PDF_PATH,
    OCR_LANGUAGE,
    TESSERACT_CMD,
    PDF_OCR_ENABLED,
    PDF_OCR_MIN_CHARS,
)
from processors.source_metadata import ensure_source_metadata

logger = logging.getLogger(__name__)

try:
    import pypdfium2 as pdfium
except Exception:
    pdfium = None  # type: ignore

try:
    import pytesseract
except Exception:
    pytesseract = None  # type: ignore


# ---------------------------------------------------------------------------
# Text quality helpers
# ---------------------------------------------------------------------------

def _normalize_text(text: str) -> str:
    return " ".join((text or "").split())


def _compute_text_quality(text: str) -> float:
    """
    Rough quality score in [0, 1].
    Counts words that look like real English/numeric tokens vs total words.
    Values below ~0.3 suggest garbage OCR or heavily symbolic content.
    """
    if not text:
        return 0.0
    words = text.split()
    if not words:
        return 0.0
    good = sum(
        1 for w in words
        if re.match(r"^[a-zA-Z0-9][a-zA-Z0-9\.\,\:\-\/\']*$", w) and len(w) >= 2
    )
    return good / len(words)


def _needs_ocr(text: str) -> bool:
    """
    Return True when the extracted text is too short OR its quality score is
    low enough that OCR may produce better output.
    """
    if len(text) < PDF_OCR_MIN_CHARS:
        return True
    quality = _compute_text_quality(text)
    # Trigger OCR when less than 30 % of tokens look like real words
    return quality < 0.30


# ---------------------------------------------------------------------------
# OCR helper
# ---------------------------------------------------------------------------

def _ocr_pdf_page(filepath: str, page_index: int) -> str:
    if not PDF_OCR_ENABLED or pdfium is None or pytesseract is None:
        return ""
    try:
        if TESSERACT_CMD:
            pytesseract.pytesseract.tesseract_cmd = TESSERACT_CMD
        pdf = pdfium.PdfDocument(filepath)
        page = pdf[page_index]
        bitmap = page.render(scale=2.0)
        image = bitmap.to_pil()
        text = pytesseract.image_to_string(image, lang=OCR_LANGUAGE)
        page.close()
        pdf.close()
        return _normalize_text(text)
    except Exception as e:
        logger.debug("OCR failed for page %d of %s: %s", page_index, filepath, e)
        return ""


# ---------------------------------------------------------------------------
# Public processor
# ---------------------------------------------------------------------------

def process_pdfs() -> list[Document]:
    """Load PDF files as Document objects (not yet chunked)."""
    documents: list[Document] = []
    if not os.path.exists(RAW_PDF_PATH):
        print("PDF folder not found")
        return documents

    for filename in os.listdir(RAW_PDF_PATH):
        if not filename.lower().endswith(".pdf"):
            continue
        filepath = os.path.join(RAW_PDF_PATH, filename)
        print(f"Processing: {filename}")
        try:
            loader = PyPDFLoader(filepath)
            docs = loader.load()
            for i, doc in enumerate(docs):
                extracted = _normalize_text(doc.page_content)
                used_ocr = False
                ocr_quality: float | None = None

                if PDF_OCR_ENABLED and _needs_ocr(extracted):
                    ocr_text = _ocr_pdf_page(filepath, i)
                    if len(ocr_text) > len(extracted):
                        ocr_quality = _compute_text_quality(ocr_text)
                        extracted = ocr_text
                        used_ocr = True

                if not extracted:
                    continue

                meta = ensure_source_metadata(doc.metadata, source_type="pdf")
                if used_ocr:
                    meta["ocr_fallback"] = True
                    if ocr_quality is not None:
                        meta["ocr_quality"] = round(ocr_quality, 3)

                documents.append(Document(page_content=extracted, metadata=meta))
        except Exception as e:
            logger.warning("Skipping unreadable PDF: %s (%s)", filename, e)
            print(f"[WARN] Skipping unreadable PDF: {filename} ({e})")
            continue

    print(f"PDF loading complete: {len(documents)} pages")
    return documents


if __name__ == "__main__":
    docs = process_pdfs()