brain-of-reese/app/rag/pdf.py

"""PDF parsing utilities (phase 126 — PDF support, locked decision A1).

Pure functions (I/O = reading the PDF file only) for extracting text
and embedded images from PDF files using PyMuPDF (the ``pymupdf``
package, imported as ``fitz`` — its stable API name).

* :func:`extract_pdf_text` — the text of every page that has text,
  each prefixed with :data:`PAGE_SEPARATOR` (``---PAGE-N---``, N is the
  1-indexed page number). The PDF chunker (phase 126 task 03) splits on
  these markers so every chunk keeps its page anchor. Pages without
  text (image-only pages) contribute nothing, so a PDF with no
  extractable text returns ``""``.
* :func:`extract_pdf_page_texts` — one entry per page that has text:
  1-indexed page number → the page's stripped text (the same page read
  as :func:`extract_pdf_text`, so the two agree page-for-page; the
  value for page N is exactly that page's separator-block body). The
  phase-137 image-context block is built from this map (page N's entry
  is the document context for an image embedded on page N).
* :func:`extract_pdf_images` — one :class:`PdfImage` per embedded-image
  occurrence, in page order then 0-based index on the page. ``data`` is
  the image's original encoded bytes as stored in the PDF (no
  re-encoding), ``mime`` is derived via :func:`_mime_from_image_info`.
* :func:`_mime_from_image_info` — maps a PyMuPDF image-info dict to a
  MIME type: the embedded file extension (``ext``) first, then the
  image type code (``type``: 2=JPEG, 3=PNG, 4=GIF, 8=TIFF, 9=WebP),
  else the ``application/octet-stream`` fallback.
* :func:`is_usable_pdf_text` — the phase-133 garbage gate: a pure,
  deterministic predicate over the extracted text layer (no PDF, no
  I/O) — REAL content vs. garbage (the pre-126 XML/binary soup, the
  empty text layer). It is the single trigger for the page-screenshot
  vision fallback (phase 133 task 03).
* :func:`render_pdf_pages` — the fallback's visual input (phase 133
  task 03): one PNG per page (150 DPI, page order); an unreadable page
  yields ``b""`` (the caller's placeholder — the page index never
  shifts), an unopenable file yields ``[]`` (the caller's existing
  empty-content path).

Every failure mode is fail-soft: an unreadable file returns an empty
result, an unreadable page or image is logged and skipped — this module
never raises for malformed input (the importer owns the loud failures).
"""
from __future__ import annotations

import logging
import re
from collections.abc import Iterator
from dataclasses import dataclass
from pathlib import Path
from typing import Any

import pymupdf as fitz  # PyMuPDF (locked A1); `fitz` is the stable API name

logger = logging.getLogger("app.rag.pdf")

#: Page separator used in extracted text (the chunker splits on this).
#: ``{page}`` is the 1-indexed page number.
PAGE_SEPARATOR = "---PAGE-{page}---"

#: Mapping from PyMuPDF image type codes to MIME types
#: (2=JPEG, 3=PNG, 4=GIF, 8=TIFF, 9=WebP — phase-126 design).
_TYPE_TO_MIME: dict[int, str] = {
    2: "image/jpeg",
    3: "image/png",
    4: "image/gif",
    8: "image/tiff",
    9: "image/webp",
}

#: Mapping from embedded-image file extensions to MIME types.
_EXT_TO_MIME: dict[str, str] = {
    "png": "image/png",
    "jpg": "image/jpeg",
    "jpeg": "image/jpeg",
    "gif": "image/gif",
    "webp": "image/webp",
    "tif": "image/tiff",
    "tiff": "image/tiff",
}


@dataclass
class PdfImage:
    """One embedded image from a PDF page."""

    page: int  #: 1-indexed page number
    index: int  #: image index on the page (0-based, occurrence order)
    width: int  #: native pixel width
    height: int  #: native pixel height
    data: bytes  #: original encoded bytes as stored in the PDF
    mime: str  #: MIME type (e.g. ``image/png``)


def _mime_from_image_info(info: dict[str, Any]) -> str:
    """Map a PyMuPDF image-info dict to a MIME type.

    The embedded file extension (``ext``) wins over the image type code
    (``type``); unknown/missing both fall back to
    ``application/octet-stream``.
    """
    ext = str(info.get("ext") or "").lower()
    if ext in _EXT_TO_MIME:
        return _EXT_TO_MIME[ext]
    type_code = info.get("type")
    if isinstance(type_code, int):
        return _TYPE_TO_MIME.get(type_code, "application/octet-stream")
    return "application/octet-stream"


def _page_image_entries(page: Any) -> list[dict[str, Any]]:
    """One info dict per embedded-image occurrence on *page*.

    Prefers ``get_image_info(xrefs=True)`` (one entry per occurrence,
    with the ``xref`` object number); older PyMuPDF releases without the
    ``xrefs`` keyword are served from ``get_images(full=True)`` (one
    entry per unique image).
    """
    try:
        infos: list[dict[str, Any]] = page.get_image_info(xrefs=True)
    except TypeError:
        infos = [
            {"xref": entry[0], "width": entry[2], "height": entry[3]}
            for entry in page.get_images(full=True)
        ]
    return [dict(info) for info in infos]


