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_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.

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
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 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 ""

    parts: list[str] = []
    with doc:
        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
            # get_text("text") always returns a str; the union return type
            # comes from pymupdf's other flavors (dict/words/blocks).
            text = raw.strip() if isinstance(raw, str) else ""
            if text:
                parts.append(f"{PAGE_SEPARATOR.format(page=page_num + 1)}\n{text}")
    return "\n\n".join(parts)


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