mirror of
https://github.com/VectifyAI/PageIndex.git
synced 2026-07-15 21:11:05 +02:00
4e6a13576d
Addresses items 9-13 and a/b/c/f from the max-effort review of PR #272. - pdf.py: image colorspace check was `pix.n > 4`, which treats CMYK-without- alpha (n==4, same as RGBA) as not needing RGB conversion; pix.save() as .png then raises "unsupported colorspace", silently dropped by the surrounding except. Fixed to `pix.n - pix.alpha >= 4` (correctly converts CMYK, leaves RGBA untouched). - pipeline.py: detect_strategy([]) (an empty/whitespace-only source file) returned "content_based", routing into the PDF-oriented TOC-detection pipeline -- wasting a real LLM call before raising IndexingError. Empty node lists now route to level_based, whose build_tree_from_levels([]) returns an empty structure instantly with zero LLM calls. - page_index.py (shim): pageindex/__init__.py binds the canonical `page_index` function as the package attribute, but this file is ALSO a real submodule of the same name -- importing it anywhere (import machinery, unconditional) overwrites that attribute with the module object, breaking `from pageindex import page_index; page_index(x)` for the rest of the process. Made the shim module itself callable (delegates to the real function via a ModuleType subclass), so whichever object ends up in that slot is callable regardless of import order. - storage/sqlite.py: create_collection let a raw sqlite3.IntegrityError escape on a duplicate name (new CollectionAlreadyExistsError); the collections table's CHECK constraint only validated the name's first character (GLOB '*' is a wildcard, not a regex quantifier over the preceding class) -- fixed to validate the whole string, and SQLiteStorage now also validates in Python (it's a public StorageEngine usable directly, bypassing LocalBackend's own check). - tests/test_review_fixes_2.py: two tests used a ContentNode with no `level` set, so build_index took the content_based path and made real (retried, slow, and -- with a valid key -- billable) LLM calls instead of testing the text-stripping logic they claimed to. Mocked out _content_based_pipeline. - retrieve.py: _parse_pages/_get_pdf_page_content were independent copies of the canonical parse_pages/get_pdf_page_content that had already drifted (missing the p>=1 filter and 1000-page DoS cap) -- delegate to canonical now, so the legacy pageindex.get_page_content path can't silently regress again. - parser/markdown.py: a leading UTF-8 BOM broke first-header detection (not whitespace, .strip() doesn't remove it) -- decode utf-8-sig. Only backtick fences were recognized as code blocks, so a '#'-prefixed line inside a ~~~-fenced block (valid CommonMark) was misparsed as a heading -- recognize both fence styles. - run_pageindex.py: --if-thinning wasn't migrated to the bare-flag + legacy-yes/no convention the other four --if-add-* flags got; bare usage raised an argparse error and it never went through the shared coercion. - types.py: DocumentDetail's `structure` field was inside the class's total=False body, so TypedDict rules made it optional even though every backend always populates it. Split into a required base class. Adds regression tests for all of the above. Full suite: 244 passed, 2 skipped (one pre-existing, unrelated flaky cloud-streaming test). Claude-Session: https://claude.ai/code/session_01Kx5DgKbhK1N8autqXH8SmS
105 lines
4 KiB
Python
105 lines
4 KiB
Python
import pymupdf
|
|
from pathlib import Path
|
|
from .protocol import ContentNode, ParsedDocument
|
|
from ..tokens import count_tokens
|
|
|
|
# Minimum image dimension to keep (skip icons/artifacts)
|
|
_MIN_IMAGE_SIZE = 32
|
|
|
|
|
|
class PdfParser:
|
|
def supported_extensions(self) -> list[str]:
|
|
return [".pdf"]
|
|
|
|
def parse(self, file_path: str, **kwargs) -> ParsedDocument:
|
|
path = Path(file_path)
|
|
model = kwargs.get("model")
|
|
images_dir = kwargs.get("images_dir")
|
|
nodes = []
|
|
|
|
with pymupdf.open(str(path)) as doc:
|
|
for i, page in enumerate(doc):
|
|
page_num = i + 1
|
|
if images_dir:
|
|
content, images = self._extract_page_with_images(
|
|
doc, page, page_num, images_dir)
|
|
else:
|
|
content = page.get_text()
|
|
images = None
|
|
|
|
tokens = count_tokens(content, model=model)
|
|
nodes.append(ContentNode(
|
|
content=content or "",
|
|
tokens=tokens,
|
|
index=page_num,
|
|
images=images if images else None,
|
|
))
|
|
|
|
return ParsedDocument(doc_name=path.stem, nodes=nodes)
|
|
|
|
@staticmethod
|
|
def _extract_page_with_images(doc, page, page_num: int,
|
|
images_dir: str) -> tuple[str, list[dict]]:
|
|
"""Extract text and images from a page, preserving their relative order.
|
|
|
|
Uses get_text("dict") to iterate blocks in reading order.
|
|
Text blocks become text; image blocks are saved to disk and replaced
|
|
with an inline placeholder: 
|
|
"""
|
|
images_path = Path(images_dir)
|
|
images_path.mkdir(parents=True, exist_ok=True)
|
|
# Store an absolute path so the  reference resolves
|
|
# regardless of the process's cwd at query time. (cwd-relative paths
|
|
# break as soon as the query runs from a different directory.)
|
|
abs_images_path = images_path.resolve()
|
|
|
|
parts: list[str] = []
|
|
images: list[dict] = []
|
|
img_idx = 0
|
|
|
|
for block in page.get_text("dict")["blocks"]:
|
|
if block["type"] == 0: # text block
|
|
lines = []
|
|
for line in block["lines"]:
|
|
spans_text = "".join(span["text"] for span in line["spans"])
|
|
lines.append(spans_text)
|
|
parts.append("\n".join(lines))
|
|
|
|
elif block["type"] == 1: # image block
|
|
width = block.get("width", 0)
|
|
height = block.get("height", 0)
|
|
if width < _MIN_IMAGE_SIZE or height < _MIN_IMAGE_SIZE:
|
|
continue
|
|
|
|
image_bytes = block.get("image")
|
|
if not image_bytes:
|
|
continue
|
|
|
|
try:
|
|
pix = pymupdf.Pixmap(image_bytes)
|
|
# n includes the alpha channel, so a plain RGBA pixmap also
|
|
# has n==4 — subtract alpha before comparing. Without this,
|
|
# a CMYK image with no alpha (n==4, same as RGBA) skips the
|
|
# RGB conversion, and pix.save() as .png then raises
|
|
# "unsupported colorspace for 'png'", silently dropping the
|
|
# image via the bare except below.
|
|
if pix.n - pix.alpha >= 4:
|
|
pix = pymupdf.Pixmap(pymupdf.csRGB, pix)
|
|
filename = f"p{page_num}_img{img_idx}.png"
|
|
save_path = images_path / filename
|
|
pix.save(str(save_path))
|
|
pix = None
|
|
except Exception:
|
|
continue
|
|
|
|
img_path = str(abs_images_path / filename)
|
|
images.append({
|
|
"path": img_path,
|
|
"width": width,
|
|
"height": height,
|
|
})
|
|
parts.append(f"")
|
|
img_idx += 1
|
|
|
|
content = "\n".join(parts)
|
|
return content, images
|