From f2b407f62f17b81f77060fa7fdc838c2bcf4cd89 Mon Sep 17 00:00:00 2001 From: Ray Date: Sun, 19 Jul 2026 15:47:47 +0800 Subject: [PATCH] fix: restore PyPDF2 as the PDF text extractor in the SDK path MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Text extraction (and therefore tree building) now matches the CLI / pre-SDK default. Image extraction keeps PyMuPDF — the only part that needs it — with ![image](path) references appended per page. --- README.md | 2 +- pageindex/parser/pdf.py | 131 +++++++++++++++++++--------------------- 2 files changed, 63 insertions(+), 70 deletions(-) diff --git a/README.md b/README.md index e9410f1..5327af7 100644 --- a/README.md +++ b/README.md @@ -205,7 +205,7 @@ Omitting `doc_ids` queries the **entire collection** and lets the agent pick whi # ⚙️ Package Usage -> **Note:** This package uses standard PDF parsing. The CLI parses PDFs with PyPDF2 (as in previous releases), while the SDK's local mode parses with PyMuPDF and also extracts images — extracted text can differ slightly, so the two entry points may produce different trees for the same PDF. For use cases with complex PDFs, our [cloud service](https://pageindex.ai/developer) (via MCP and API) offers enhanced OCR, tree building, and retrieval. +> **Note:** This package uses standard PDF parsing. For use cases with complex PDFs, our [cloud service](https://pageindex.ai/developer) (via MCP and API) offers enhanced OCR, tree building, and retrieval. You can follow these steps to generate a PageIndex tree from a PDF document. diff --git a/pageindex/parser/pdf.py b/pageindex/parser/pdf.py index 971b5d9..b95506b 100644 --- a/pageindex/parser/pdf.py +++ b/pageindex/parser/pdf.py @@ -1,3 +1,4 @@ +import PyPDF2 import pymupdf from pathlib import Path from .protocol import ContentNode, ParsedDocument @@ -15,39 +16,42 @@ class PdfParser: path = Path(file_path) model = kwargs.get("model") images_dir = kwargs.get("images_dir") + + # Images are extracted with PyMuPDF (PyPDF2 cannot); text stays with + # PyPDF2 below so the extracted text — and therefore the tree — matches + # the CLI / pre-SDK default. + page_images: dict[int, list[dict]] = {} + if images_dir: + with pymupdf.open(str(path)) as doc: + for i, page in enumerate(doc): + images = self._extract_page_images(page, i + 1, images_dir) + if images: + page_images[i + 1] = images + + reader = PyPDF2.PdfReader(str(path)) nodes = [] + for i, page in enumerate(reader.pages): + page_num = i + 1 + content = page.extract_text() or "" + images = page_images.get(page_num) + if images: + refs = "\n".join(f"![image]({img['path']})" for img in images) + content = f"{content}\n{refs}" if content else refs - with pymupdf.open(str(path)) as doc: - page_count = doc.page_count - 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, - )) + tokens = count_tokens(content, model=model) + nodes.append(ContentNode( + content=content, + tokens=tokens, + index=page_num, + images=images, + )) return ParsedDocument(doc_name=path.stem, nodes=nodes, - metadata={"page_count": page_count}) + metadata={"page_count": len(reader.pages)}) @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: ![image](path) - """ + def _extract_page_images(page, page_num: int, images_dir: str) -> list[dict]: + """Save a page's images to disk and return their metadata.""" images_path = Path(images_dir) images_path.mkdir(parents=True, exist_ok=True) # Store an absolute path so the ![image](...) reference resolves @@ -55,53 +59,42 @@ class PdfParser: # 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)) + if block["type"] != 1: # image blocks only + continue + width = block.get("width", 0) + height = block.get("height", 0) + if width < _MIN_IMAGE_SIZE or height < _MIN_IMAGE_SIZE: + continue - 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 - 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" + pix.save(str(images_path / filename)) + pix = None + except Exception: + 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 + images.append({ + "path": str(abs_images_path / filename), + "width": width, + "height": height, + }) + img_idx += 1 - img_path = str(abs_images_path / filename) - images.append({ - "path": img_path, - "width": width, - "height": height, - }) - parts.append(f"![image]({img_path})") - img_idx += 1 - - content = "\n".join(parts) - return content, images + return images