fix: restore PyPDF2 as the PDF text extractor in the SDK path

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.
This commit is contained in:
Ray 2026-07-19 15:47:47 +08:00
parent b2b0901579
commit f2b407f62f
2 changed files with 63 additions and 70 deletions

View file

@ -205,7 +205,7 @@ Omitting `doc_ids` queries the **entire collection** and lets the agent pick whi
# ⚙️ Package Usage # ⚙️ 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. You can follow these steps to generate a PageIndex tree from a PDF document.

View file

@ -1,3 +1,4 @@
import PyPDF2
import pymupdf import pymupdf
from pathlib import Path from pathlib import Path
from .protocol import ContentNode, ParsedDocument from .protocol import ContentNode, ParsedDocument
@ -15,39 +16,42 @@ class PdfParser:
path = Path(file_path) path = Path(file_path)
model = kwargs.get("model") model = kwargs.get("model")
images_dir = kwargs.get("images_dir") 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 = [] 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: tokens = count_tokens(content, model=model)
page_count = doc.page_count nodes.append(ContentNode(
for i, page in enumerate(doc): content=content,
page_num = i + 1 tokens=tokens,
if images_dir: index=page_num,
content, images = self._extract_page_with_images( images=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, return ParsedDocument(doc_name=path.stem, nodes=nodes,
metadata={"page_count": page_count}) metadata={"page_count": len(reader.pages)})
@staticmethod @staticmethod
def _extract_page_with_images(doc, page, page_num: int, def _extract_page_images(page, page_num: int, images_dir: str) -> list[dict]:
images_dir: str) -> tuple[str, list[dict]]: """Save a page's images to disk and return their metadata."""
"""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)
"""
images_path = Path(images_dir) images_path = Path(images_dir)
images_path.mkdir(parents=True, exist_ok=True) images_path.mkdir(parents=True, exist_ok=True)
# Store an absolute path so the ![image](...) reference resolves # 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.) # break as soon as the query runs from a different directory.)
abs_images_path = images_path.resolve() abs_images_path = images_path.resolve()
parts: list[str] = []
images: list[dict] = [] images: list[dict] = []
img_idx = 0 img_idx = 0
for block in page.get_text("dict")["blocks"]: for block in page.get_text("dict")["blocks"]:
if block["type"] == 0: # text block if block["type"] != 1: # image blocks only
lines = [] continue
for line in block["lines"]: width = block.get("width", 0)
spans_text = "".join(span["text"] for span in line["spans"]) height = block.get("height", 0)
lines.append(spans_text) if width < _MIN_IMAGE_SIZE or height < _MIN_IMAGE_SIZE:
parts.append("\n".join(lines)) continue
elif block["type"] == 1: # image block image_bytes = block.get("image")
width = block.get("width", 0) if not image_bytes:
height = block.get("height", 0) continue
if width < _MIN_IMAGE_SIZE or height < _MIN_IMAGE_SIZE:
continue
image_bytes = block.get("image") try:
if not image_bytes: pix = pymupdf.Pixmap(image_bytes)
continue # 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: images.append({
pix = pymupdf.Pixmap(image_bytes) "path": str(abs_images_path / filename),
# n includes the alpha channel, so a plain RGBA pixmap also "width": width,
# has n==4 — subtract alpha before comparing. Without this, "height": height,
# a CMYK image with no alpha (n==4, same as RGBA) skips the })
# RGB conversion, and pix.save() as .png then raises img_idx += 1
# "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) return images
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