feat: add pdf-inspector as an optional PDF parser backend

pageindex/config.yaml

@@ -10,4 +10,8 @@ max_token_num_each_node: 20000
 if_add_node_id: "yes"
 if_add_node_summary: "yes"
 if_add_doc_description: "no"
-if_add_node_text: "no"
+if_add_node_text: "no"
+# PDF text extractor. One of: PyPDF2 (default), PyMuPDF, pdf_inspector.
+# pdf_inspector requires `pip install pdf-inspector` and preserves headings,
+# lists and tables from text-based PDFs.
+pdf_parser: "PyPDF2"

pageindex/page_index.py

@@ -1240,7 +1240,7 @@ def page_index_main(doc, opt=None):
         raise ValueError("Unsupported input type. Expected a PDF file path or BytesIO object.")
 
     print('Parsing PDF...')
-    page_list = get_page_tokens(doc, model=opt.model)
+    page_list = get_page_tokens(doc, model=opt.model, pdf_parser=getattr(opt, 'pdf_parser', 'PyPDF2'))
 
     logger.info({'total_page_number': len(page_list)})
     logger.info({'total_token': sum([page[1] for page in page_list])})

pageindex/utils.py

@@ -20,6 +20,57 @@
 # litellm is imported inside the functions that use it; eager import is slow
 # and fetches a remote model-cost map.
 
+# pdf-inspector is an optional backend imported lazily so its absence does not
+# break users on the default PyPDF2 path.
+SUPPORTED_PDF_PARSERS = ("PyPDF2", "PyMuPDF", "pdf_inspector")
+
+
+def _load_pdf_inspector():
+    try:
+        import pdf_inspector  # noqa: F401
+        return pdf_inspector
+    except ImportError as e:
+        raise ImportError(
+            "pdf_parser='pdf_inspector' requires the 'pdf-inspector' package. "
+            "Install it with `pip install pdf-inspector`."
+        ) from e
+
+
+def _pdf_inspector_pages(pdf_path, pages=None):
+    """Return the raw list of PageMarkdown items from pdf-inspector.
+
+    Accepts a filesystem path or a BytesIO buffer. `pages` is an optional list
+    of 0-indexed page numbers matching pdf-inspector's convention.
+    """
+    pi = _load_pdf_inspector()
+    if isinstance(pdf_path, BytesIO):
+        buf = pdf_path.getvalue()
+        return pi.extract_pages_markdown_bytes(buf, pages=pages).pages
+    return pi.extract_pages_markdown(pdf_path, pages=pages).pages
+
+
+def classify_pdf(pdf_path):
+    """Classify a PDF as text_based / scanned / image_based / mixed.
+
+    Returns a dict with `pdf_type`, `page_count`, `pages_needing_ocr`
+    (1-indexed), or None if pdf-inspector is not installed. Useful as a
+    preflight before running the full PageIndex pipeline.
+    """
+    try:
+        pi = _load_pdf_inspector()
+    except ImportError:
+        return None
+    if isinstance(pdf_path, BytesIO):
+        result = pi.classify_pdf_bytes(pdf_path.getvalue())
+    else:
+        result = pi.classify_pdf(pdf_path)
+    return {
+        "pdf_type": result.pdf_type,
+        "page_count": result.page_count,
+        # pdf-inspector reports 0-indexed pages; PageIndex uses 1-indexed.
+        "pages_needing_ocr": [p + 1 for p in getattr(result, "pages_needing_ocr", []) or []],
+    }
+
 # Backward compatibility: support CHATGPT_API_KEY as alias for OPENAI_API_KEY
 if not os.getenv("OPENAI_API_KEY") and os.getenv("CHATGPT_API_KEY"):
     os.environ["OPENAI_API_KEY"] = os.getenv("CHATGPT_API_KEY")
@@ -276,22 +327,62 @@ def get_last_node(structure):
     return structure[-1]
 
 
