""" Layout-aware PDF extraction: tables and OCR. Why this exists --------------- `pypdf`'s page.extract_text() returns a flat character stream. For prose that is fine. For the documents this product actually ingests it is not: MSHA regulations, equipment maintenance schedules, and dust sampling reports are largely *tables*, and flattening a table produces runs like "Bolt spacing 5 4 Roof height 6 8 Support type Resin Mechanical" which is worse than useless — it is confidently wrong, and it gets embedded, retrieved, and cited as though it were meaningful. Scanned documents are the other half of the gap. An incident report that was photocopied and re-scanned yields *zero* extractable characters, so today it becomes an empty document that silently fails ingestion. This module adds two things to the PDF path: * table extraction via pdfplumber, rendered as Markdown so the structure survives chunking and is legible to the LLM * an OCR fallback for pages whose extractable text falls below a threshold Both degrade gracefully. If pdfplumber is missing, the caller falls back to pypdf. If Tesseract is not installed, OCR logs a warning once and is skipped, so local development without system binaries still works. Licensing note: PyMuPDF (fitz) is the more capable library here and is deliberately not used — it is AGPL-3.0, which is incompatible with this project's MIT licence. pdfplumber (MIT), pypdf (BSD), pytesseract (Apache-2.0) and pdf2image (MIT) are all compatible. """ from __future__ import annotations import logging from dataclasses import dataclass from typing import List, Optional from app.config import settings logger = logging.getLogger(__name__) # Warn once per process rather than once per page. _ocr_unavailable_logged = False @dataclass class PageExtraction: """Text for one page, plus how it was obtained.""" text: str tables_found: int = 0 used_ocr: bool = False # ── Table rendering ──────────────────────────────────────────────────────────── def table_to_markdown(table: List[List[Optional[str]]]) -> str: """ Render an extracted table as a Markdown table. Markdown is used rather than CSV or raw text because it keeps the header/row relationship explicit in a form LLMs read reliably, and it survives the chunker (which splits on sentence boundaries, not pipes). Ragged rows are padded, None cells become empty, and embedded pipes and newlines are escaped so a single cell cannot break the table structure. """ if not table: return "" def clean(cell: Optional[str]) -> str: if cell is None: return "" return str(cell).replace("|", "\\|").replace("\n", " ").strip() rows = [[clean(c) for c in row] for row in table] width = max(len(r) for r in rows) rows = [r + [""] * (width - len(r)) for r in rows] # Drop rows that are entirely empty — pdfplumber emits these for ruled # lines that are not really rows. rows = [r for r in rows if any(c for c in r)] if not rows: return "" header, body = rows[0], rows[1:] lines = [ "| " + " | ".join(header) + " |", "| " + " | ".join("---" for _ in header) + " |", ] lines.extend("| " + " | ".join(r) + " |" for r in body) return "\n".join(lines) # ── OCR ──────────────────────────────────────────────────────────────────────── def ocr_available() -> bool: """True if both pytesseract and the Tesseract binary are usable.""" try: import pytesseract pytesseract.get_tesseract_version() return True except Exception: return False def ocr_page(pdf_path: str, page_number: int) -> str: """ OCR a single 1-indexed page of a PDF. Rasterises just that page rather than the whole document, so a 500-page file with one scanned insert does not cost 500 page renders. """ global _ocr_unavailable_logged try: import pytesseract from pdf2image import convert_from_path except ImportError: if not _ocr_unavailable_logged: logger.warning( "OCR requested but pytesseract/pdf2image are not installed — " "scanned pages will yield no text" ) _ocr_unavailable_logged = True return "" try: images = convert_from_path( pdf_path, dpi=settings.OCR_DPI, first_page=page_number, last_page=page_number, ) if not images: return "" return pytesseract.image_to_string( images[0], lang=settings.OCR_LANGUAGE ).strip() except Exception as e: if not _ocr_unavailable_logged: logger.warning( f"OCR unavailable ({e}). Install the tesseract-ocr and " f"poppler-utils system packages to read scanned documents." ) _ocr_unavailable_logged = True return "" # ── Page extraction ──────────────────────────────────────────────────────────── def extract_page( plumber_page, pdf_path: str, page_number: int, ocr_budget: List[int], ) -> PageExtraction: """ Extract one page: prose, then tables, then OCR if it looks scanned. `ocr_budget` is a single-element list used as a mutable counter so the per-document OCR cap is shared across pages without a class. """ text = (plumber_page.extract_text() or "").strip() tables_found = 0 if settings.ENABLE_TABLE_EXTRACTION: try: tables = plumber_page.extract_tables() or [] except Exception as e: logger.debug(f"Table extraction failed on page {page_number}: {e}") tables = [] rendered = [md for t in tables if (md := table_to_markdown(t))] if rendered: tables_found = len(rendered) # Labelled so a retrieved chunk makes clear it is tabular data and # the model does not read the pipes as prose. text = ( f"{text}\n\n" + "\n\n".join( f"[Table {i} on page {page_number}]\n{md}" for i, md in enumerate(rendered, start=1) ) ).strip() used_ocr = False if settings.ENABLE_OCR and len(text) < settings.OCR_MIN_CHARS and ocr_budget[0] > 0: ocr_text = ocr_page(pdf_path, page_number) if ocr_text: ocr_budget[0] -= 1 used_ocr = True # Keep whatever little native text there was; OCR supplements it. text = f"{text}\n{ocr_text}".strip() if text else ocr_text return PageExtraction(text=text, tables_found=tables_found, used_ocr=used_ocr)