perf: summaries run deepest-first and start while expand is still deciding

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)

pageindex/flash/api.py

@@ -68,38 +68,61 @@ 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"],
             "kept_collapsed": outcome["kept_collapsed"],
             "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