-def extract_text_from_pdf(pdf_path):
+def extract_text_from_pdf(pdf_path, pdf_parser="PyPDF2"):
+    if pdf_parser == "pdf_inspector":
+        pages = _pdf_inspector_pages(pdf_path)
+        return "".join((p.markdown or "") for p in pages)
+    if pdf_parser == "PyMuPDF":
+        if isinstance(pdf_path, BytesIO):
+            doc = pymupdf.open(stream=pdf_path, filetype="pdf")
+        else:
+            doc = pymupdf.open(pdf_path)
+        return "".join(page.get_text() for page in doc)
+    if pdf_parser != "PyPDF2":
+        raise ValueError(f"Unsupported PDF parser: {pdf_parser}")
     pdf_reader = PyPDF2.PdfReader(pdf_path)
-    ###return text not list 
+    ###return text not list
     text=""
     for page_num in range(len(pdf_reader.pages)):
         page = pdf_reader.pages[page_num]
         text+=page.extract_text()
     return text
 
-def get_pdf_title(pdf_path):
+def get_pdf_title(pdf_path, pdf_parser="PyPDF2"):
+    # Title lives in the PDF metadata dictionary regardless of the text-parser
+    # backend; PyPDF2 handles that cheaply so we keep it as the single path.
     pdf_reader = PyPDF2.PdfReader(pdf_path)
     meta = pdf_reader.metadata
     title = meta.title if meta and meta.title else 'Untitled'
     return title
 
-def get_text_of_pages(pdf_path, start_page, end_page, tag=True):
+def get_text_of_pages(pdf_path, start_page, end_page, tag=True, pdf_parser="PyPDF2"):
+    if pdf_parser == "pdf_inspector":
+        # pdf-inspector uses 0-indexed page numbers.
+        wanted = list(range(start_page - 1, end_page))
+        pages = _pdf_inspector_pages(pdf_path, pages=wanted)
+        text = ""
+        for page_num, page in zip(wanted, pages):
+            page_text = page.markdown or ""
+            if tag:
+                text += f"<start_index_{page_num+1}>\n{page_text}\n<end_index_{page_num+1}>\n"
+            else:
+                text += page_text
+        return text
+    if pdf_parser == "PyMuPDF":
+        if isinstance(pdf_path, BytesIO):
+            doc = pymupdf.open(stream=pdf_path, filetype="pdf")
+        else:
+            doc = pymupdf.open(pdf_path)
+        text = ""
+        for page_num in range(start_page - 1, end_page):
+            page_text = doc[page_num].get_text()
+            if tag:
+                text += f"<start_index_{page_num+1}>\n{page_text}\n<end_index_{page_num+1}>\n"
+            else:
+                text += page_text
+        return text
+    if pdf_parser != "PyPDF2":
+        raise ValueError(f"Unsupported PDF parser: {pdf_parser}")
     pdf_reader = PyPDF2.PdfReader(pdf_path)
     text = ""
     for page_num in range(start_page-1, end_page):
@@ -465,6 +556,17 @@ def get_page_tokens(pdf_path, model=None, pdf_parser="PyPDF2"):
             token_length = litellm.token_counter(model=model, text=page_text)
             page_list.append((page_text, token_length))
         return page_list
+    elif pdf_parser == "pdf_inspector":
+        # pdf-inspector emits per-page Markdown with heading tiers, list
+        # markers and GFM tables preserved. That structure carries into the
+        # downstream LLM prompt and node summaries.
+        pages = _pdf_inspector_pages(pdf_path)
+        page_list = []
+        for page in pages:
+            page_text = page.markdown or ""
+            token_length = litellm.token_counter(model=model, text=page_text)
+            page_list.append((page_text, token_length))
+        return page_list
     else:
         raise ValueError(f"Unsupported PDF parser: {pdf_parser}")
 
@@ -482,7 +584,16 @@ def get_text_of_pdf_pages_with_labels(pdf_pages, start_page, end_page):
         text += f"<physical_index_{page_num+1}>\n{pdf_pages[page_num][0]}\n<physical_index_{page_num+1}>\n"
     return text
 
-def get_number_of_pages(pdf_path):
+def get_number_of_pages(pdf_path, pdf_parser="PyPDF2"):
+    if pdf_parser == "pdf_inspector":
+        try:
+            pi = _load_pdf_inspector()
+        except ImportError:
+            pass
+        else:
+            if isinstance(pdf_path, BytesIO):
+                return pi.classify_pdf_bytes(pdf_path.getvalue()).page_count
+            return pi.classify_pdf(pdf_path).page_count
     pdf_reader = PyPDF2.PdfReader(pdf_path)
     num = len(pdf_reader.pages)
     return num

