ingestion/pdf_reader.py

import logging
import pdfplumber
from pathlib import Path
from dataclasses import dataclass

logger = logging.getLogger(__name__)


@dataclass
class PageContent:
    page_number: int
    text: str
    char_count: int


@dataclass
class DocumentContent:
    file_path: str
    file_name: str
    total_pages: int
    pages: list[PageContent]

    @property
    def full_text(self) -> str:
        return "\n\n".join(p.text for p in self.pages if p.text)

    @property
    def total_chars(self) -> int:
        return sum(p.char_count for p in self.pages)


def extract_pdf(
    file_path: str | Path,
    parser: str = "pymupdf4llm",
) -> DocumentContent:
    """Extract text from a PDF, returning a DocumentContent with per-page text.

    Args:
        file_path: Path to the PDF file.
        parser:    "pymupdf4llm" (default) uses structured markdown extraction
                   that preserves headers, tables, and lists.  Falls back to
                   "pdfplumber" automatically if pymupdf4llm fails.
                   Pass "pdfplumber" explicitly to always use the flat extractor.
    """
    path = Path(file_path)
    if not path.exists():
        raise FileNotFoundError(f"PDF not found: {path}")
    if path.suffix.lower() != ".pdf":
        raise ValueError(f"Expected a .pdf file, got: {path.suffix}")

    if parser == "pymupdf4llm":
        try:
            return extract_pdf_structured(path)
        except Exception as exc:
            logger.warning(
                "pymupdf4llm extraction failed for %s (%s) — falling back to pdfplumber.",
                path.name, exc,
            )
            return _extract_pdf_pdfplumber(path)

    return _extract_pdf_pdfplumber(path)


def extract_pdf_structured(file_path: str | Path) -> DocumentContent:
    """Extract a PDF to per-page markdown using pymupdf4llm.

    pymupdf4llm preserves document structure as markdown:
    - Section headers become ## / ### headings
    - Tables become markdown tables
    - Bullet lists become markdown lists

    Args:
        file_path: Path to an existing .pdf file.

    Returns:
        DocumentContent whose page texts are markdown strings.
    """
    import pymupdf4llm  # lazy import — optional dependency

    path = Path(file_path)
    page_dicts: list[dict] = pymupdf4llm.to_markdown(str(path), page_chunks=True)

    pages: list[PageContent] = []
    for i, page_dict in enumerate(page_dicts, start=1):
        text = page_dict.get("text", "").strip()
        pages.append(PageContent(page_number=i, text=text, char_count=len(text)))

    return DocumentContent(
        file_path=str(path.resolve()),
        file_name=path.name,
        total_pages=len(page_dicts),
        pages=pages,
    )


def _extract_pdf_pdfplumber(path: Path) -> DocumentContent:
    """Extract flat text from a PDF using pdfplumber (original implementation)."""
    pages: list[PageContent] = []

    with pdfplumber.open(path) as pdf:
        total_pages = len(pdf.pages)
        for i, page in enumerate(pdf.pages, start=1):
            raw = page.extract_text() or ""
            text = _clean(raw)
            pages.append(PageContent(page_number=i, text=text, char_count=len(text)))

    return DocumentContent(
        file_path=str(path.resolve()),
        file_name=path.name,
        total_pages=total_pages,
        pages=pages,
    )


def _clean(text: str) -> str:
    lines = (line.strip() for line in text.splitlines())
    non_empty = (line for line in lines if line)
    return "\n".join(non_empty)