phase: 126_pdf_support
All gates verified green — no code changes were needed on this pass. Final report: ## Phase 126 — PDF support: final verification pass **Verified (all tasks 01–05 already complete in working tree):** - `pymupdf` dep in `pyproject.toml`/`uv.lock`; `.env.example` documents the `pdf` token - `app/rag/pdf.py` (text + image extraction, MIME mapping), page-aware `_chunk_pdf`, importer `_index_pdf_file` + embedded-image pass, archives covered via extension walk **Test / lint / coverage results:** - `uv run pytest` → 2997 passed - `uv run pytest tests/e2e/test_pdf_support.py -v --no-cov` (isolation) → 6 passed - `uv run pytest --cov=app --cov-report=term-missing` → TOTAL **99%** (>90% gate), 2997 passed - `uv run ruff check . && uv run pyright` → clean (0 errors) **Completion criteria:** 1. PDF → doc with text/chunks/embedding/summary — ✅ (integration + E2E Sources-page tests) 2. `BOR_IMAGES=true`: embedded images → standalone image docs — ✅ (integration + E2E chat-chip test) 3. `BOR_IMAGES=false`: images skipped, text indexed — ✅ (both suites) 4. PDF renders as extracted text in viewer — ✅ (E2E) 5. Test/coverage/lint gates — ✅ (above) 6. Commit + phase move — harness handles (per protocol, no `git` on this pass) **Decisions:** none new; inline `[Embedded image: …]` placeholders were a documented task-04 deviation (simplified spec wins) — not a defect. No defects found; repo left functional with changes in the working tree. **Next pending phase:** none — `todo/` is empty once 126 is moved.
This commit is contained in:
+210
@@ -0,0 +1,210 @@
|
||||
"""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
|
||||
Reference in New Issue
Block a user