run_pageindex.py

@@ -39,7 +39,17 @@
                       help='Whether to add doc description to the doc')
     parser.add_argument('--if-add-node-text', type=str, default=None,
                       help='Whether to add text to the node')
-
+
+    parser.add_argument('--pdf-parser', type=str, default=None,
+                      choices=['PyPDF2', 'PyMuPDF', 'pdf_inspector'],
+                      help='PDF text extractor backend (PDF only, non-flash). '
+                           '`pdf_inspector` preserves headings, lists and tables '
+                           'and requires `pip install pdf-inspector`.')
+    parser.add_argument('--check-ocr', action='store_true',
+                      help='Classify the PDF with pdf-inspector before parsing '
+                           'and refuse image-based files that would need OCR '
+                           '(requires `pip install pdf-inspector`).')
+
     # Markdown specific arguments
     parser.add_argument('--if-thinning', type=str, default='no',
                       help='Whether to apply tree thinning for markdown (markdown only)')
@@ -67,8 +77,27 @@
             raise ValueError("PDF file must have .pdf extension")
         if not os.path.isfile(args.pdf_path):
             raise ValueError(f"PDF file not found: {args.pdf_path}")
-
+
+        if args.check_ocr:
+            from pageindex.utils import classify_pdf
+            info = classify_pdf(args.pdf_path)
+            if info is None:
+                raise SystemExit(
+                    "--check-ocr requires the `pdf-inspector` package. "
+                    "Install it with `pip install pdf-inspector`."
+                )
+            print(f"PDF type: {info['pdf_type']} ({info['page_count']} pages)")
+            if info['pdf_type'] in ('scanned', 'image_based'):
+                raise SystemExit(
+                    f"Refusing to parse: this PDF is {info['pdf_type']} and "
+                    "needs OCR. PageIndex does not run OCR locally."
+                )
+            if info['pages_needing_ocr']:
+                print(f"Warning: pages needing OCR: {info['pages_needing_ocr']}")
+
         if args.flash:
+            if args.pdf_parser is not None:
+                raise ValueError("--pdf-parser is not supported with --flash")
             from pageindex.flash import page_index_flash
             if args.optimize == 'full':
                 from pageindex.tree_optimize import default_model
@@ -102,6 +131,7 @@
                 'if_add_node_summary': args.if_add_node_summary,
                 'if_add_doc_description': args.if_add_doc_description,
                 'if_add_node_text': args.if_add_node_text,
+                'pdf_parser': args.pdf_parser,
             }
             opt = ConfigLoader().load({k: v for k, v in user_opt.items() if v is not None})
             toc_with_page_number = page_index_main(args.pdf_path, opt)

tests/test_pdf_parser_backend.py

