diff --git a/pageindex/client.py b/pageindex/client.py
index 709b92d22..da326e459 100644
--- a/pageindex/client.py
+++ b/pageindex/client.py
@@ -58,7 +58,9 @@ def _agents_sdk_model_name(model: str) -> str:
return f"litellm/{model}"
-_LOCAL_INDEX_KEYS = ("model", "summary_model", "backend", "storage_path")
+_LOCAL_INDEX_KEYS = ("model", "summary_model", "backend", "storage_path",
+ "summary_max_words", "summary_concurrency",
+ "use_embedded_toc", "optimize")
# Near-synonyms of "cloud" that would otherwise parse as model names —
# a silent wrong mode. They error, pointing at the real word.
@@ -83,7 +85,8 @@ def _env_cloud_key(spelling: str, inline: str = "api_key=...") -> str:
# there as a PageIndexAPIError — never later, never silently.
_ARG_TYPES: "dict[str, tuple[type, ...]]" = {
"model": (str,), "index_model": (str,), "summary_model": (str,),
- "chat_model": (str,), "retrieve_model": (str,),
+ "chat_model": (str,), "retrieve_model": (str,), "summary_max_words": (int,),
+ "summary_concurrency": (int,), "use_embedded_toc": (bool,), "optimize": (str,),
"storage_path": (str, os.PathLike), "index_backend": (dict,),
"chat_backend": (dict,)}
@@ -128,8 +131,8 @@ def _resolve_index_slot(index) -> "tuple[_CloudKey, dict[str, Any]]":
'or "cloud".')
if isinstance(index, Mapping):
# None-valued keys mean "absent", exactly like the flat arguments.
- conf = {name: value for name, value in index.items()
- if value is not None}
+ conf: dict[str, Any] = {name: value for name, value in index.items()
+ if value is not None}
declared = _declared_mode(conf.pop("mode", None), "index")
if not conf:
if declared == "cloud":
@@ -167,12 +170,8 @@ def _resolve_index_slot(index) -> "tuple[_CloudKey, dict[str, Any]]":
'index declares mode "cloud" but carries local keys '
f"({', '.join(sorted(conf))}) — the cloud pipeline does "
'its own indexing; cloud takes "api_key" only.')
- mapped = {"index_model": conf.get("model"),
- "summary_model": conf.get("summary_model"),
- "index_backend": conf.get("backend"),
- "storage_path": conf.get("storage_path")}
- return None, {name: value for name, value in mapped.items()
- if value is not None}
+ rename = {"model": "index_model", "backend": "index_backend"}
+ return None, {rename.get(name, name): value for name, value in conf.items()}
raise PageIndexAPIError("index must be a string or a dict.")
@@ -264,7 +263,9 @@ class PageIndexClient:
``"cloud"`` / ``"pageindex-cloud"`` (cloud, key from the
environment), ``"local"``, a local index model name, or a
dict: ``{"api_key": ...}`` for cloud, ``{"model",
- "summary_model", "backend", "storage_path"}`` for local. An
+ "summary_model", "backend", "storage_path",
+ "summary_max_words", "summary_concurrency", "use_embedded_toc",
+ "optimize"}`` for local. An
optional ``"mode"`` key (``"cloud"`` / ``"local"``) states
the side and must agree with the other keys; ``{"mode":
"cloud"}`` alone reads the key from the environment. Not
@@ -306,6 +307,18 @@ class PageIndexClient:
summary_model (str, optional): Local mode only — legacy: overrides
the model used for node summaries and document descriptions;
``index_model`` covers this.
+ summary_max_words (int, optional): Local mode only — the word cap
+ each node summary is asked to stay within. Defaults to 150.
+ summary_concurrency (int, optional): Local mode only — cap on
+ simultaneous indexing model calls per lane: the summaries, and
+ expand up to its own ceiling of 32. The lanes overlap, so up to
+ cap + min(32, cap) calls run at once. Defaults to 64.
+ use_embedded_toc (bool, optional): Local mode only — whether flash
+ indexing consumes the PDF's embedded bookmarks when they look
+ trustworthy. Defaults to True.
+ optimize (str, optional): Local mode only — the flash tree
+ refinement pass: ``"full"`` (merge + model expand, the
+ default), ``"merge"`` (deterministic merge only) or ``"off"``.
retrieve_model (str, optional): Legacy name for ``chat_model`` —
same meaning everywhere, cloud clients included.
storage_path (str or os.PathLike, optional): Local mode only —
@@ -346,6 +359,10 @@ def __init__(
chat_model: Optional[str] = None,
model: Optional[str] = None,
summary_model: Optional[str] = None,
+ summary_max_words: Optional[int] = None,
+ summary_concurrency: Optional[int] = None,
+ use_embedded_toc: Optional[bool] = None,
+ optimize: Optional[str] = None,
retrieve_model: Optional[str] = None,
storage_path: Optional[Union[str, os.PathLike[str]]] = None,
index_backend: Optional[dict[str, Any]] = None,
@@ -363,6 +380,10 @@ def __init__(
(("api_key", api_key),
("index_model", index_model),
("summary_model", summary_model),
+ ("summary_max_words", summary_max_words),
+ ("summary_concurrency", summary_concurrency),
+ ("use_embedded_toc", use_embedded_toc),
+ ("optimize", optimize),
("index_backend", index_backend),
("storage_path", storage_path), ("model", model))
if value is not None}
@@ -443,10 +464,13 @@ def __init__(
f"got {type(value).__name__}.")
if isinstance(value, str):
value = conf[name] = value.strip()
- if not value:
+ if not value and not isinstance(value, bool):
raise PageIndexAPIError(
f"{shown} is empty — it configures nothing. Pass a "
"real value, or drop the argument.")
+ if name == "optimize" and value not in ("full", "merge", "off"):
+ raise PageIndexAPIError(
+ f'{shown} must be "full", "merge" or "off", got {value!r}.')
if cloud_key is not None:
if index_conf:
@@ -506,6 +530,10 @@ def __init__(
model=self.model,
summary_model=self.summary_model,
index_backend=index_conf.get("index_backend"),
+ summary_max_words=index_conf.get("summary_max_words"),
+ summary_concurrency=index_conf.get("summary_concurrency"),
+ use_embedded_toc=index_conf.get("use_embedded_toc", True),
+ optimize=index_conf.get("optimize", "full"),
)
# LiteLLM's multi-second import would otherwise land on the
# first chat call; failures resurface there with real context.
@@ -1651,6 +1679,10 @@ def __init__(
chat_model: Optional[str] = None,
model: Optional[str] = None,
summary_model: Optional[str] = None,
+ summary_max_words: Optional[int] = None,
+ summary_concurrency: Optional[int] = None,
+ use_embedded_toc: Optional[bool] = None,
+ optimize: Optional[str] = None,
retrieve_model: Optional[str] = None,
storage_path: Optional[Union[str, os.PathLike[str]]] = None,
index_backend: Optional[dict[str, Any]] = None,
@@ -1659,5 +1691,8 @@ def __init__(
super().__init__(None, index=index, chat=chat,
index_model=index_model, chat_model=chat_model,
model=model, summary_model=summary_model,
+ summary_max_words=summary_max_words,
+ summary_concurrency=summary_concurrency,
+ use_embedded_toc=use_embedded_toc, optimize=optimize,
retrieve_model=retrieve_model, storage_path=storage_path,
index_backend=index_backend, chat_backend=chat_backend)
diff --git a/pageindex/flash/api.py b/pageindex/flash/api.py
index 6f7361e27..1399b43b6 100644
--- a/pageindex/flash/api.py
+++ b/pageindex/flash/api.py
@@ -68,26 +68,28 @@ def _validate_pdf(pdf):
return pdf
-async def _summarize(structure, page_list, model, concurrency=None):
+async def _summarize(structure, page_list, model, concurrency=None, max_words=None):
from ..utils import summarize_tree
- await summarize_tree(structure, page_list, model=model, concurrency=concurrency)
+ await summarize_tree(structure, page_list, model=model, concurrency=concurrency,
+ max_words=max_words)
-def _optimize(structure, page_texts, do_expand, model):
- """Merge/expand refinement between extraction and summaries.
+async def _optimize_async(structure, page_texts, do_expand, model, on_final=None,
+ concurrency=None):
+ """Merge/expand refinement after extraction, overlapped with the summaries
+ when `on_final` is passed; without it the caller runs them after.
Beyond the merge the default path runs anyway, this adds LLM expand and
- reports before/after search-cost metrics. Summaries run after, so they
- describe the final tree. Expand reads the same page text the summaries use.
+ reports before/after search-cost metrics. Expand reads the same page text
+ the summaries use.
"""
- import asyncio
from ..tree_optimize import optimize
lines = [[line_text.strip() for line_text in (page_text or "").splitlines()
if line_text.strip()]
for page_text in page_texts]
- outcome = asyncio.run(optimize(structure, page_texts, lines, model=model,
- do_expand=do_expand,
- page_count=len(page_texts)))
+ outcome = await optimize(structure, page_texts, lines, model=model,
+ do_expand=do_expand, page_count=len(page_texts),
+ on_final=on_final, concurrency=concurrency)
return {"merges": outcome["merges"], "expands": outcome["expands"],
"same_page_merges": outcome["same_page_merges"],
"same_page_dropped": outcome["same_page_dropped"],
@@ -95,11 +97,32 @@ def _optimize(structure, page_texts, do_expand, model):
"before": outcome["before"], "after": outcome["after"]}
+def _optimize(structure, page_texts, do_expand, model, concurrency=None):
+ import asyncio
+ return asyncio.run(_optimize_async(structure, page_texts, do_expand, model,
+ concurrency=concurrency))
+
+
+async def _optimize_and_summarize(structure, page_texts, optimize_model, summary_model,
+ concurrency, max_words=None):
+ """Expand and summarize on one loop: a node is summarized as soon as
+ expand can no longer change it, a parent once its children are done."""
+ from ..utils import SummaryScheduler
+ scheduler = SummaryScheduler(structure, [(text, 0) for text in page_texts],
+ model=summary_model, concurrency=concurrency,
+ max_words=max_words)
+ report = await _optimize_async(structure, page_texts, True, optimize_model,
+ on_final=scheduler.mark_final,
+ concurrency=concurrency)
+ await scheduler.finish()
+ return report
+
+
def page_index_flash(pdf, summary=True, summary_model=None,
optimize: str | bool | None = None, 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. The tree extraction itself uses no LLM; by default an LLM writes node summaries and expands the tree (``summary=False, optimize=False`` runs fully LLM-free). 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 (a model unreachable after the retry ladder — a missing credential included — fails the run loudly from expand itself; a per-prompt rejection leaves just that node collapsed), ``"merge"`` for deterministic merge only, ``False`` to disable. ``True`` is accepted as ``"full"`` for backward compatibility; defaults to ``"full"``. Expand needs readable page text, so a bookmark-only or scanned PDF runs the merge half only (``expands`` reports 0). optimize_expand: deprecated — use ``optimize``. Honored only when ``optimize`` is not passed (or is the legacy ``True``): ``False`` maps to ``"merge"``, ``True`` to ``"full"``. 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. """
+ use_embedded_toc=True, summary_max_words=None) -> dict:
+ """Build a PageIndex tree structure from a PDF using layout statistics. The tree extraction itself uses no LLM; by default an LLM writes node summaries and expands the tree (``summary=False, optimize=False`` runs fully LLM-free). 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 (a model unreachable after the retry ladder — a missing credential included — fails the run loudly from expand itself; a per-prompt rejection leaves just that node collapsed), ``"merge"`` for deterministic merge only, ``False`` to disable. ``True`` is accepted as ``"full"`` for backward compatibility; defaults to ``"full"``. Expand needs readable page text, so a bookmark-only or scanned PDF runs the merge half only (``expands`` reports 0). optimize_expand: deprecated — use ``optimize``. Honored only when ``optimize`` is not passed (or is the legacy ``True``): ``False`` maps to ``"merge"``, ``True`` to ``"full"``. optimize_model: the LLM model for expand (defaults to the summary model). summary_concurrency: cap on simultaneous indexing model calls per lane: the summaries, and expand up to its own ceiling of 32 (the lanes overlap, so up to cap + min(32, cap) calls run at once); None uses the library defaults (64 and 32). 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. summary_max_words: word cap each node summary is asked to stay within; None uses the library default (150). 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_expand is not None:
import warnings
warnings.warn(
@@ -118,28 +141,34 @@ def page_index_flash(pdf, summary=True, summary_model=None,
f"optimize must be 'full', 'merge', or False, got {optimize!r}")
result = extract_toc(_validate_pdf(pdf), use_embedded_toc=use_embedded_toc)
structure = result.get("structure", [])
+ if summary and structure and summary_model is None:
+ from ..utils import ConfigLoader
+ cfg = ConfigLoader().load()
+ summary_model = getattr(cfg, 'summary_model', None) or cfg.model
+ # bookmark-only extractions carry no page_texts and scanned ones
+ # only empty strings; expand needs text
+ pages = result.pop("page_texts", None) or []
+ do_expand = optimize == "full" and any(pages)
+ if optimize and structure and summary and do_expand:
+ import asyncio
+ result["optimize"] = asyncio.run(_optimize_and_summarize(
+ structure, pages, optimize_model=optimize_model or summary_model,
+ summary_model=summary_model, concurrency=summary_concurrency,
+ max_words=summary_max_words))
+ return result
if optimize and structure:
- # bookmark-only extractions carry no page_texts and scanned ones
- # only empty strings; expand needs text
- pages = result.get("page_texts") or []
- result["optimize"] = _optimize(structure, pages,
- optimize == "full" and any(pages),
- optimize_model or summary_model)
+ result["optimize"] = _optimize(structure, pages, do_expand,
+ optimize_model or summary_model,
+ concurrency=summary_concurrency)
if summary and structure:
import asyncio
- from ..utils import ConfigLoader
- if summary_model is None:
- cfg = ConfigLoader().load()
- summary_model = getattr(cfg, 'summary_model', None) or cfg.model
- page_texts = result.pop("page_texts", [])
- page_list = [(text, 0) for text in page_texts]
+ page_list = [(text, 0) for text in pages]
asyncio.run(_summarize(structure, page_list, summary_model,
- concurrency=summary_concurrency))
- else:
- result.pop("page_texts", None)
- if structure:
- from ..utils import strip_internal_keys
- strip_internal_keys(structure) # summarize_tree does this on its way out
+ concurrency=summary_concurrency,
+ max_words=summary_max_words))
+ elif structure:
+ from ..utils import strip_internal_keys
+ strip_internal_keys(structure) # summarize_tree does this on its way out
return result
diff --git a/pageindex/local_api.py b/pageindex/local_api.py
index 40aa3aedc..463909eac 100644
--- a/pageindex/local_api.py
+++ b/pageindex/local_api.py
@@ -12,7 +12,7 @@
from .errors import PageIndexAPIError
from .local_store import DocStore
-from .utils import run_off_loop
+from .utils import count_tokens, run_off_loop
logger = logging.getLogger(__name__)
@@ -37,11 +37,19 @@ class LocalAPI:
"""Backs PageIndexClient's local mode. One instance per client."""
def __init__(self, storage_path: str, model: str, summary_model: str,
- index_backend: dict | None = None):
+ index_backend: dict | None = None,
+ summary_max_words: int | None = None,
+ summary_concurrency: int | None = None,
+ use_embedded_toc: bool = True,
+ optimize: str = "full"):
self._store = DocStore(storage_path)
self._model = model
self._summary_model = summary_model
self._index_backend = index_backend
+ self._summary_max_words = summary_max_words
+ self._summary_concurrency = summary_concurrency
+ self._use_embedded_toc = use_embedded_toc
+ self._optimize = optimize
from .utils import ConfigLoader
self._config_loader = ConfigLoader()
@@ -206,9 +214,7 @@ def _extract_page_texts(file_path: str) -> list[str]:
def _index_standard(self, file_path: str, page_texts: list[str]) -> tuple[list, str | None]:
from .page_index_classic import page_index_main
- import litellm
- page_list = [(text, litellm.token_counter(model=self._model, text=text))
- for text in page_texts]
+ page_list = [(text, count_tokens(text, model=self._model)) for text in page_texts]
opt = self._config_loader.load({
"model": self._model,
"summary_model": self._summary_model,
@@ -231,8 +237,11 @@ def _index_flash(self, file_path: str) -> tuple[list, str | None]:
generate_doc_description, write_node_id)
result = page_index_flash(file_path, summary=True,
summary_model=self._summary_model,
- optimize="full",
- optimize_model=self._summary_model)
+ optimize=False if self._optimize == "off" else self._optimize,
+ optimize_model=self._summary_model,
+ summary_concurrency=self._summary_concurrency,
+ summary_max_words=self._summary_max_words,
+ use_embedded_toc=self._use_embedded_toc)
structure = result.get("structure", [])
if not structure:
raise PageIndexAPIError(
diff --git a/pageindex/tree_optimize.py b/pageindex/tree_optimize.py
index f1638ff4e..ddf397c7d 100644
--- a/pageindex/tree_optimize.py
+++ b/pageindex/tree_optimize.py
@@ -36,7 +36,8 @@
`key_items`: the pages stay reachable by scanning the parent, but the titles
are routing information that would otherwise be lost.
-merge_same_page() runs first, as a special case of the same idea. Retrieval is
+merge_same_page() runs at every seam of optimize() where same-page duplicates can
+appear, as a special case of the same idea. Retrieval is
page-granular, so frontier siblings covering identical pages cannot be told apart:
an agent routed to any of them reads the same text, and because the leaf summary
prompt sees only that text, their summaries come back near-identical. They collapse
@@ -491,9 +492,10 @@ def union_title(titles, node):
def merge_same_page(structure, log):
"""Collapse frontier siblings that cover exactly the same pages.
- Deterministic and free. Runs before merge() because a narrower tree changes
- its ancestors' tree_cost, and before expand() because children an expand pass
- lands on one page are the same redundancy arriving later.
+ Deterministic and free. Runs at every optimize() seam where the redundancy can appear:
+ before merge() because a narrower tree changes its ancestors' tree_cost,
+ again after it because a collapsed subtree can land on a sibling's exact
+ pages, and on a node's children right after expand attaches them.
"""
changed = False
@@ -659,7 +661,7 @@ async def expand(structure, pages, lines, args, log, frozen):
priced with expand_cost and the cheapest is kept.
"""
changed = False
- semaphore = asyncio.Semaphore(EXPAND_CONCURRENCY)
+ semaphore = asyncio.Semaphore(args.concurrency)
async def proposals_for(node):
"""The model half of one node's lookahead: the empty-retry ladder and
@@ -705,6 +707,7 @@ async def process(node):
log.append({"op": "expand", "node_id": node.get("node_id"),
"decision": "no_children", "S": span, "attempts": attempts})
frozen.add(node.get("node_id"))
+ args.settled([node])
return
scored = []
@@ -719,6 +722,8 @@ async def process(node):
cost = best["expand_cost"]
gain = span - cost
ratio = gain / span if span else 0.0
+ # strict: at cost == span the next round's merge (span <= cost) would
+ # fold the node right back, touching children already marked final
keep = cost < span and ratio >= args.min_gain_ratio
note(args.progress,
@@ -739,15 +744,21 @@ async def process(node):
for c in best["children"]]})
frozen.add(node.get("node_id"))
- if keep:
- changed = True
- attach_children(node, best["children"], lines)
- results = await asyncio.gather(*(process(child)
- for child in node["nodes"]),
- return_exceptions=True)
- for result in results:
- if isinstance(result, BaseException):
- raise result
+ if not keep:
+ args.settled([node])
+ return
+ changed = True
+ attach_children(node, best["children"], lines)
+ if args.do_merge:
+ merge_same_page([node], log)
+ # settle after the fusion: mark_final snapshots the children, finish() rejects a later change
+ args.settled([node])
+ results = await asyncio.gather(*(process(child)
+ for child in node["nodes"]),
+ return_exceptions=True)
+ for result in results:
+ if isinstance(result, BaseException):
+ raise result
results = await asyncio.gather(*(process(node)
for node, _ in flatten(structure)),
@@ -768,11 +779,20 @@ def default_model():
return getattr(opt, "summary_model", None) or opt.model
+def final_nodes(nodes, trigger, frozen):
+ """Nodes whose children will not change any more: everything but a
+ collapsed node over the trigger that expand has not judged yet."""
+ return [node for node, _ in flatten(nodes)
+ if not is_frontier(node) or S(node) <= trigger
+ or node.get("node_id") in frozen]
+
+
async def optimize(structure, pages, lines, model=None, routing=ROUTING_COST,
trigger_pages=TRIGGER_PAGES, min_gain_ratio=0.0,
do_merge=True, do_expand=True, max_rounds=3, page_count=None,
cache=None, kinds=("section", "table"), empty_retries=1,
- do_relabel=True, progress=False):
+ do_relabel=True, progress=False, on_final=None,
+ concurrency=None):
"""Run merge and expand over a tree until neither changes anything.
Mutates `structure` in place and returns a summary.
@@ -782,30 +802,52 @@ async def optimize(structure, pages, lines, model=None, routing=ROUTING_COST,
ancestors' tree_cost. Nodes decided by either operator are frozen for the rest
of the run, so a node cannot be collapsed and re-expanded in alternating
rounds.
+
+ `on_final(nodes)` hears, as the run goes, which nodes will not change any
+ more (see final_nodes): after each round's merges, as expand decides each
+ candidate, and for the whole tree at the end.
"""
if do_expand and pages is None:
raise ValueError("expand needs the PDF pages; pass pages/lines or do_expand=False")
+ log, frozen = [], set()
+
+ def settled(nodes):
+ if on_final is not None:
+ on_final(final_nodes(nodes, trigger_pages, frozen))
opts = SimpleNamespace(model=model or default_model(), routing=routing,
trigger_pages=trigger_pages,
min_gain_ratio=min_gain_ratio, cache=cache,
kinds=set(kinds) if kinds else None,
- empty_retries=empty_retries, progress=progress)
+ empty_retries=empty_retries, progress=progress,
+ settled=settled, do_merge=do_merge,
+ concurrency=min(EXPAND_CONCURRENCY,
+ concurrency or EXPAND_CONCURRENCY))
baseline = set(validate(structure, page_count)) if page_count else set()
before = complexity(structure, page_count, routing=routing) if page_count else {}
- log, frozen = [], set()
rounds = 0
for round_no in range(1, max_rounds + 1):
rounds = round_no
note(progress, f" round {round_no}")
- same_page = merge_same_page(structure, log) if do_merge else False
+ round_start = len(log)
+ if do_merge:
+ merge_same_page(structure, log)
merged = merge(structure, routing, log, frozen, progress) if do_merge else False
+ if merged:
+ # a collapsed subtree can land on a sibling's exact pages
+ merge_same_page(structure, log)
+ settled(structure)
expanded = await expand(structure, pages, lines, opts, log, frozen) \
if do_expand else False
+ # derived from this round's log slice, so the flag counts the fusions
+ # inside expand too and cannot disagree with the run counters
+ same_page = any(e["op"] == "merge_same_page" for e in log[round_start:])
log.append({"op": "round", "round": round_no, "same_page": same_page,
"merged": merged, "expanded": expanded})
if not (same_page or merged or expanded):
break
+ if on_final is not None:
+ on_final([node for node, _ in flatten(structure)])
id_map = relabel(structure) if do_relabel else {}
after = complexity(structure, page_count, routing=routing) if page_count else {}
diff --git a/pageindex/types.py b/pageindex/types.py
index b02962899..8a7da077f 100644
--- a/pageindex/types.py
+++ b/pageindex/types.py
@@ -32,6 +32,10 @@ class LocalIndexConfig(TypedDict, total=False):
mode: Literal["local"]
model: str
summary_model: str
+ summary_max_words: int
+ summary_concurrency: int
+ use_embedded_toc: bool
+ optimize: Literal["full", "merge", "off"]
backend: dict
storage_path: Union[str, os.PathLike[str]]
diff --git a/pageindex/utils.py b/pageindex/utils.py
index 83e9d0e79..30799762a 100644
--- a/pageindex/utils.py
+++ b/pageindex/utils.py
@@ -9,6 +9,8 @@
import PyPDF2
import copy
import asyncio
+import heapq
+from contextlib import asynccontextmanager
from io import BytesIO
from dotenv import find_dotenv, load_dotenv
load_dotenv(find_dotenv(usecwd=True))
@@ -69,7 +71,10 @@ def count_tokens(text, model=None):
if not text:
return 0
import litellm
- return litellm.token_counter(model=model, text=text)
+ try:
+ return litellm.token_counter(model=model, text=text)
+ except Exception:
+ return litellm.token_counter(model=None, text=text)
def _strip_prefix(s, prefix):
@@ -513,14 +518,13 @@ def add_preface_if_needed(data):
def get_page_tokens(pdf_path, model=None, pdf_parser="PyPDF2"):
- import litellm
if pdf_parser == "PyPDF2":
pdf_reader = PyPDF2.PdfReader(pdf_path)
page_list = []
for page_num in range(len(pdf_reader.pages)):
page = pdf_reader.pages[page_num]
page_text = page.extract_text()
- token_length = litellm.token_counter(model=model, text=page_text)
+ token_length = count_tokens(page_text, model=model)
page_list.append((page_text, token_length))
return page_list
elif pdf_parser == "PyMuPDF":
@@ -533,7 +537,7 @@ def get_page_tokens(pdf_path, model=None, pdf_parser="PyPDF2"):
page_list = []
for page in doc:
page_text = page.get_text()
- token_length = litellm.token_counter(model=model, text=page_text)
+ token_length = count_tokens(page_text, model=model)
page_list.append((page_text, token_length))
return page_list
else:
@@ -743,6 +747,48 @@ async def generate_summaries_for_structure(structure, model=None):
SUMMARY_CONCURRENCY = 64 # simultaneous summary model calls
SUMMARY_RAW_TEXT_TOKENS = 200 # leaves under this reuse their raw text as the summary
SUMMARY_INTRO_MAX_PAGES = 3 # cap on leading pages fed into a parent summary
+SUMMARY_MAX_WORDS = 150 # word cap the summary prompts ask for
+
+
+class _PriorityGate:
+ """Semaphore that admits the highest-priority waiter first, FIFO within a priority."""
+
+ def __init__(self, permits):
+ if permits < 1:
+ raise ValueError("permits must be >= 1")
+ self._free = permits
+ self._waiters = []
+ self._seq = 0
+
+ @asynccontextmanager
+ async def slot(self, prio):
+ await self.acquire(prio)
+ try:
+ yield
+ finally:
+ self.release()
+
+ async def acquire(self, prio):
+ if self._free > 0 and not self._waiters:
+ self._free -= 1
+ return
+ fut = asyncio.get_running_loop().create_future()
+ heapq.heappush(self._waiters, (-prio, self._seq, fut))
+ self._seq += 1
+ try:
+ await fut
+ except asyncio.CancelledError:
+ if fut.done() and not fut.cancelled():
+ self.release() # granted while cancelling: pass the permit on
+ raise
+
+ def release(self):
+ while self._waiters:
+ _, _, fut = heapq.heappop(self._waiters)
+ if not fut.done():
+ fut.set_result(None) # the permit moves straight to this waiter
+ return
+ self._free += 1
def get_intro_text(node, pdf_pages, max_pages=SUMMARY_INTRO_MAX_PAGES):
@@ -823,29 +869,83 @@ def strip_internal_keys(structure):
return structure
-async def summarize_tree(structure, pdf_pages, model=None,
- small_node_tokens=SUMMARY_RAW_TEXT_TOKENS,
- max_intro_pages=SUMMARY_INTRO_MAX_PAGES, concurrency=None):
- """Bottom-up summaries: leaves from their own pages, parents composed from
- child summaries plus the pages no child covers. A parent's summary describes
- its whole subtree (end_index union semantics). Nodes that already carry a
- summary are left untouched; leaves under `small_node_tokens` use their raw
- text as the summary without a model call."""
- semaphore = asyncio.Semaphore(concurrency or SUMMARY_CONCURRENCY)
- asked = answered = False
-
- async def ask(prompt):
- nonlocal asked, answered
- asked = True
- async with semaphore:
- reply = await llm_acompletion(model, prompt)
+def _subtree(nodes):
+ for node in nodes:
+ yield node
+ yield from _subtree(node.get('nodes') or [])
+
+
+class SummaryScheduler:
+ """Bottom-up summaries, taking nodes as they are marked final.
+
+ A node's task waits for its mark (its children will not change any more),
+ then for its children's tasks, then makes its own call: leaves from their
+ own pages, parents composed from child summaries plus the pages no child
+ covers. Marked subtrees get their tasks deepest node first and queued
+ calls leave the gate deepest first: depth counts the calls left on a
+ node's path to the root, its own included."""
+
+ def __init__(self, structure, pdf_pages, model=None,
+ small_node_tokens=SUMMARY_RAW_TEXT_TOKENS,
+ max_intro_pages=SUMMARY_INTRO_MAX_PAGES, concurrency=None,
+ max_words=None):
+ self.structure = structure
+ self._pdf_pages = pdf_pages
+ self._model = model
+ self._small_node_tokens = small_node_tokens
+ self._max_intro_pages = max_intro_pages
+ self._max_words = max_words or SUMMARY_MAX_WORDS
+ self._gate = _PriorityGate(concurrency or SUMMARY_CONCURRENCY)
+ self._asked = self._answered = False
+ self._marks = {} # id(node) -> future resolved once the node is final
+ self._tasks = {} # id(node) -> its summary task
+ self._finals = [] # (node, ids of its children when it was marked)
+
+ def mark_final(self, nodes):
+ """These nodes will not gain, lose or swap children: their summaries
+ may start. Their subtrees get tasks, deepest node first."""
+ nodes = list(nodes)
+ for node in nodes:
+ mark = self._mark(node)
+ if not mark.done():
+ mark.set_result(None)
+ self._finals.append((node, tuple(id(c) for c in node.get('nodes') or [])))
+ marked = {id(node) for node in nodes}
+ order = []
+
+ def walk(nodes, depth, inside):
+ for node in nodes:
+ inside_here = inside or id(node) in marked
+ if inside_here:
+ order.append((depth, node))
+ walk(node.get('nodes') or [], depth + 1, inside_here)
+ walk(self.structure, 1, False)
+ for depth, node in sorted(order, key=lambda pair: -pair[0]):
+ self._task(node, depth)
+
+ def _mark(self, node):
+ mark = self._marks.get(id(node))
+ if mark is None:
+ mark = self._marks[id(node)] = asyncio.get_running_loop().create_future()
+ return mark
+
+ def _task(self, node, depth):
+ task = self._tasks.get(id(node))
+ if task is None:
+ task = self._tasks[id(node)] = asyncio.create_task(self._visit(node, depth))
+ return task
+
+ async def _ask(self, prompt, prio):
+ self._asked = True
+ async with self._gate.slot(prio):
+ reply = await llm_acompletion(self._model, prompt)
if reply:
- answered = True
+ self._answered = True
return reply
- async def leaf_summary(node):
- text = get_text_of_pdf_pages(pdf_pages, node['start_index'], node['end_index'])
- if count_tokens(text, model="gpt-4o") < small_node_tokens:
+ async def _leaf_summary(self, node, prio):
+ text = get_text_of_pdf_pages(self._pdf_pages, node['start_index'], node['end_index'])
+ if count_tokens(text, model=self._model) < self._small_node_tokens:
return text.strip()
# A node merged from same-page siblings carries a title joined from theirs.
@@ -862,34 +962,33 @@ async def leaf_summary(node):
prompt = f"""You are given a text chunk from a document.
Your task is to generate a concise description of everything that is covered in the text, summarizing all its points without omitting any type of content.
- Keep the description concise and to the point, avoiding unnecessary details.{ask_title}
+ Keep the description concise and to the point, avoiding unnecessary details, within {self._max_words} words.{ask_title}
Given Text: {text}
Reply strictly in the following JSON format:
{{{title_field}
- "points": ,
"summary":
}}
Follow strictly the above JSON return format. Do not include any other text!
"""
- reply = await ask(prompt)
+ reply = await self._ask(prompt, prio)
if retitle:
written = parse_title(reply)
if written:
node['title'] = written
return parse_summary(reply)
- async def parent_summary(node):
+ async def _parent_summary(self, node, prio):
children = node['nodes']
- intro = get_intro_text(node, pdf_pages, max_pages=max_intro_pages)
+ intro = get_intro_text(node, self._pdf_pages, max_pages=self._max_intro_pages)
listing = json.dumps(
[{'title': c.get('title', ''), 'summary': c.get('summary', '')} for c in children],
ensure_ascii=False)
prompt = f"""You are given a section of a document: the text that opens the section (possibly empty) and the titles and summaries of its subsections.
Your task is to generate a concise description of everything that is covered in the whole section, summarizing all its points without omitting any type of content.
- Keep the description concise and to the point, avoiding unnecessary details.
+ Keep the description concise and to the point, avoiding unnecessary details, within {self._max_words} words.
Section Title: {node.get('title', '')}
@@ -899,18 +998,18 @@ async def parent_summary(node):
Reply strictly in the following JSON format:
{{
- "points": ,
"summary":
}}
Follow strictly the above JSON return format. Do not include any other text!
"""
- return parse_summary(await ask(prompt))
+ return parse_summary(await self._ask(prompt, prio))
- async def visit(node):
+ async def _visit(self, node, depth):
+ await self._mark(node)
children = node.get('nodes') or []
if children:
- done = await asyncio.gather(*(visit(child) for child in children),
+ done = await asyncio.gather(*(self._task(child, depth + 1) for child in children),
return_exceptions=True)
for result in done:
if isinstance(result, Exception) and _is_unrecoverable(result):
@@ -918,32 +1017,69 @@ async def visit(node):
if node.get('summary'):
return
try:
- node['summary'] = await (parent_summary(node) if children else leaf_summary(node))
+ node['summary'] = await (self._parent_summary(node, depth) if children
+ else self._leaf_summary(node, depth))
except Exception as e:
node['summary'] = ""
if _is_unrecoverable(e):
raise
- results = await asyncio.gather(*(visit(root) for root in structure),
- return_exceptions=True)
- for r in results:
- if isinstance(r, Exception) and _is_unrecoverable(r):
- raise r
-
- # Raw-text leaves summarize without the model, so they cannot vouch for
- # it: a run whose every model call failed still fails loud.
- def _any_summary(nodes):
- return any(n.get('summary') or _any_summary(n.get('nodes') or [])
- for n in nodes)
- if (asked and not answered) or not _any_summary(structure):
- raise RuntimeError(
- "Summary generation failed for all nodes "
- "(every summary call failed or returned empty; "
- "check the model and its context limits)"
- )
+ async def finish(self):
+ """Wait for every summary; fails loud if the model never answered."""
+ # mark_final is a promise: the node stays in the tree and keeps its
+ # children. Verify it before awaiting anything - a broken promise
+ # leaves tasks whose marks can never arrive, a hang with no diagnosis.
+ live = {id(node) for node in _subtree(self.structure)}
+ for node, children in self._finals:
+ if id(node) not in live or tuple(id(c) for c in node.get('nodes') or []) != children:
+ name = node.get('title') or node.get('node_id') or '?'
+ raise RuntimeError(f"node {name!r} was dropped or changed "
+ "after it was marked final")
+ undecided = sum(1 for nid in live
+ if not (mark := self._marks.get(nid)) or not mark.done())
+ if undecided:
+ raise RuntimeError(f"{undecided} node(s) were never marked final; "
+ "their summaries would wait forever")
+ results = await asyncio.gather(*(self._task(root, 1) for root in self.structure),
+ return_exceptions=True)
+ for r in results:
+ if isinstance(r, Exception) and _is_unrecoverable(r):
+ raise r
+
+ # Raw-text leaves summarize without the model, so they cannot vouch for
+ # it: a run whose every model call failed still fails loud.
+ def _any_summary(nodes):
+ return any(n.get('summary') or _any_summary(n.get('nodes') or [])
+ for n in nodes)
+ if (self._asked and not self._answered) or not _any_summary(self.structure):
+ raise RuntimeError(
+ "Summary generation failed for all nodes "
+ "(every summary call failed or returned empty; "
+ "check the model and its context limits)"
+ )
+
+ strip_internal_keys(self.structure)
+ return self.structure
- strip_internal_keys(structure)
- return structure
+
+async def summarize_tree(structure, pdf_pages, model=None,
+ small_node_tokens=SUMMARY_RAW_TEXT_TOKENS,
+ max_intro_pages=SUMMARY_INTRO_MAX_PAGES, concurrency=None,
+ max_words=None):
+ """Bottom-up summaries: leaves from their own pages, parents composed from
+ child summaries plus the pages no child covers. A parent's summary describes
+ its whole subtree (end_index union semantics). Nodes that already carry a
+ summary are left untouched; leaves under `small_node_tokens` use their raw
+ text as the summary without a model call; every prompt asks for at most
+ `max_words` words. Model calls run deepest node first, both in starting
+ order and in leaving the queue: depth counts the calls left on a node's
+ path to the root, its own included."""
+ scheduler = SummaryScheduler(structure, pdf_pages, model=model,
+ small_node_tokens=small_node_tokens,
+ max_intro_pages=max_intro_pages,
+ concurrency=concurrency, max_words=max_words)
+ scheduler.mark_final(list(_subtree(structure)))
+ return await scheduler.finish()
def create_clean_structure_for_description(structure):
diff --git a/run_pageindex.py b/run_pageindex.py
index f2ae111be..bf4b386a4 100644
--- a/run_pageindex.py
+++ b/run_pageindex.py
@@ -33,6 +33,10 @@
help='(legacy) Same as --index-model')
parser.add_argument('--summary-model', type=str, default=None,
help='Model for node summaries (falls back to config.yaml summary_model, then --index-model, then --model)')
+ parser.add_argument('--summary-max-words', type=int, default=None,
+ help='Word cap for each node summary (flash mode; default 150)')
+ parser.add_argument('--summary-concurrency', type=int, default=None,
+ help='Cap on simultaneous indexing model calls per lane (flash mode; default 64, expand tops out at 32)')
parser.add_argument('--toc-check-pages', type=int, default=None,
help='Number of pages to check for table of contents (PDF only)')
@@ -66,14 +70,15 @@
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 is not None and not (args.pdf_path and args.mode == 'flash'):
- raise ValueError("--optimize requires Flash mode with --pdf_path")
+ for flag, value in (('--optimize', args.optimize),
+ ('--embedded-toc', args.embedded_toc),
+ ('--summary', args.summary),
+ ('--summary-max-words', args.summary_max_words),
+ ('--summary-concurrency', args.summary_concurrency)):
+ if value is not None and not (args.pdf_path and args.mode == 'flash'):
+ raise ValueError(f"{flag} 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),
@@ -107,6 +112,8 @@
summary_model=summary_model,
use_embedded_toc=args.embedded_toc if args.embedded_toc is not None else True,
summary=will_summarize,
+ summary_max_words=args.summary_max_words,
+ summary_concurrency=args.summary_concurrency,
)
if not toc_with_page_number.get('structure'):
raise ValueError("PageIndex Flash could not extract a structure from this PDF; "
diff --git a/tests/test_client.py b/tests/test_client.py
index ec7ed845e..9eacbe220 100644
--- a/tests/test_client.py
+++ b/tests/test_client.py
@@ -1187,6 +1187,151 @@ async def unexpected(model, prompt):
assert out[0]["summary"] == "tiny"
+def test_summarize_tree_starts_the_deepest_leaf_first(monkeypatch):
+ """Visited level by level, the shallow leaves A and D would take both
+ permits before the deep leaf C even reached the gate. C must start first:
+ the parents still owed above it are what the run ends on."""
+ started = []
+
+ async def fake(model, prompt):
+ started.append(next(w for w in ("alpha", "gamma", "delta", "Section Title")
+ if w in prompt))
+ for _ in range(8): # enough loop turns for every leaf to reach the gate
+ await asyncio.sleep(0)
+ return '{"points": [], "summary": "ok"}'
+ monkeypatch.setattr(pageindex.utils, "llm_acompletion", fake)
+ pdf_pages = [("alpha " * 5, 5), ("gamma " * 5, 5), ("delta " * 5, 5)]
+ structure = [{"title": "R", "start_index": 1, "end_index": 3, "nodes": [
+ {"title": "A", "start_index": 1, "end_index": 1},
+ {"title": "B", "start_index": 2, "end_index": 2, "nodes": [
+ {"title": "C", "start_index": 2, "end_index": 2}]},
+ {"title": "D", "start_index": 3, "end_index": 3}]}]
+ asyncio.run(pageindex.utils.summarize_tree(
+ structure, pdf_pages, small_node_tokens=0, concurrency=2))
+ assert started.index("gamma") < started.index("alpha")
+ assert started.index("gamma") < started.index("delta")
+ assert started.count("Section Title") == 2
+
+
+def test_summarize_tree_admits_a_ready_parent_before_a_queued_shallow_leaf(monkeypatch):
+ """With one permit the leaf D is already queued when C's child finishes
+ and C becomes ready. C has two calls left above it, D one, so C goes
+ first even though D asked earlier."""
+ started = []
+
+ async def fake(model, prompt):
+ started.append(next(w for w in ("Section Title: C", "Section Title",
+ "alpha", "epsilon", "delta")
+ if w in prompt))
+ for _ in range(8):
+ await asyncio.sleep(0)
+ return '{"points": [], "summary": "ok"}'
+ monkeypatch.setattr(pageindex.utils, "llm_acompletion", fake)
+ pdf_pages = [("alpha " * 5, 5), ("epsilon " * 5, 5), ("delta " * 5, 5)]
+ structure = [{"title": "R", "start_index": 1, "end_index": 3, "nodes": [
+ {"title": "B", "start_index": 2, "end_index": 2, "nodes": [
+ {"title": "C", "start_index": 2, "end_index": 2, "nodes": [
+ {"title": "E", "start_index": 2, "end_index": 2}]}]},
+ {"title": "A", "start_index": 1, "end_index": 1},
+ {"title": "D", "start_index": 3, "end_index": 3}]}]
+ asyncio.run(pageindex.utils.summarize_tree(
+ structure, pdf_pages, small_node_tokens=0, concurrency=1))
+ assert started.index("Section Title: C") < started.index("delta")
+
+
+def test_summary_scheduler_starts_a_marked_node_without_waiting_for_the_rest(monkeypatch):
+ """A node marked final summarizes right away while its siblings are still
+ unmarked; marking it again, or marking its parent later, never repeats
+ its call: the parent composes the summary already written."""
+ calls = []
+
+ async def fake(model, prompt):
+ calls.append(prompt)
+ return '{"points": [], "summary": "ok"}'
+ monkeypatch.setattr(pageindex.utils, "llm_acompletion", fake)
+ pdf_pages = [("alpha " * 5, 5), ("beta " * 5, 5)]
+ leaf = {"title": "A", "start_index": 1, "end_index": 1}
+ root = {"title": "R", "start_index": 1, "end_index": 2, "nodes": [
+ leaf, {"title": "B", "start_index": 2, "end_index": 2}]}
+
+ async def scenario():
+ scheduler = pageindex.utils.SummaryScheduler([root], pdf_pages, small_node_tokens=0)
+ scheduler.mark_final([leaf])
+ scheduler.mark_final([leaf])
+ for _ in range(3):
+ await asyncio.sleep(0)
+ assert len(calls) == 1 and "alpha" in calls[0]
+ scheduler.mark_final(list(pageindex.utils._subtree([root])))
+ return await scheduler.finish()
+ asyncio.run(scenario())
+ assert len(calls) == 3
+ assert [n["summary"] for n in (leaf, root["nodes"][1], root)] == ["ok"] * 3
+
+
+def test_finish_fails_loud_when_a_final_node_changes_afterwards(monkeypatch):
+ """mark_final is a promise: the node stays in the tree and keeps its
+ children. finish() verifies it, so a run that broke the promise fails
+ with a diagnosis instead of shipping a tree missing a summarized
+ subtree."""
+ async def fake(model, prompt):
+ return '{"points": [], "summary": "ok"}'
+ monkeypatch.setattr(pageindex.utils, "llm_acompletion", fake)
+ pdf_pages = [("alpha " * 5, 5), ("beta " * 5, 5)]
+ leaf = {"title": "A", "start_index": 1, "end_index": 1}
+ root = {"title": "R", "start_index": 1, "end_index": 2, "nodes": [leaf]}
+
+ async def scenario():
+ scheduler = pageindex.utils.SummaryScheduler([root], pdf_pages, small_node_tokens=0)
+ scheduler.mark_final(list(pageindex.utils._subtree([root])))
+ root["nodes"] = [] # the promise, broken
+ await scheduler.finish()
+ with pytest.raises(RuntimeError, match="marked final"):
+ asyncio.run(asyncio.wait_for(scenario(), 5))
+
+
+def test_finish_fails_loud_instead_of_waiting_on_a_node_never_marked(monkeypatch):
+ """A tasked node whose mark never arrives would park finish() forever,
+ indistinguishable from a slow model. Marking a proper subset must be
+ named as the protocol breach it is."""
+ async def fake(model, prompt):
+ return '{"points": [], "summary": "ok"}'
+ monkeypatch.setattr(pageindex.utils, "llm_acompletion", fake)
+ pdf_pages = [("alpha " * 5, 5), ("beta " * 5, 5)]
+ leaf = {"title": "A", "start_index": 1, "end_index": 1}
+ root = {"title": "R", "start_index": 1, "end_index": 2, "nodes": [leaf]}
+
+ async def scenario():
+ scheduler = pageindex.utils.SummaryScheduler([root], pdf_pages, small_node_tokens=0)
+ scheduler.mark_final([root]) # tasks the subtree, marks only the root
+ await scheduler.finish()
+ with pytest.raises(RuntimeError, match="marked final"):
+ asyncio.run(asyncio.wait_for(scenario(), 5))
+
+ async def root_scenario():
+ other = {"title": "B", "start_index": 2, "end_index": 2}
+ scheduler = pageindex.utils.SummaryScheduler([root, other], pdf_pages, small_node_tokens=0)
+ scheduler.mark_final([root, leaf]) # a root nobody marked: tasked by finish() itself
+ await scheduler.finish()
+ with pytest.raises(RuntimeError, match="marked final"):
+ asyncio.run(asyncio.wait_for(root_scenario(), 5))
+
+
+def test_priority_gate_passes_the_permit_on_when_its_taker_is_cancelled():
+ """A waiter cancelled after the permit was granted but before it ran
+ must hand the permit on, or the pool shrinks by one for good."""
+ async def scenario():
+ gate = pageindex.utils._PriorityGate(1)
+ await gate.acquire(1)
+ waiter = asyncio.create_task(gate.acquire(1))
+ await asyncio.sleep(0) # queued
+ gate.release() # granted to `waiter`, not yet resumed
+ waiter.cancel()
+ with pytest.raises(asyncio.CancelledError):
+ await waiter
+ await asyncio.wait_for(gate.acquire(1), 1) # hangs on a leaked permit
+ asyncio.run(scenario())
+
+
def test_summarize_tree_partial_exhaustion_fails_loud(monkeypatch):
"""One lucky call must not vouch for a model that then went away: a
ladder-exhausted node raises instead of silently blanking."""
@@ -1866,6 +2011,217 @@ async def slow_empty(model, prompt):
assert inflight["peak"] <= 32
+def _expand_fixture():
+ """Twelve pages; X spans nine of them, so expand looks at it and, when
+ the model offers 'Sub One' (page 4) and 'Sub Two' (page 8), splits it
+ into two leaves under the trigger. Bodies are long enough that every
+ leaf summary needs the model."""
+ body = "body " * 250
+ pages = [body] * 12
+ pages[3] = "Sub One\n" + body
+ pages[7] = "Sub Two\n" + body
+ lines = [[l for l in p.splitlines() if l.strip()] for p in pages]
+ tree = [{"title": "R", "start_index": 1, "end_index": 12, "node_id": "0000", "nodes": [
+ {"title": "A", "start_index": 1, "end_index": 3, "node_id": "0001"},
+ {"title": "X", "start_index": 4, "end_index": 12, "node_id": "0002"}]}]
+ return tree, pages, lines
+
+
+def test_final_nodes_holds_back_only_undecided_expand_candidates():
+ import pageindex.tree_optimize as tree_optimize
+
+ tree, _, _ = _expand_fixture()
+ titles = lambda nodes: sorted(n["title"] for n in nodes)
+ assert titles(tree_optimize.final_nodes(tree, 5, set())) == ["A", "R"]
+ assert titles(tree_optimize.final_nodes(tree, 5, {"0002"})) == ["A", "R", "X"]
+ assert titles(tree_optimize.final_nodes(tree, 9, set())) == ["A", "R", "X"]
+
+
+def test_optimize_reports_final_nodes_as_expand_decides(monkeypatch):
+ """Before the first expand call, everything expand cannot touch is
+ reported final; the candidate and what it grows are reported the moment
+ it is decided; the closing report covers the whole tree."""
+ import pageindex.tree_optimize as tree_optimize
+
+ tree, pages, lines = _expand_fixture()
+ reports, replied_after = [], []
+
+ async def propose(model, prompt):
+ await asyncio.sleep(0.01)
+ replied_after.append(len(reports))
+ return {"subsections": [{"title": "Sub One", "page": 4},
+ {"title": "Sub Two", "page": 8}]}
+ monkeypatch.setattr(tree_optimize, "ask_model", propose)
+ asyncio.run(tree_optimize.optimize(
+ tree, pages, lines, model="m", do_expand=True,
+ on_final=lambda nodes: reports.append(sorted(n["title"] for n in nodes))))
+ before_reply = [t for report in reports[:replied_after[0]] for t in report]
+ assert "R" in before_reply and "A" in before_reply and "X" not in before_reply
+ decided = next(r for r in reports[replied_after[0]:] if "X" in r)
+ assert decided == ["Sub One", "Sub Two", "X"]
+ assert reports[-1] == ["A", "R", "Sub One", "Sub Two", "X"]
+
+
+def test_optimize_reports_a_candidate_kept_collapsed_once_decided(monkeypatch):
+ """A candidate the model finds nothing in is final the moment its retry
+ ladder ends, not at the end of the run."""
+ import pageindex.tree_optimize as tree_optimize
+
+ tree, pages, lines = _expand_fixture()
+ reports = []
+
+ async def nothing(model, prompt):
+ return {"subsections": []}
+ monkeypatch.setattr(tree_optimize, "ask_model", nothing)
+ asyncio.run(tree_optimize.optimize(
+ tree, pages, lines, model="m", do_expand=True,
+ on_final=lambda nodes: reports.append(sorted(n["title"] for n in nodes))))
+ assert ["X"] in reports[:-1]
+
+
+def test_optimize_reports_a_candidate_kept_collapsed_for_too_little_gain_once_decided(monkeypatch):
+ """The other way to stay collapsed: the model proposes a split that does
+ not pay for itself. That node is final right then too."""
+ import pageindex.tree_optimize as tree_optimize
+
+ tree, pages, lines = _expand_fixture()
+ reports = []
+
+ async def propose(model, prompt):
+ return {"subsections": [{"title": "Sub One", "page": 4}, {"title": "Sub Two", "page": 8}]}
+ monkeypatch.setattr(tree_optimize, "ask_model", propose)
+ outcome = asyncio.run(tree_optimize.optimize(
+ tree, pages, lines, model="m", do_expand=True, min_gain_ratio=0.99,
+ on_final=lambda nodes: reports.append(sorted(n["title"] for n in nodes))))
+ assert outcome["kept_collapsed"] == 1 and ["X"] in reports[:-1]
+
+
+def test_merge_fuses_the_same_page_frontier_it_creates_at_once():
+ """Collapsing F leaves it on exactly G's page; the two must fuse right
+ then, not one round later (with a single round they never would)."""
+ import pageindex.tree_optimize as tree_optimize
+
+ tree = [{"title": "R", "start_index": 1, "end_index": 6, "node_id": "0000", "nodes": [
+ {"title": "F", "start_index": 1, "end_index": 1, "node_id": "0001", "nodes": [
+ {"title": "f1", "start_index": 1, "end_index": 1, "node_id": "0002"}]},
+ {"title": "G", "start_index": 1, "end_index": 1, "node_id": "0003"},
+ {"title": "H", "start_index": 2, "end_index": 5, "node_id": "0004"}]}]
+ outcome = asyncio.run(tree_optimize.optimize(
+ tree, None, None, do_expand=False, max_rounds=1))
+ assert outcome["merges"] == 1 and outcome["same_page_merges"] == 1
+ kept = tree[0]["nodes"]
+ assert [n["title"] for n in kept][1:] == ["H"] and kept[0].get("_same_page")
+
+
+def test_expand_fuses_same_page_children_before_reporting_them(monkeypatch):
+ """Two proposed headings on one page become one node as soon as they
+ are attached, so the report never carries a duplicate."""
+ import pageindex.tree_optimize as tree_optimize
+
+ tree, pages, lines = _expand_fixture()
+ pages[3] = "Sub One\nSub Two\n" + pages[3]
+ pages[4] = "Sub Three\n" + pages[4]
+ pages[8] = "Sub Four\n" + pages[8]
+ lines = [[l for l in p.splitlines() if l.strip()] for p in pages]
+
+ async def propose(model, prompt):
+ return {"subsections": [{"title": "Sub One", "page": 4}, {"title": "Sub Two", "page": 4},
+ {"title": "Sub Three", "page": 5}, {"title": "Sub Four", "page": 9}]}
+ monkeypatch.setattr(tree_optimize, "ask_model", propose)
+
+ async def summarize(model, prompt):
+ return '{"summary": "ok"}'
+ monkeypatch.setattr(pageindex.utils, "llm_acompletion", summarize)
+ reports = []
+
+ async def run():
+ # the real scheduler: mark_final snapshots each node's children and
+ # finish() rejects a later change, so settling before the fusion fails here
+ scheduler = pageindex.utils.SummaryScheduler(tree, [(p, 0) for p in pages], model="m")
+
+ def on_final(nodes):
+ reports.append([(n["title"], len(n.get("nodes") or [])) for n in nodes])
+ scheduler.mark_final(nodes)
+ outcome = await tree_optimize.optimize(tree, pages, lines, model="m", do_expand=True,
+ max_rounds=1, on_final=on_final)
+ await scheduler.finish()
+ return outcome
+ outcome = asyncio.run(run())
+ assert outcome["expands"] == 1 and outcome["same_page_merges"] == 1
+ assert [n["title"] for n in tree[0]["nodes"][1]["nodes"]][1:] == ["Sub Three", "Sub Four"]
+ assert all(children != 4 for report in reports for _, children in report)
+ # the round entry agrees with the run counters: the fusion happened this round
+ assert [e["same_page"] for e in outcome["log"] if e["op"] == "round"] == [True]
+
+
+def test_expand_keeps_same_page_children_when_merging_is_off(monkeypatch):
+ """do_merge=False turns off every merge, the fusion inside expand
+ included: a caller who disabled merging keeps every proposed node and
+ every document title."""
+ import pageindex.tree_optimize as tree_optimize
+
+ tree, pages, lines = _expand_fixture()
+ pages[3] = "Sub One\nSub Two\n" + pages[3]
+ pages[4] = "Sub Three\n" + pages[4]
+ pages[8] = "Sub Four\n" + pages[8]
+ lines = [[l for l in p.splitlines() if l.strip()] for p in pages]
+
+ async def propose(model, prompt):
+ return {"subsections": [{"title": "Sub One", "page": 4}, {"title": "Sub Two", "page": 4},
+ {"title": "Sub Three", "page": 5}, {"title": "Sub Four", "page": 9}]}
+ monkeypatch.setattr(tree_optimize, "ask_model", propose)
+ outcome = asyncio.run(tree_optimize.optimize(
+ tree, pages, lines, model="m", do_expand=True, do_merge=False, max_rounds=1))
+ children = tree[0]["nodes"][1]["nodes"]
+ assert [n["title"] for n in children] == ["Sub One", "Sub Two", "Sub Three", "Sub Four"]
+ assert outcome["same_page_merges"] == 0
+ assert not any(n.get("_same_page") for n in children)
+
+
+def test_flash_summaries_start_while_expand_is_still_deciding(tmp_path, monkeypatch):
+ """With expand on, summaries run on the same loop: a leaf expand cannot
+ touch is already being summarized when the model answers about X. Every
+ node is summarized exactly once and the tree matches the merge+expand
+ pass run on its own."""
+ from conftest import build_pdf
+ import copy
+ import pageindex.flash.api as flash_api
+ import pageindex.tree_optimize as tree_optimize
+
+ tree, pages, _ = _expand_fixture()
+ pdf = tmp_path / "doc.pdf"
+ pdf.write_bytes(build_pdf(["x"]))
+ monkeypatch.setattr(flash_api, "extract_toc", lambda pdf, **kw: {
+ "structure": copy.deepcopy(tree), "page_texts": list(pages)})
+ started, replied_after, models = [], [], []
+
+ async def propose(model, prompt):
+ models.append(("expand", model))
+ await asyncio.sleep(0.05)
+ replied_after.append(len(started))
+ return {"subsections": [{"title": "Sub One", "page": 4},
+ {"title": "Sub Two", "page": 8}]}
+ monkeypatch.setattr(tree_optimize, "ask_model", propose)
+
+ async def summarize(model, prompt):
+ models.append(("summary", model))
+ started.append(prompt)
+ await asyncio.sleep(0.01)
+ return '{"points": [], "summary": "ok"}'
+ monkeypatch.setattr(pageindex.utils, "llm_acompletion", summarize)
+
+ result = flash_api.page_index_flash(str(pdf), summary_model="m", optimize_model="x")
+ assert replied_after[0] >= 1
+ assert result["optimize"]["expands"] == 1
+ nodes = list(pageindex.utils._subtree(result["structure"]))
+ assert [n["summary"] for n in nodes] == ["ok"] * 5 and len(started) == 5
+ assert set(models) == {("expand", "x"), ("summary", "m")}
+ shape = lambda nodes: [(n["title"], n["start_index"], n["end_index"],
+ shape(n.get("nodes") or [])) for n in nodes]
+ alone = flash_api.page_index_flash(str(pdf), summary=False)
+ assert shape(result["structure"]) == shape(alone["structure"])
+
+
def test_mode_declaration_top_level(monkeypatch):
"""mode= states where documents live; always optional, always
checked, and mode="cloud" alone reads the env key."""
@@ -1967,3 +2323,159 @@ def test_blank_chat_model_carries_no_model_into_agent_config():
client = PageIndexClient()
client.chat_model = " "
assert "model" not in client.openai_agent_config()
+
+
+def test_summary_prompts_cap_words_and_omit_points(monkeypatch):
+ """Both summary prompts ask for the summary alone, within the word cap.
+ The points list the model wrote first was parsed and thrown away, and
+ with it gone the summary swallows its content unless a cap holds it."""
+ prompts = []
+
+ async def capture(model, prompt):
+ prompts.append(prompt)
+ return '{"summary": "ok"}'
+ monkeypatch.setattr(pageindex.utils, "llm_acompletion", capture)
+ pdf_pages = [("alpha " * 5, 5), ("beta " * 5, 5)]
+
+ def tree():
+ return [{"title": "R", "start_index": 1, "end_index": 2,
+ "nodes": [{"title": "A", "start_index": 1, "end_index": 1},
+ {"title": "B", "start_index": 2, "end_index": 2}]}]
+
+ asyncio.run(pageindex.utils.summarize_tree(tree(), pdf_pages, small_node_tokens=0))
+ assert len(prompts) == 3
+ assert all("within 150 words" in p and '"points"' not in p for p in prompts)
+
+ prompts.clear()
+ asyncio.run(pageindex.utils.summarize_tree(tree(), pdf_pages, small_node_tokens=0,
+ max_words=80))
+ assert len(prompts) == 3 and all("within 80 words" in p for p in prompts)
+
+
+def test_page_index_flash_plumbs_summary_max_words(tmp_path, monkeypatch):
+ """summary_max_words reaches the prompts on both summary paths: the
+ plain one and the one overlapped with expand."""
+ from conftest import build_pdf
+ import copy
+ import pageindex.flash.api as flash_api
+ import pageindex.tree_optimize as tree_optimize
+
+ tree, pages, _ = _expand_fixture()
+ pdf = tmp_path / "doc.pdf"
+ pdf.write_bytes(build_pdf(["x"]))
+ monkeypatch.setattr(flash_api, "extract_toc", lambda pdf, **kw: {
+ "structure": copy.deepcopy(tree), "page_texts": list(pages)})
+
+ async def propose(model, prompt):
+ return {"subsections": [{"title": "Sub One", "page": 4},
+ {"title": "Sub Two", "page": 8}]}
+ monkeypatch.setattr(tree_optimize, "ask_model", propose)
+ prompts = []
+
+ async def capture(model, prompt):
+ prompts.append(prompt)
+ return '{"summary": "ok"}'
+ monkeypatch.setattr(pageindex.utils, "llm_acompletion", capture)
+
+ for optimize in (False, "full"):
+ prompts.clear()
+ flash_api.page_index_flash(str(pdf), summary_model="m", optimize=optimize,
+ summary_max_words=80)
+ assert prompts and all("within 80 words" in p for p in prompts), optimize
+
+
+def test_summary_max_words_reaches_the_local_indexer():
+ """index["summary_max_words"] and the flat argument both land on the
+ local indexer; a non-int refuses in the constructor like every slot."""
+ from pageindex import PageIndexLocalClient
+ assert PageIndexLocalClient(index={"summary_max_words": 80})._api._summary_max_words == 80
+ assert PageIndexLocalClient(summary_max_words=80)._api._summary_max_words == 80
+ assert PageIndexLocalClient()._api._summary_max_words is None
+ with pytest.raises(PageIndexAPIError, match=r'index\["summary_max_words"\] must be a'):
+ PageIndexLocalClient(index={"summary_max_words": "80"})
+
+
+def test_flash_knobs_reach_the_local_indexer(tmp_path, sample_pdf, monkeypatch):
+ """summary_concurrency, use_embedded_toc and optimize are settable on the
+ client, flat or in the index slot, and land on page_index_flash the way
+ the CLI's flags do ("off" is no optimize pass at all)."""
+ from pageindex import PageIndexLocalClient
+ captured = {}
+ monkeypatch.setattr(pageindex.flash, "page_index_flash",
+ lambda p, **kw: captured.update(kw) or {
+ "structure": [{"title": "T", "start_index": 1,
+ "end_index": 1, "summary": "s", "nodes": []}]})
+ monkeypatch.setattr(pageindex.utils, "llm_completion",
+ lambda model, prompt, **kw: "d.")
+ monkeypatch.chdir(tmp_path) # the default .pageindex store lands here
+
+ def run(**kwargs):
+ captured.clear()
+ PageIndexLocalClient(**kwargs).submit_document(sample_pdf)
+ return (captured.get("summary_concurrency"), captured["use_embedded_toc"],
+ captured["optimize"])
+
+ assert run(summary_concurrency=8, use_embedded_toc=False, optimize="merge") == (8, False, "merge")
+ assert run(index={"summary_concurrency": 8, "use_embedded_toc": False,
+ "optimize": "off"}) == (8, False, False)
+ assert run() == (None, True, "full")
+ for kwargs, msg in (({"index": {"use_embedded_toc": "no"}},
+ r'index\["use_embedded_toc"\] must be a bool'),
+ ({"optimize": "sometimes"},
+ r'optimize must be "full", "merge" or "off"'),
+ ({"summary_concurrency": "8"}, "summary_concurrency must be a")):
+ with pytest.raises(PageIndexAPIError, match=msg):
+ PageIndexLocalClient(**kwargs)
+
+
+def test_summary_concurrency_caps_expand_too(tmp_path, monkeypatch):
+ """A user who lowers summary_concurrency for a tight quota gets the whole
+ indexing lane lowered: expand's own gate takes the same cap."""
+ from conftest import build_pdf
+ import pageindex.flash.api as flash_api
+ import pageindex.tree_optimize as tree_optimize
+
+ body = "body " * 250
+ pages = [body] * 30
+
+ def roots():
+ return [{"title": f"X{i}", "start_index": 1 + 10 * i, "end_index": 10 + 10 * i,
+ "node_id": f"000{i}"} for i in range(3)]
+ pdf = tmp_path / "doc.pdf"
+ pdf.write_bytes(build_pdf(["x"]))
+ monkeypatch.setattr(flash_api, "extract_toc", lambda pdf, **kw: {
+ "structure": roots(), "page_texts": list(pages)})
+ in_flight, peak = 0, 0
+
+ async def propose(model, prompt):
+ nonlocal in_flight, peak
+ in_flight += 1
+ peak = max(peak, in_flight)
+ await asyncio.sleep(0.02)
+ in_flight -= 1
+ return {"subsections": []}
+ monkeypatch.setattr(tree_optimize, "ask_model", propose)
+
+ async def summarize(model, prompt):
+ return '{"summary": "ok"}'
+ monkeypatch.setattr(pageindex.utils, "llm_acompletion", summarize)
+
+ flash_api.page_index_flash(str(pdf), summary_model="m", summary_concurrency=1)
+ assert peak == 1
+ peak = 0
+ flash_api.page_index_flash(str(pdf), summary_model="m") # control: the three overlap
+ assert peak == 3
+
+
+def test_count_tokens_falls_back_to_the_default_tokenizer(monkeypatch):
+ import litellm
+
+ def token_counter(model=None, text=None, **_):
+ if model is not None:
+ raise TypeError("TextInputSequence must be str") # HF tokenizer on a lone surrogate
+ return 7
+ monkeypatch.setattr(litellm, "token_counter", token_counter)
+ assert pageindex.utils.count_tokens("x", model="groq/llama-3.1-8b-instant") == 7
+
+ monkeypatch.setattr(litellm, "token_counter", lambda model=None, text=None, **_: 3 if model else 7)
+ assert pageindex.utils.count_tokens("x", model="m") == 3 # the model's own count wins when it works
diff --git a/tests/test_flash_extraction.py b/tests/test_flash_extraction.py
index 28c0a963e..c1e423e53 100644
--- a/tests/test_flash_extraction.py
+++ b/tests/test_flash_extraction.py
@@ -237,7 +237,7 @@ def test_optimize_wins_over_deprecated_optimize_expand(tmp_path, monkeypatch):
from pageindex.flash import api as flash_api
seen = {}
- def fake_optimize(structure, pages, do_expand, model):
+ def fake_optimize(structure, pages, do_expand, model, concurrency=None):
seen["do_expand"] = do_expand
return {"merges": 0}
@@ -356,7 +356,7 @@ def test_optimize_full_skips_expand_without_page_texts(tmp_path, monkeypatch):
monkeypatch.setenv("OPENAI_API_KEY", "k")
calls = {}
- def fake_optimize(structure, pages, do_expand, model):
+ def fake_optimize(structure, pages, do_expand, model, concurrency=None):
calls["pages"] = pages
calls["do_expand"] = do_expand
return {"merges": 0}
@@ -382,7 +382,7 @@ def test_optimize_full_skips_expand_on_textless_pages(tmp_path, monkeypatch):
monkeypatch.setenv("OPENAI_API_KEY", "k")
calls = {}
- def fake_optimize(structure, pages, do_expand, model):
+ def fake_optimize(structure, pages, do_expand, model, concurrency=None):
calls["do_expand"] = do_expand
return {"merges": 0}
@@ -467,3 +467,15 @@ def test_flash_cli_rejects_empty_structure(monkeypatch, tmp_path):
with pytest.raises(ValueError, match="try --mode standard"):
_run_flash_cli(monkeypatch, tmp_path, [], [])
assert not (tmp_path / "results").exists()
+
+
+def test_flash_cli_summary_concurrency_reaches_the_indexer(monkeypatch, tmp_path):
+ captured = _run_flash_cli(monkeypatch, tmp_path, ["--summary-concurrency", "8"],
+ [{"title": "A", "start_index": 1, "end_index": 1}])
+ assert captured["summary_concurrency"] == 8
+
+
+def test_flash_cli_summary_max_words_reaches_the_indexer(monkeypatch, tmp_path):
+ captured = _run_flash_cli(monkeypatch, tmp_path, ["--summary-max-words", "80"],
+ [{"title": "A", "start_index": 1, "end_index": 1}])
+ assert captured["summary_max_words"] == 80