diff --git a/README.md b/README.md
index 5ce0ca5e6..27e3084a4 100644
--- a/README.md
+++ b/README.md
@@ -173,9 +173,10 @@ python3 run_pageindex.py --pdf_path /path/to/your/document.pdf
Optional parameters
-You can customize the processing with additional optional arguments:
+You can customize the processing with additional optional arguments (the structure-tuning flags below require --mode standard):
```
+--mode Processing mode: flash (default) or standard
--model LLM model to use (default: gpt-4o-2024-11-20)
--toc-check-pages Pages to check for table of contents (default: 20)
--max-pages-per-node Max pages per node (default: 10)
@@ -199,13 +200,13 @@ python3 run_pageindex.py --md_path /path/to/your/document.md
> ### ⚡ PageIndex Flash *(preview)*
-> **PageIndex Flash** ([`pageindex/flash`](pageindex/flash)) generates tree structures from PDFs in seconds. Structure extraction is purely heuristic-based, no LLM needed. LLM is only used to generate node summaries.
+> **PageIndex Flash** ([`pageindex/flash`](pageindex/flash)) generates tree structures from PDFs in seconds. Structure extraction is purely heuristic-based, no LLM needed. An LLM is used only for node summaries and the optimization's expansion pass.
>
> ```bash
-> python3 run_pageindex.py --flash --pdf_path /path/to/your/document.pdf
+> python3 run_pageindex.py --mode flash --pdf_path /path/to/your/document.pdf
> ```
>
-> Add `--optimize` to refine the tree structure for more efficient retrieval (with an LLM expansion pass).
+> Tree optimization for retrieval (a deterministic merge, then an LLM expansion pass) is on by default; pass `--optimize off` to disable.
## 🚀 Agentic Vectorless RAG: An Example
diff --git a/pageindex/client.py b/pageindex/client.py
index eebd4e7bc..009b7cf41 100644
--- a/pageindex/client.py
+++ b/pageindex/client.py
@@ -149,19 +149,19 @@ def submit_document(
``wait=True`` to block until the document is ready, or poll
``get_document(doc_id)['status']`` yourself.
- Local: indexes the document in this call (it blocks while your LLM
- builds the tree — minutes for a standard index of a long document),
- then stores it under ``storage_path``. Pass ``mode="flash"`` to build
- the tree with PageIndex Flash (layout-based extraction, no LLM calls
- for the structure; node summaries and the document description still
- use ``summary_model``). ``beta_headers`` and ``folder_id`` are
+ Local: indexes the document in this call and stores it under
+ ``storage_path``. Defaults to Flash indexing: layout-based extraction,
+ refined for retrieval (a deterministic merge, then an LLM expansion
+ pass); node summaries, the expansion pass, and the document
+ description use ``summary_model``. Pass ``mode="standard"`` for a
+ full LLM-built tree (slower). ``beta_headers`` and ``folder_id`` are
cloud-only.
Args:
file_path (str): Path to the PDF file.
- mode (str, optional): Processing mode. Local mode supports
- "standard" and "flash"; omit it for standard indexing. Cloud
- modes are passed through (e.g. "mcp").
+ mode (str, optional): Processing mode. Local defaults to "flash";
+ pass "standard" for a full LLM-built tree. Cloud modes are
+ passed through (e.g. "mcp").
beta_headers (list[str], optional): Cloud-only beta feature headers.
folder_id (str, optional): Cloud-only folder (workspace) ID.
metadata (dict, optional): Your own JSON-serializable tags for the
diff --git a/pageindex/flash/README.md b/pageindex/flash/README.md
index 99d236181..0d23a3c35 100644
--- a/pageindex/flash/README.md
+++ b/pageindex/flash/README.md
@@ -11,9 +11,9 @@ an LLM.
```python
from pageindex.flash import page_index_flash
-tree = page_index_flash("paper.pdf")
-tree = page_index_flash("paper.pdf", summary=False) # tree structure only, no LLM
-tree = page_index_flash("paper.pdf", optimize=True) # refined tree for retrieval
+tree = page_index_flash("paper.pdf") # optimized tree + summaries
+tree = page_index_flash("paper.pdf", summary=False, optimize=False) # raw tree only, no LLM
+tree = page_index_flash("paper.pdf", optimize="merge") # deterministic merge, no LLM expand
```
Takes a file path or an `io.BytesIO` stream and returns the tree as a dict.
@@ -22,12 +22,11 @@ Summaries are on by default and need an LLM API key.
### Command line
```bash
-python3 run_pageindex.py --pdf_path document.pdf --flash
-python3 run_pageindex.py --pdf_path document.pdf --flash --no-summary # tree structure only, no LLM
-python3 run_pageindex.py --pdf_path document.pdf --flash --optimize # refined tree for retrieval
+python3 run_pageindex.py --mode flash --pdf_path document.pdf # optimized tree + summaries
+python3 run_pageindex.py --mode flash --pdf_path document.pdf --no-summary --optimize off # raw tree only, no LLM
```
-Writes the tree to `results/_structure_flash.json`.
+Writes the tree to `results/_structure.json`.
## Output
diff --git a/pageindex/flash/api.py b/pageindex/flash/api.py
index 74656d162..bf62d9657 100644
--- a/pageindex/flash/api.py
+++ b/pageindex/flash/api.py
@@ -96,15 +96,24 @@ def _optimize(structure, page_texts, do_expand, model):
def page_index_flash(pdf, summary=True, summary_model=None,
- optimize=False, optimize_expand=True,
+ optimize: str | bool = "full", optimize_expand=None,
optimize_model=None, summary_concurrency=None,
use_embedded_toc=True) -> dict:
- """Build a PageIndex tree structure from a PDF using layout statistics, without an LLM. Args: pdf: path to a PDF file (``str`` or ``pathlib.Path``) or an in-memory binary stream (``io.BytesIO``). summary: if True, generate LLM summaries for each node (requires ``summary_model``). summary_model: the LLM model identifier to use for summary generation. optimize: if True, refine the tree for search cost before summaries: a deterministic merge collapses subtrees whose structure does not beat a linear scan, keeping the removed titles on the parent as ``key_items``, then an LLM pass expands oversized sections. Without it the extracted tree is returned unchanged. optimize_expand: if False, run the merge but skip the LLM expansion. optimize_model: the LLM model for expand (defaults to the summary model). summary_concurrency: maximum simultaneous summary model calls; None uses the library default. use_embedded_toc: if True, consume the PDF's embedded bookmarks when trustworthy: deep bookmarks become the frame and the detected sections they lack are grafted back in after noise filtering, coarse ones become the chapter frame with detected nodes re-hung under them (deeper sparse entries are filled in when the page text confirms them, and garbled extracted titles are repaired from the bookmark strings), garbage ones are ignored; adds a ``toc_source`` key to the result. On by default; pass False for the pure detected structure. Returns: dict with keys ``doc_name``, ``doc_title``, ``structure`` (a list of nested ``{"title", "start_index", "end_index", "nodes"}`` dicts; page indexes are 1-based) and ``has_abstract_or_references_section`` (True when a top-level entry is an abstract or references heading). With ``optimize`` an ``optimize`` key reports merge/expand counts and before/after search-cost metrics. """
+ """Build a PageIndex tree structure from a PDF using layout statistics, without an LLM. Args: pdf: path to a PDF file (``str`` or ``pathlib.Path``) or an in-memory binary stream (``io.BytesIO``). summary: if True, generate LLM summaries for each node (requires ``summary_model``). summary_model: the LLM model identifier to use for summary generation. optimize: ``"full"`` for merge + LLM expand, ``"merge"`` for deterministic merge only, ``False`` to disable. ``True`` is accepted as ``"full"`` for backward compatibility. optimize_model: the LLM model for expand (defaults to the summary model). summary_concurrency: maximum simultaneous summary model calls; None uses the library default. use_embedded_toc: if True, consume the PDF's embedded bookmarks when trustworthy: deep bookmarks become the frame and the detected sections they lack are grafted back in after noise filtering, coarse ones become the chapter frame with detected nodes re-hung under them (deeper sparse entries are filled in when the page text confirms them, and garbled extracted titles are repaired from the bookmark strings), garbage ones are ignored; adds a ``toc_source`` key to the result. On by default; pass False for the pure detected structure. Returns: dict with keys ``doc_name``, ``doc_title``, ``structure`` (a list of nested ``{"title", "start_index", "end_index", "nodes"}`` dicts; page indexes are 1-based) and ``has_abstract_or_references_section`` (True when a top-level entry is an abstract or references heading). With ``optimize`` an ``optimize`` key reports merge/expand counts and before/after search-cost metrics. """
+ if optimize is True:
+ optimize = "full"
+ if not optimize:
+ optimize = False
+ elif optimize not in ("full", "merge"):
+ raise ValueError(
+ f"optimize must be 'full', 'merge', or False, got {optimize!r}")
+ if optimize_expand is not None and optimize:
+ optimize = "full" if optimize_expand else "merge"
result = extract_toc(_validate_pdf(pdf), use_embedded_toc=use_embedded_toc)
structure = result.get("structure", [])
if optimize and structure:
result["optimize"] = _optimize(structure, result.get("page_texts") or [],
- optimize_expand,
+ optimize == "full",
optimize_model or summary_model)
if summary and structure:
import asyncio
diff --git a/pageindex/local_api.py b/pageindex/local_api.py
index 8b1e6f184..9a82c48f9 100644
--- a/pageindex/local_api.py
+++ b/pageindex/local_api.py
@@ -75,8 +75,10 @@ def submit_document(
if mode not in (None, "standard", "flash"):
raise PageIndexAPIError(
f"Failed to submit document: unknown local processing mode {mode!r}. "
- "Supported: None or 'standard' for standard indexing, or 'flash'."
+ "Supported: 'flash' (default) or 'standard'."
)
+ if mode is None:
+ mode = "flash"
file_path = os.path.abspath(os.path.expanduser(str(file_path)))
if not os.path.isfile(file_path):
raise FileNotFoundError(f"No such file: {file_path}")
@@ -123,7 +125,7 @@ def submit_document(
"pageNum": len(page_texts),
"folderId": None,
"metadata": metadata,
- "mode": mode or "standard",
+ "mode": mode,
}
pages = [{"page_index": i + 1, "markdown": text}
for i, text in enumerate(page_texts)]
@@ -180,8 +182,16 @@ def _index_flash(self, file_path: str, page_texts: list[str]) -> tuple[list, str
from .flash import page_index_flash
from .utils import (add_node_text, create_clean_structure_for_description,
generate_doc_description, write_node_id)
+ import litellm
+ env = litellm.validate_environment(self._summary_model)
+ if not env["keys_in_environment"]:
+ raise PageIndexAPIError(
+ f"Failed to submit document: missing API key for "
+ f"{self._summary_model}: {', '.join(env['missing_keys'])}")
result = page_index_flash(file_path, summary=True,
- summary_model=self._summary_model)
+ summary_model=self._summary_model,
+ optimize="full",
+ optimize_model=self._summary_model)
structure = result.get("structure", [])
if not structure:
raise PageIndexAPIError(
diff --git a/pageindex/local_chat.py b/pageindex/local_chat.py
index 3dddcf18c..df890ebb6 100644
--- a/pageindex/local_chat.py
+++ b/pageindex/local_chat.py
@@ -120,6 +120,10 @@ def _run_sync(coro):
try:
asyncio.get_running_loop()
except RuntimeError:
+ has_loop = False
+ else:
+ has_loop = True
+ if not has_loop:
return asyncio.run(coro)
with concurrent.futures.ThreadPoolExecutor(max_workers=1) as pool:
return pool.submit(asyncio.run, coro).result()
diff --git a/run_pageindex.py b/run_pageindex.py
index 452f08174..80c01f16f 100644
--- a/run_pageindex.py
+++ b/run_pageindex.py
@@ -10,15 +10,18 @@
parser = argparse.ArgumentParser(description='Process PDF or Markdown document and generate structure')
parser.add_argument('--pdf_path', type=str, help='Path to the PDF file')
parser.add_argument('--md_path', type=str, help='Path to the Markdown file')
- parser.add_argument('--flash', action='store_true', help='Use PageIndex Flash (with --pdf_path)')
+ parser.add_argument('--mode', choices=['flash', 'standard'], default='flash',
+ help='Processing mode (default: flash)')
+ parser.add_argument('--flash', action='store_true', default=False,
+ help=argparse.SUPPRESS)
parser.add_argument('--embedded-toc', action=argparse.BooleanOptionalAction, default=None,
- help='Use the PDF\'s embedded bookmarks when trustworthy (default: on with --flash)')
+ help='Use the PDF\'s embedded bookmarks when trustworthy (default: on in flash mode)')
parser.add_argument('--summary', action=argparse.BooleanOptionalAction, default=None,
- help='Generate node summaries with an LLM (default: on with --flash)')
- parser.add_argument('--optimize', nargs='?', const='full', choices=['full', 'merge'],
+ help='Generate node summaries with an LLM (default: on in flash mode)')
+ parser.add_argument('--optimize', nargs='?', const='full', choices=['full', 'merge', 'off'],
default=None,
- help='Refine the tree for search cost: a deterministic merge, then an '
- 'LLM expansion pass; pass `merge` to run the merge alone (PDF only)')
+ help='Refine the tree for search cost (default: full in flash mode). '
+ '`merge` for deterministic merge only; `off` to disable')
parser.add_argument('--model', type=str, default=None, help='Model to use (overrides config.yaml)')
parser.add_argument('--summary-model', type=str, default=None,
@@ -48,18 +51,32 @@
parser.add_argument('--summary-token-threshold', type=int, default=200,
help='Token threshold for generating summaries (markdown only)')
args = parser.parse_args()
-
+ if args.flash:
+ args.mode = 'flash'
+
# Validate that exactly one file type is specified
if not args.pdf_path and not args.md_path:
raise ValueError("Either --pdf_path or --md_path must be specified")
if args.pdf_path and args.md_path:
raise ValueError("Only one of --pdf_path or --md_path can be specified")
- if args.optimize and not (args.pdf_path and args.flash):
- raise ValueError("--optimize requires --flash with --pdf_path")
- if args.embedded_toc is not None and not (args.pdf_path and args.flash):
- raise ValueError("--embedded-toc requires --flash with --pdf_path")
- if args.summary is not None and not (args.pdf_path and args.flash):
- raise ValueError("--summary requires --flash with --pdf_path")
+ if args.optimize in ('full', 'merge') and not (args.pdf_path and args.mode == 'flash'):
+ raise ValueError("--optimize requires Flash mode with --pdf_path")
+ if args.optimize is None:
+ args.optimize = 'full' if args.mode == 'flash' else 'off'
+ if args.embedded_toc is not None and not (args.pdf_path and args.mode == 'flash'):
+ raise ValueError("--embedded-toc requires Flash mode with --pdf_path")
+ if args.summary is not None and not (args.pdf_path and args.mode == 'flash'):
+ raise ValueError("--summary requires Flash mode with --pdf_path")
+ if args.pdf_path and args.mode == 'flash':
+ for flag, value in (('--toc-check-pages', args.toc_check_pages),
+ ('--max-pages-per-node', args.max_pages_per_node),
+ ('--max-tokens-per-node', args.max_tokens_per_node),
+ ('--if-add-node-id', args.if_add_node_id),
+ ('--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)):
+ if value is not None:
+ raise ValueError(f"{flag} is not supported in flash mode; use --mode standard")
if args.pdf_path:
# Validate PDF file
@@ -68,22 +85,23 @@
if not os.path.isfile(args.pdf_path):
raise ValueError(f"PDF file not found: {args.pdf_path}")
- if args.flash:
+ if args.mode == 'flash':
from pageindex.flash import page_index_flash
- if args.optimize == 'full':
- from pageindex.tree_optimize import default_model
- from pageindex.utils import _is_openai_model
- expand_model = args.model or default_model()
- if _is_openai_model(expand_model) and not os.getenv("OPENAI_API_KEY"):
- raise SystemExit(f"OPENAI_API_KEY is not set (expand model: {expand_model}).")
+ summary_model = args.summary_model or args.model
+ will_summarize = args.summary if args.summary is not None else True
+ if summary_model and (will_summarize or args.optimize == 'full'):
+ import litellm
+ env = litellm.validate_environment(summary_model)
+ if not env["keys_in_environment"]:
+ raise SystemExit(
+ f"Missing API key for {summary_model}: {', '.join(env['missing_keys'])}")
toc_with_page_number = page_index_flash(
args.pdf_path,
- optimize=args.optimize is not None,
- optimize_expand=args.optimize == 'full',
- optimize_model=args.model,
- summary_model=args.summary_model or args.model,
+ optimize=args.optimize if args.optimize != 'off' else False,
+ optimize_model=summary_model,
+ summary_model=summary_model,
use_embedded_toc=args.embedded_toc if args.embedded_toc is not None else True,
- summary=args.summary if args.summary is not None else True,
+ summary=will_summarize,
)
if 'optimize' in toc_with_page_number:
o = toc_with_page_number['optimize']
@@ -110,7 +128,7 @@
# Save results
pdf_name = os.path.splitext(os.path.basename(args.pdf_path))[0]
- suffix = '_structure_flash' if args.flash else '_structure'
+ suffix = '_structure'
output_dir = './results'
output_file = f'{output_dir}/{pdf_name}{suffix}.json'
os.makedirs(output_dir, exist_ok=True)
diff --git a/tests/test_client.py b/tests/test_client.py
index 60375e0df..aad4e3979 100644
--- a/tests/test_client.py
+++ b/tests/test_client.py
@@ -47,7 +47,7 @@ def fake_page_index_main(doc, opt=None, logger=None, page_list=None):
"doc_description": "A test document.",
"structure": json.loads(json.dumps(STRUCTURE))}
monkeypatch.setattr(page_index_module, "page_index_main", fake_page_index_main)
- return local_client.submit_document(sample_pdf)["doc_id"]
+ return local_client.submit_document(sample_pdf, mode="standard")["doc_id"]
# ── constructor ──
@@ -168,7 +168,7 @@ def fake_page_index_main(doc, opt=None, logger=None, page_list=None):
return {"doc_name": "sample.pdf", "doc_description": None,
"structure": json.loads(json.dumps(STRUCTURE))}
monkeypatch.setattr(page_index_module, "page_index_main", fake_page_index_main)
- local_client.submit_document(sample_pdf)
+ local_client.submit_document(sample_pdf, mode="standard")
assert not (tmp_path / "logs").exists()
@@ -179,10 +179,10 @@ def fake_page_index_main(doc, opt=None, logger=None, page_list=None):
return {"doc_name": "sample.pdf", "doc_description": "d",
"structure": json.loads(json.dumps(STRUCTURE))}
monkeypatch.setattr(page_index_module, "page_index_main", fake_page_index_main)
- first = local_client.submit_document(sample_pdf)
+ first = local_client.submit_document(sample_pdf, mode="standard")
assert first["name"] == "sample.pdf"
with pytest.warns(UserWarning, match='stored as "sample_1.pdf"'):
- second = local_client.submit_document(sample_pdf)
+ second = local_client.submit_document(sample_pdf, mode="standard")
assert second["name"] == "sample_1.pdf"
names = {d["id"]: d["name"]
for d in local_client.list_documents()["documents"]}
@@ -212,7 +212,7 @@ def test_submit_name_exhaustion_rejects_before_indexing(
"indexer ran despite name exhaustion"),
)
with pytest.raises(PageIndexAPIError, match="Too many files"):
- local_client.submit_document(sample_pdf)
+ local_client.submit_document(sample_pdf, mode="standard")
def test_submit_flash(local_client, sample_pdf, monkeypatch):
@@ -220,6 +220,8 @@ def test_submit_flash(local_client, sample_pdf, monkeypatch):
def fake_flash(pdf, summary=True, summary_model=None, **kwargs):
calls["summary"] = summary
calls["summary_model"] = summary_model
+ calls["optimize"] = kwargs.get("optimize")
+ calls["optimize_model"] = kwargs.get("optimize_model")
return {"doc_name": "sample.pdf",
"structure": [{"title": "Flash Root", "start_index": 1,
"end_index": 2, "summary": "s", "nodes": []}]}
@@ -227,13 +229,34 @@ def fake_flash(pdf, summary=True, summary_model=None, **kwargs):
monkeypatch.setattr(pageindex.utils, "llm_completion",
lambda model, prompt, **kw: "Flash description.")
doc_id = local_client.submit_document(sample_pdf, mode="flash")["doc_id"]
- assert calls == {"summary": True, "summary_model": local_client.summary_model}
+ assert calls == {"summary": True, "summary_model": local_client.summary_model,
+ "optimize": "full",
+ "optimize_model": local_client.summary_model}
root = local_client.get_tree(doc_id)["result"][0]
assert root["node_id"] == "0000"
assert "Hello page one" in root["text"]
assert local_client.get_document(doc_id)["description"] == "Flash description."
+def test_submit_defaults_to_flash(local_client, sample_pdf, monkeypatch):
+ monkeypatch.setattr(
+ pageindex.flash, "page_index_flash",
+ lambda pdf, **kwargs: {
+ "doc_name": "sample.pdf",
+ "structure": [{"title": "Flash Root", "start_index": 1,
+ "end_index": 2, "summary": "s", "nodes": []}]})
+ monkeypatch.setattr(pageindex.utils, "llm_completion",
+ lambda model, prompt, **kw: "Flash description.")
+ doc_id = local_client.submit_document(sample_pdf)["doc_id"]
+ assert local_client._api._store.get_meta(doc_id)["mode"] == "flash"
+
+
+def test_page_index_flash_rejects_unknown_optimize():
+ from pageindex.flash import page_index_flash
+ with pytest.raises(ValueError, match="optimize must be"):
+ page_index_flash("never-opened.pdf", optimize="off")
+
+
def test_llm_completion_missing_key_raises_immediately(monkeypatch):
import openai
monkeypatch.delenv("OPENAI_API_KEY", raising=False)
@@ -335,7 +358,7 @@ def test_submit_with_metadata(local_client, sample_pdf, monkeypatch):
"doc_name": "sample.pdf", "doc_description": None,
"structure": json.loads(json.dumps(STRUCTURE))})
tags = {"project": "alpha", "year": 2026}
- doc_id = local_client.submit_document(sample_pdf, metadata=tags)["doc_id"]
+ doc_id = local_client.submit_document(sample_pdf, mode="standard", metadata=tags)["doc_id"]
assert local_client.get_tree(doc_id)["metadata"] == tags
assert local_client.get_ocr(doc_id)["metadata"] == tags
assert local_client.list_documents()["documents"][0]["metadata"] == tags
@@ -486,7 +509,7 @@ def test_torn_delete_never_lists_ghost(local_client, indexed_doc, tmp_path):
def test_corrupt_doc_json_is_contained(local_client, indexed_doc, sample_pdf, tmp_path):
with pytest.warns(UserWarning): # same-name resubmit → stored as sample_1.pdf
- second = local_client.submit_document(sample_pdf)["doc_id"]
+ second = local_client.submit_document(sample_pdf, mode="standard")["doc_id"]
(tmp_path / "store" / "docs" / indexed_doc / "doc.json").write_text("{truncated")
# manifest still holds a good copy of the meta — served consistently