def _iter_page_texts(doc: Any, pdf_path: Path) -> Iterator[tuple[int, str]]:
    """Yield ``(page, text)`` for every page of *doc* that has text.

    ``page`` is the 1-indexed page number and ``text`` the page's
    stripped text (the ``get_text("text")`` read — it always returns a
    str; the union return type comes from pymupdf's other flavors,
    dict/words/blocks). A page that cannot be read is logged and
    skipped; a page without text contributes nothing. The shared page
    iteration behind :func:`extract_pdf_text` and
    :func:`extract_pdf_page_texts` — the two callers share the exact
    same read and cannot drift.
    """
    for page_num in range(len(doc)):
        try:
            raw = doc[page_num].get_text("text")
        except Exception as exc:
            logger.warning(
                "pdf: could not read text of page %d of %s — %s",
                page_num + 1, pdf_path, exc,
            )
            continue
        text = raw.strip() if isinstance(raw, str) else ""
        if text:
            yield page_num + 1, text


def extract_pdf_text(pdf_path: Path) -> str:
    """Extract the text of all pages of *pdf_path*.

    Every page that has text is prefixed with :data:`PAGE_SEPARATOR`
    (``---PAGE-N---``, N 1-indexed) and the page blocks are joined with
    a blank line. Pages without text contribute nothing, so a PDF with
    no extractable text (or a zero-page file) returns ``""``.

    Fail-soft: an unreadable file returns ``""``; a page that cannot be
    read is logged and skipped.
    """
    try:
        doc = fitz.open(str(pdf_path))
    except Exception as exc:
        logger.warning("pdf: could not open %s — %s", pdf_path, exc)
        return ""

    with doc:
        parts = [
            f"{PAGE_SEPARATOR.format(page=page)}\n{text}"
            for page, text in _iter_page_texts(doc, pdf_path)
        ]
    return "\n\n".join(parts)


def extract_pdf_page_texts(pdf_path: Path) -> dict[int, str]:
    """One entry per page of *pdf_path* that has text.

    Keys are the 1-indexed page numbers, values the page's stripped
    text — the shared :func:`_iter_page_texts` read, so this map agrees
    with :func:`extract_pdf_text` page-for-page: the value for page N
    is exactly the text after that page's ``---PAGE-N---`` separator.
    Pages without text (image-only pages) contribute nothing — no key.

    The phase-137 document context for an embedded image is built from
    this map: one image on page N gets page N's text as the context
    block's body.

    Fail-soft: an unreadable file returns ``{}`` (one warning — the
    :func:`extract_pdf_text` open contract); a page that cannot be
    read is logged and skipped.
    """
    try:
        doc = fitz.open(str(pdf_path))
    except Exception as exc:
        logger.warning("pdf: could not open %s for page texts — %s", pdf_path, exc)
        return {}

    with doc:
        return dict(_iter_page_texts(doc, pdf_path))


#: A line "looks like XML" when it carries a tag-shaped token — the
#: pre-126 garbage's signature (raw XMP blocks, PDF-structure noise).
_XMLISH_LINE_RE = re.compile(r"</?[A-Za-z_][\w:.-]*[\s/>=]")

#: Garbage-gate thresholds (phase 133 task 02): an extracted text
#: layer is UNUSABLE when its XML-ish line fraction reaches this bar.
#: The FINDING's reproduced pre-126 read of the owner's file is 46%
#: XML-ish; a real document ABOUT XML that reaches the bar is a
#: documented false positive (ASSUMPTION A8 — the fallback still
#: extracts it via vision; the cost is a few vision calls).
XMLISH_LINE_FRACTION = 0.30

#: Garbage-gate thresholds (phase 133 task 02): the non-printable
#: fraction (U+FFFD replacement chars + C0 controls over the whole
#: text) at which the layer is UNUSABLE — the ``errors="replace"``
#: binary-soup signal (the pre-126 read is ~19% non-printable; real
#: extracted text is ~0.05%).
NONPRINT_CHAR_FRACTION = 0.05


