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)