perf(retrieval): batch map-nav payload loads and add stage timing

apps/api/tests/contract/test_retrieval_lazy_snapshot_quality_contract.py

@@ -27,6 +27,32 @@
 @dataclass
 class _FakeChunkStore:
     units_by_section: dict[str, list[UnitRow]]
+    document_loads: int = 0
+
+    def load_document_units(
+        self,
+        document_id: str,
+        section_ids: Sequence[str],
+        extra_chunk_ids_by_section: dict[str, Sequence[str]] | None = None,
+    ) -> list[UnitRow]:
+        del document_id
+        self.document_loads += 1
+        selected = {str(section_id) for section_id in section_ids}
+        units = [
+            unit
+            for section_id, section_units in self.units_by_section.items()
+            if section_id in selected
+            for unit in section_units
+        ]
+        known = {unit.chunk_id for unit in units}
+        for chunk_ids in (extra_chunk_ids_by_section or {}).values():
+            for chunk_id in chunk_ids:
+                for section_units in self.units_by_section.values():
+                    for unit in section_units:
+                        if unit.chunk_id == chunk_id and unit.chunk_id not in known:
+                            units.append(unit)
+                            known.add(unit.chunk_id)
+        return units
 
     def load_section_units(
         self,
@@ -47,7 +73,7 @@ def close(self) -> None:
         return None
 
 
-def _providers() -> tuple[ProviderToolSpace, ProviderToolSpace]:
+def _providers() -> tuple[ProviderToolSpace, ProviderToolSpace, _FakeChunkStore]:
     sections = [
         SectionRow("root", None, "Root", "Root", 0, "", 0),
         SectionRow("section", "root", "Root / Section", "Section", 1, "", 1),
@@ -89,11 +115,11 @@ def _providers() -> tuple[ProviderToolSpace, ProviderToolSpace]:
         titles={"doc": "document"},
         chunk_owner_by_id={"duplicate-chunk": "doc", "asset-1": "doc"},
     )
-    return ProviderToolSpace(eager), ProviderToolSpace(lazy)
+    return ProviderToolSpace(eager), ProviderToolSpace(lazy), store
 
 
 def test_lazy_provider_preserves_score_units_and_scores() -> None:
-    eager, lazy = _providers()
+    eager, lazy, store = _providers()
 
     assert build_score_units(eager, "doc") == build_score_units(lazy, "doc")
     assert compute_corpus_map_and_unit_scores(
@@ -103,6 +129,10 @@ def test_lazy_provider_preserves_score_units_and_scores() -> None:
     )
 
     lazy_provider = lazy._provider
+    store.document_loads = 0
+    prefetch = getattr(lazy_provider, "prefetch_document_units")
+    prefetch("doc")
+    assert store.document_loads == 1
     self_units = getattr(lazy_provider, "self_units")
     assert [unit.chunk_id for unit in self_units("leaf")] == [
         "duplicate-chunk",

packages/shared-python/shared/services/retrieval/execution/routes.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import asyncio
+import time
 from contextlib import AbstractAsyncContextManager
 
 from loguru import logger
@@ -180,6 +181,7 @@ async def _run_mapnav_route(
         episode_token_count,
         episode_workflow_plan,
     )
+    snapshot_started = time.perf_counter()
     snapshot = await load_nav_snapshot(
         context.db,
         user_id=context.user_id,
@@ -188,6 +190,14 @@ async def _run_mapnav_route(
         exclude_sections=context.exclude_sections,
         lazy=True,
     )
+    snapshot_seconds = time.perf_counter() - snapshot_started
+    logger.info(
+        "retrieval mapnav stage=snapshot_load seconds={:.3f} documents={} refs={}".format(
+            snapshot_seconds,
+            len(snapshot.document_ids),
+            len(snapshot.chunk_ref_index),
+        )
+    )
 
     # Small-corpus count / snapshot reads may leave a checkout; drop it before
     # the sync LLM episode (same pattern as the retired workflow route).
@@ -197,6 +207,7 @@ async def _run_mapnav_route(
     cfg = build_nav_config()
     toolspace = ProviderToolSpace(snapshot.provider)
 
+    episode_started = time.perf_counter()
     try:
         episode = await asyncio.to_thread(
             run_nav_episode,
@@ -211,9 +222,16 @@ async def _run_mapnav_route(
         )
 
         refs, score_by_chunk_id = build_referenced_chunks(episode, snapshot)
+        logger.info(
+            "retrieval mapnav stage=episode seconds={:.3f} refs={}".format(
+                time.perf_counter() - episode_started,
+                len(refs),
+            )
+        )
     finally:
         snapshot.close()
 
+    hydration_started = time.perf_counter()
     async with open_fresh_database_context() as final_db:
         resolved = await resolve_workflow_references(
             db=final_db,
@@ -260,6 +278,12 @@ async def _run_mapnav_route(
             selected_paths=selected_paths,
             selected_doc_ids=selected_docs,
         )
+    logger.info(
+        "retrieval mapnav stage=hydration seconds={:.3f} results={}".format(
+            time.perf_counter() - hydration_started,
+            len(assembled_rows),
+        )
+    )
 
     stop_reason = str(getattr(episode, "stop_reason", "") or "completed")
     evidence_text = str(getattr(episode, "evidence_text", "") or "")

packages/shared-python/shared/services/retrieval/nav/nav_hierarchy.py

@@ -237,6 +237,30 @@ def release_section_units(self, section_id: str) -> None:
         if callable(release):
             release(section_id)
 
+    def prefetch_document_units(self, doc_id: str) -> None:
+        """Forward a provider's bounded document payload prefetch capability."""
+        provider = self._provider
+        fn = getattr(provider, "prefetch_document_units", None)
+        if not callable(fn):
+            return
+        if callable(getattr(provider, "document_ids", None)):
+            fn(doc_id)
+            return
+        if str(getattr(provider, "doc_id", "")) == str(doc_id):
+            fn()
+
+    def release_document_units(self, doc_id: str) -> None:
+        """Forward release of one document's prefetched payloads."""
+        provider = self._provider
+        fn = getattr(provider, "release_document_units", None)
+        if not callable(fn):
+            return
+        if callable(getattr(provider, "document_ids", None)):
+            fn(doc_id)
+            return
+        if str(getattr(provider, "doc_id", "")) == str(doc_id):
+            fn()
+
 
 @dataclass
 class InMemoryNode:

packages/shared-python/shared/services/retrieval/nav/nav_knowhere.py

@@ -141,6 +141,14 @@ def knowhere_database_url() -> str:
 
 
 class ChunkStore(Protocol):
+    def load_document_units(
+        self,
+        document_id: str,
+        section_ids: Sequence[str],
+        extra_chunk_ids_by_section: Optional[Dict[str, Sequence[str]]] = None,
+    ) -> List[UnitRow]:
+        raise NotImplementedError
+
     def load_section_units(
         self,
         document_id: str,
@@ -211,6 +219,55 @@ def load_section_units(
         finally:
             cur.close()
 
+    def load_document_units(
+        self,
+        document_id: str,
+        section_ids: Sequence[str],
+        extra_chunk_ids_by_section: Optional[Dict[str, Sequence[str]]] = None,
+    ) -> List[UnitRow]:
+        """Load one document's section payloads in a single ordered query."""
+        doc_id = str(document_id).strip()
+        job_result_id = self._revisions.get(doc_id)
+        section_values = [str(section_id).strip() for section_id in section_ids if str(section_id).strip()]
+        extra_ids = [
+            str(chunk_id).strip()
+            for chunk_ids in (extra_chunk_ids_by_section or {}).values()
+            for chunk_id in chunk_ids
+            if str(chunk_id).strip()
+        ]
+        if not doc_id or not job_result_id or (not section_values and not extra_ids):
+            return []
+
+        predicates: list[str] = []
+        params: list[object] = [doc_id, job_result_id]
+        if section_values:
+            predicates.append("section_id = ANY(%s)")
+            params.append(section_values)
+        if extra_ids:
+            predicates.append("chunk_id = ANY(%s)")
+            params.append(extra_ids)
+
+        cur = self._connection().cursor()
+        try:
+            cur.execute(
+                "SELECT chunk_id, section_id, chunk_type, content, sort_order, "
+                "source_chunk_path, file_path, chunk_metadata "
+                "FROM document_chunks "
+                "WHERE document_id = %s AND job_result_id = %s AND ("
+                + " OR ".join(predicates)
+                + ") ORDER BY section_id, sort_order, chunk_id, id",
+                params,
+            )
+            units: list[UnitRow] = []
+            for row in cur.fetchall():
+                section_id = str(row[1]) if row[1] else None
+                if section_id and (doc_id, section_id) in self._excluded_sections:
+                    continue
+                units.append(_unit_from_row(row))
+            return units
+        finally:
+            cur.close()
+
     def close(self) -> None:
         if self._conn is not None:
             self._conn.close()
@@ -538,6 +595,53 @@ def __init__(
         )
         self._chunk_store = chunk_store
         self._root_asset_ids = {str(chunk_id) for chunk_id in root_asset_ids}
+        self._remounted_assets_by_section = {
+            str(section_id): [str(chunk_id) for chunk_id in chunk_ids]
+            for section_id, chunk_ids in (remounted_assets_by_section or {}).items()
+        }
+
+    def prefetch_document_units(self) -> None:
+        """Load this document's section payloads with one bounded SQL query."""
+        section_ids = list(self._sections)
+        loaded = self._chunk_store.load_document_units(
+            self.doc_id,
+            section_ids,
+            self._remounted_assets_by_section,
+        )
+        by_section: Dict[str, List[UnitRow]] = {}
+        for unit in loaded:
+            sid = str(unit.section_id or "").strip()
+            if sid and sid in self._sections:
+                by_section.setdefault(sid, []).append(unit)
+
+        units_by_id = {unit.chunk_id: unit for unit in loaded if unit.chunk_id}
+        for section_id, asset_ids in self._remounted_assets_by_section.items():
+            target = by_section.setdefault(section_id, [])
+            known = {unit.chunk_id for unit in target}
+            for asset_id in asset_ids:
+                asset = units_by_id.get(asset_id)
+                if asset is not None and asset.chunk_id not in known:
+                    target.append(asset)
+                    known.add(asset.chunk_id)
+
+        for section_id in section_ids:
+            units = by_section.get(section_id, [])
+            units.sort(key=lambda unit: (unit.sort_order, unit.chunk_id))
+            self._units_by_section[section_id] = units
+            self._loaded_sections.add(section_id)
+
+        for section_id in section_ids:
+            if is_root_section_path(self.section_path(section_id)):
+                self._units_by_section[section_id] = [
+                    unit
+                    for unit in self._units_by_section.get(section_id, ())
+                    if unit.chunk_id not in self._root_asset_ids
+                ]
+
+    def release_document_units(self) -> None:
+        """Release all payloads loaded by the document batch."""
+        self._units_by_section.clear()
+        self._loaded_sections.clear()
 
     def _ensure_section_loaded(self, section_id: str) -> None:
         super()._ensure_section_loaded(section_id)
@@ -762,6 +866,18 @@ def release_section_units(self, section_id: str) -> None:
         if callable(release):
             release(section_id)
 
+    def prefetch_document_units(self, doc_id: str) -> None:
+        provider = self._docs.get(str(doc_id).strip())
+        prefetch = getattr(provider, "prefetch_document_units", None)
+        if callable(prefetch):
+            prefetch()
+
+    def release_document_units(self, doc_id: str) -> None:
+        provider = self._docs.get(str(doc_id).strip())
+        release = getattr(provider, "release_document_units", None)
+        if callable(release):
+            release()
+
     def address_level(self, node_id: str) -> Optional[NavLevel]:
         sid = str(node_id or "").strip()
         if not sid:

packages/shared-python/shared/services/retrieval/nav/nav_map_scores.py

@@ -385,13 +385,21 @@ def compute_corpus_map_and_unit_scores(
     def unit_factory() -> Iterator[ScoreUnitRow]:
         for document_id in valid_doc_ids:
             children_map, leaves, titles = tree_by_doc[document_id]
-            yield from iter_score_units(
-                ts,
-                document_id,
-                children_map=children_map,
-                leaves=leaves,
-                titles=titles,
-            )
+            prefetch = getattr(ts, "prefetch_document_units", None)
+            release = getattr(ts, "release_document_units", None)
+            if callable(prefetch):
+                prefetch(document_id)
+            try:
+                yield from iter_score_units(
+                    ts,
+                    document_id,
+                    children_map=children_map,
+                    leaves=leaves,
+                    titles=titles,
+                )
+            finally:
+                if callable(release):
+                    release(document_id)
 
     unit_scores = score_unit_stream_hybrid_all(unit_factory, query)