From b984f5234e2b124f23d89317bf1260e40192f178 Mon Sep 17 00:00:00 2001 From: suguanYang Date: Fri, 28 Aug 2026 11:49:52 +0800 Subject: [PATCH 1/2] perf: batch lazy map-nav payload loads --- ...etrieval_lazy_snapshot_quality_contract.py | 35 +++++- .../services/retrieval/execution/routes.py | 24 ++++ .../services/retrieval/nav/nav_hierarchy.py | 14 +++ .../services/retrieval/nav/nav_knowhere.py | 116 ++++++++++++++++++ .../services/retrieval/nav/nav_map_scores.py | 22 ++-- 5 files changed, 201 insertions(+), 10 deletions(-) diff --git a/apps/api/tests/contract/test_retrieval_lazy_snapshot_quality_contract.py b/apps/api/tests/contract/test_retrieval_lazy_snapshot_quality_contract.py index 13c5bcf4..7f9d9e99 100644 --- a/apps/api/tests/contract/test_retrieval_lazy_snapshot_quality_contract.py +++ b/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,9 @@ def test_lazy_provider_preserves_score_units_and_scores() -> None: ) lazy_provider = lazy._provider + 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", diff --git a/packages/shared-python/shared/services/retrieval/execution/routes.py b/packages/shared-python/shared/services/retrieval/execution/routes.py index ad2268ab..97961074 100644 --- a/packages/shared-python/shared/services/retrieval/execution/routes.py +++ b/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 "") diff --git a/packages/shared-python/shared/services/retrieval/nav/nav_hierarchy.py b/packages/shared-python/shared/services/retrieval/nav/nav_hierarchy.py index 93b35598..a3f0c915 100644 --- a/packages/shared-python/shared/services/retrieval/nav/nav_hierarchy.py +++ b/packages/shared-python/shared/services/retrieval/nav/nav_hierarchy.py @@ -237,6 +237,20 @@ 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 callable(fn) and 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 callable(fn) and str(getattr(provider, "doc_id", "")) == str(doc_id): + fn() + @dataclass class InMemoryNode: diff --git a/packages/shared-python/shared/services/retrieval/nav/nav_knowhere.py b/packages/shared-python/shared/services/retrieval/nav/nav_knowhere.py index c2c13187..a1387a64 100644 --- a/packages/shared-python/shared/services/retrieval/nav/nav_knowhere.py +++ b/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: diff --git a/packages/shared-python/shared/services/retrieval/nav/nav_map_scores.py b/packages/shared-python/shared/services/retrieval/nav/nav_map_scores.py index b88df1fa..0223e232 100644 --- a/packages/shared-python/shared/services/retrieval/nav/nav_map_scores.py +++ b/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) From 2f5b9aa806e1c90bb30317e15ed50cdccfe9fba9 Mon Sep 17 00:00:00 2001 From: suguanYang Date: Fri, 28 Aug 2026 12:00:31 +0800 Subject: [PATCH 2/2] fix: forward document prefetch through namespace adapter --- ...est_retrieval_lazy_snapshot_quality_contract.py | 1 + .../shared/services/retrieval/nav/nav_hierarchy.py | 14 ++++++++++++-- 2 files changed, 13 insertions(+), 2 deletions(-) diff --git a/apps/api/tests/contract/test_retrieval_lazy_snapshot_quality_contract.py b/apps/api/tests/contract/test_retrieval_lazy_snapshot_quality_contract.py index 7f9d9e99..cfb82c9e 100644 --- a/apps/api/tests/contract/test_retrieval_lazy_snapshot_quality_contract.py +++ b/apps/api/tests/contract/test_retrieval_lazy_snapshot_quality_contract.py @@ -129,6 +129,7 @@ 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 diff --git a/packages/shared-python/shared/services/retrieval/nav/nav_hierarchy.py b/packages/shared-python/shared/services/retrieval/nav/nav_hierarchy.py index a3f0c915..0609cbc2 100644 --- a/packages/shared-python/shared/services/retrieval/nav/nav_hierarchy.py +++ b/packages/shared-python/shared/services/retrieval/nav/nav_hierarchy.py @@ -241,14 +241,24 @@ 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 callable(fn) and str(getattr(provider, "doc_id", "")) == str(doc_id): + 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 callable(fn) and str(getattr(provider, "doc_id", "")) == str(doc_id): + 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()