def is_usable_pdf_text(text: str) -> bool:
    """Whether an extracted PDF text layer is REAL content (not garbage).

    Unusable (→ the page-screenshot fallback) when ANY of:

    * ``text.strip() == ""`` — the empty text layer (a scanned/image-
      only PDF has nothing to read);
    * ≥ :data:`XMLISH_LINE_FRACTION` (0.30) of the non-empty lines
      match :data:`_XMLISH_LINE_RE` — the raw-XMP / binary-soup
      signature (the FINDING's reproduced pre-126 read is 46% —
      ``read_text(errors="replace")`` on the owner's file left
      ``<illustrator:StartupProfile>Web</illustrator:StartupProfile>``
      lines in the preview);
    * ≥ :data:`NONPRINT_CHAR_FRACTION` (0.05) of all characters are
      non-printable — the ``errors="replace"`` binary-soup signal.
      The non-printable set is ``U+FFFD`` + C0 control chars
      (0x00–0x1F) EXCEPT the structural ``\\t``/``\\n``/``\\r`` the
      extractor legitimately emits (the explicit set is deliberate: a
      ``unicodedata`` Cc-category check would catch those too, and
      counting them would flag every prose document).

    Pure + deterministic (no PDF, no I/O) — the unit tests pin the
    thresholds against the REAL reproduced garbage string, not a
    hand-drawn one. A real document ABOUT XML whose tag-like lines
    reach 30% is a documented false positive (ASSUMPTION A8 — the
    fallback still extracts it correctly via vision; the cost is a
    few vision calls).
    """
    if not text.strip():
        return False
    non_empty_lines = [line for line in text.splitlines() if line.strip()]
    xml_ish = sum(1 for line in non_empty_lines if _XMLISH_LINE_RE.search(line))
    if xml_ish / len(non_empty_lines) >= XMLISH_LINE_FRACTION:
        return False
    non_printable = sum(
        1
        for ch in text
        if ch == "\ufffd" or (ord(ch) < 32 and ch not in "\t\n\r")
    )
    return non_printable / len(text) < NONPRINT_CHAR_FRACTION


def render_pdf_pages(pdf_path: Path, dpi: int = 150) -> list[bytes]:
    """Render every page of *pdf_path* to a PNG (page order — list index
    *i* is page *i*+1).

    The fallback's visual input (phase 133 task 03, ASSUMPTION A8:
    fixed 150 DPI — the balance of vision legibility vs data-URL
    payload). A page that cannot be rendered yields ``b""`` (the
    caller's placeholder — the page index NEVER shifts, so
    ``---PAGE-N---`` stays truthful; the failure is logged, the
    :func:`extract_pdf_text` fail-soft style). A file that cannot be
    OPENED yields ``[]`` (the caller's existing empty-content path —
    the doc lands with the no-extractable-text behavior, unchanged).
    """
    try:
        doc = fitz.open(str(pdf_path))
    except Exception as exc:
        logger.warning("pdf: could not open %s for page rendering — %s", pdf_path, exc)
        return []

    pages: list[bytes] = []
    with doc:
        for page_num in range(len(doc)):
            try:
                pages.append(doc[page_num].get_pixmap(dpi=dpi).tobytes("png"))
            except Exception as exc:
                logger.warning(
                    "pdf: could not render page %d of %s — %s",
                    page_num + 1, pdf_path, exc,
                )
                pages.append(b"")
    return pages


def extract_pdf_images(pdf_path: Path) -> list[PdfImage]:
    """Extract all embedded images from *pdf_path*.

    Returns one :class:`PdfImage` per image occurrence, in page order
    then 0-based index on the page. ``data`` is the original encoded
    bytes as stored in the PDF; ``mime`` is derived from the image's
    extension/type via :func:`_mime_from_image_info`.

    Fail-soft: an unreadable file returns ``[]``; an image without an
    ``xref`` or that cannot be extracted is logged and skipped.
    """
    try:
        doc = fitz.open(str(pdf_path))
    except Exception as exc:
        logger.warning("pdf: could not open %s for image extraction — %s", pdf_path, exc)
        return []

    images: list[PdfImage] = []
    with doc:
        for page_num in range(len(doc)):
            try:
                entries = _page_image_entries(doc[page_num])
            except Exception as exc:
                logger.warning(
                    "pdf: could not list images on page %d of %s — %s",
                    page_num + 1, pdf_path, exc,
                )
                continue
            for index, info in enumerate(entries):
                xref = int(info.get("xref") or 0)
                if not xref:
                    logger.warning(
                        "pdf: skipping image page=%d index=%d of %s — no xref",
                        page_num + 1, index, pdf_path,
                    )
                    continue
                try:
                    extracted = doc.extract_image(xref)
                except Exception as exc:
                    logger.warning(
                        "pdf: could not extract image page=%d index=%d of %s — %s",
                        page_num + 1, index, pdf_path, exc,
                    )
                    continue
                data = extracted.get("image") or b""
                if not data:
                    logger.warning(
                        "pdf: skipping image page=%d index=%d of %s — empty data",
                        page_num + 1, index, pdf_path,
                    )
                    continue
                # The image info's own `ext` is unreliable in some
                # PyMuPDF versions (None); `extract_image` reports the
                # actual format of the returned bytes, so it wins.
                mime = _mime_from_image_info({**info, "ext": extracted.get("ext") or ""})
                images.append(
                    PdfImage(
                        page=page_num + 1,
                        index=index,
                        width=int(extracted.get("width") or info.get("width") or 0),
                        height=int(extracted.get("height") or info.get("height") or 0),
                        data=data,
                        mime=mime,
                    )
                )
    return images