@@ -0,0 +1,130 @@
+import os
+import sys
+import unittest
+from unittest.mock import MagicMock, patch
+
+sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))
+
+from pageindex import utils
+from pageindex.utils import (
+    SUPPORTED_PDF_PARSERS,
+    classify_pdf,
+    extract_text_from_pdf,
+    get_number_of_pages,
+    get_page_tokens,
+    get_text_of_pages,
+)
+
+
+FIXTURE_PDF = os.path.join(
+    os.path.dirname(__file__),
+    "..",
+    "examples",
+    "documents",
+    "q1-fy25-earnings.pdf",
+)
+
+
+def _has_pdf_inspector():
+    try:
+        import pdf_inspector  # noqa: F401
+        return True
+    except ImportError:
+        return False
+
+
+REQUIRES_PDF_INSPECTOR = unittest.skipUnless(
+    _has_pdf_inspector(), "pdf-inspector not installed"
+)
+
+
+class BackendDispatchTest(unittest.TestCase):
+    def test_supported_parsers_advertised(self):
+        self.assertIn("PyPDF2", SUPPORTED_PDF_PARSERS)
+        self.assertIn("PyMuPDF", SUPPORTED_PDF_PARSERS)
+        self.assertIn("pdf_inspector", SUPPORTED_PDF_PARSERS)
+
+    def test_get_page_tokens_rejects_unknown_parser(self):
+        with self.assertRaises(ValueError):
+            get_page_tokens("dummy.pdf", model="gpt-4o", pdf_parser="not_a_parser")
+
+    def test_extract_text_from_pdf_rejects_unknown_parser(self):
+        with self.assertRaises(ValueError):
+            extract_text_from_pdf("dummy.pdf", pdf_parser="not_a_parser")
+
+    def test_get_text_of_pages_rejects_unknown_parser(self):
+        with self.assertRaises(ValueError):
+            get_text_of_pages("dummy.pdf", 1, 1, pdf_parser="not_a_parser")
+
+    def test_load_pdf_inspector_raises_friendly_error_when_missing(self):
+        # Simulate the package being absent even when it is installed on the
+        # dev machine, so this test runs uniformly in CI.
+        with patch.dict(sys.modules, {"pdf_inspector": None}):
+            with self.assertRaises(ImportError) as ctx:
+                utils._load_pdf_inspector()
+            self.assertIn("pdf-inspector", str(ctx.exception))
+
+    def test_classify_pdf_returns_none_when_pdf_inspector_missing(self):
+        with patch.dict(sys.modules, {"pdf_inspector": None}):
+            self.assertIsNone(classify_pdf("dummy.pdf"))
+
+    def test_get_number_of_pages_falls_back_when_pdf_inspector_missing(self):
+        # Backend requested but unavailable: fall back to PyPDF2 rather than
+        # raise, so opt.pdf_parser stays a soft preference.
+        fake_reader = MagicMock()
+        fake_reader.pages = [object(), object(), object()]
+        with patch.dict(sys.modules, {"pdf_inspector": None}), \
+             patch("pageindex.utils.PyPDF2.PdfReader", return_value=fake_reader):
+            self.assertEqual(
+                get_number_of_pages("dummy.pdf", pdf_parser="pdf_inspector"),
+                3,
+            )
+
+
+@REQUIRES_PDF_INSPECTOR
+class PdfInspectorBackendTest(unittest.TestCase):
+    @classmethod
+    def setUpClass(cls):
+        if not os.path.isfile(FIXTURE_PDF):
+            raise unittest.SkipTest(f"fixture PDF missing: {FIXTURE_PDF}")
+
+    def test_classify_pdf_returns_expected_fields(self):
+        info = classify_pdf(FIXTURE_PDF)
+        self.assertIsNotNone(info)
+        self.assertIn("pdf_type", info)
+        self.assertIn("page_count", info)
+        self.assertIn("pages_needing_ocr", info)
+        self.assertGreater(info["page_count"], 0)
+        self.assertIn(info["pdf_type"], {"text_based", "mixed", "scanned", "image_based"})
+
+    def test_get_number_of_pages_matches_pypdf2(self):
+        expected = get_number_of_pages(FIXTURE_PDF, pdf_parser="PyPDF2")
+        actual = get_number_of_pages(FIXTURE_PDF, pdf_parser="pdf_inspector")
+        self.assertEqual(actual, expected)
+
+    def test_get_page_tokens_returns_all_pages(self):
+        pages = get_page_tokens(FIXTURE_PDF, model="gpt-4o", pdf_parser="pdf_inspector")
+        self.assertGreater(len(pages), 0)
+        expected = get_number_of_pages(FIXTURE_PDF, pdf_parser="PyPDF2")
+        self.assertEqual(len(pages), expected)
+        # Every entry must be (text, token_count).
+        for text, tokens in pages:
+            self.assertIsInstance(text, str)
+            self.assertIsInstance(tokens, int)
+
+    def test_get_text_of_pages_tags_selected_range(self):
+        text = get_text_of_pages(
+            FIXTURE_PDF, 3, 4, tag=True, pdf_parser="pdf_inspector"
+        )
+        self.assertIn("<start_index_3>", text)
+        self.assertIn("<start_index_4>", text)
+        self.assertIn("<end_index_3>", text)
+        self.assertIn("<end_index_4>", text)
+
+    def test_extract_text_from_pdf_returns_nonempty(self):
+        text = extract_text_from_pdf(FIXTURE_PDF, pdf_parser="pdf_inspector")
+        self.assertGreater(len(text), 100)
+
+
+if __name__ == "__main__":
+    unittest.main()