Validate PDF parser CLI option

pageindex/utils.py

@@ -391,9 +391,9 @@ SUPPORTED_PDF_PARSERS = ("PyPDF2", "pypdfium2", "PyMuPDF")
 DEFAULT_PDF_PARSER = os.getenv("PAGEINDEX_PDF_PARSER") or SUPPORTED_PDF_PARSERS[0]
 
 
-def read_pdf_pages(doc):
-    """Return a list of per-page text strings using the currently configured parser."""
-    parser = DEFAULT_PDF_PARSER
+def read_pdf_pages(doc, pdf_parser=None):
+    """Return a list of per-page text strings using the selected parser."""
+    parser = pdf_parser or DEFAULT_PDF_PARSER
 
     if parser == "PyPDF2":
         reader = PyPDF2.PdfReader(doc)
@@ -435,15 +435,14 @@ def read_pdf_pages(doc):
             d.close()
 
     raise ValueError(
-        f"Unsupported DEFAULT_PDF_PARSER={parser!r}. Choose from {SUPPORTED_PDF_PARSERS}."
+        f"Unsupported pdf_parser={parser!r}. Choose from {SUPPORTED_PDF_PARSERS}."
     )
 
 
-def get_page_tokens(pdf_path, model=None):
-    pages = read_pdf_pages(pdf_path)
+def get_page_tokens(pdf_path, model=None, pdf_parser=None):
+    pages = read_pdf_pages(pdf_path, pdf_parser=pdf_parser)
     return [(text, litellm.token_counter(model=model, text=text)) for text in pages]
 
-        
 
 def get_text_of_pdf_pages(pdf_pages, start_page, end_page):
     text = ""

requirements.txt

@@ -1,7 +1,7 @@
 litellm==1.83.7
 # openai-agents  # optional: required for examples/agentic_vectorless_rag_demo.py
 pymupdf==1.26.4
-# pypdfium2     # optional: enables pdf_parser="pypdfium2" (cleaner text, faster, Apache 2.0)
+# pypdfium2     # optional: alternative PDF parser
 PyPDF2==3.0.1
 python-dotenv==1.2.2
 pyyaml==6.0.2

run_pageindex.py

@@ -29,7 +29,8 @@ if __name__ == "__main__":
     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,
-                      help='PDF text extractor: PyPDF2 (default), pypdfium2 (requires `pip install pypdfium2`), or PyMuPDF')
+                      choices=pageindex_utils.SUPPORTED_PDF_PARSERS,
+                      help='PDF text extractor to use')
 
     # Markdown specific arguments
     parser.add_argument('--if-thinning', type=str, default='no',
@@ -136,4 +137,4 @@ if __name__ == "__main__":
         with open(output_file, 'w', encoding='utf-8') as f:
             json.dump(toc_with_page_number, f, indent=2, ensure_ascii=False)
         
-        print(f'Tree structure saved to: {output_file}')
+        print(f'Tree structure saved to: {output_file}')