From 66c71250666a6a5a0c282570d9d446b013a4edbf Mon Sep 17 00:00:00 2001 From: suguanYang Date: Fri, 28 Aug 2026 13:19:45 +0800 Subject: [PATCH 1/3] perf: batch map-nav corpus scoring and reads --- ...etrieval_lazy_snapshot_quality_contract.py | 158 +++++++++++++++++- .../services/retrieval/nav/knowhere_hybrid.py | 85 ++++++++-- .../services/retrieval/nav/nav_hierarchy.py | 9 + .../services/retrieval/nav/nav_knowhere.py | 131 ++++++++++++++- .../services/retrieval/nav/nav_map_scores.py | 132 +++++++++++++-- .../services/retrieval/nav/nav_orchestrate.py | 67 ++++++-- 6 files changed, 533 insertions(+), 49 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 cfb82c9e6..ae6d78ed3 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 @@ -1,7 +1,7 @@ from __future__ import annotations from dataclasses import dataclass -from collections.abc import Sequence +from collections.abc import Mapping, Sequence from typing import Any from shared.services.retrieval.nav.nav_hierarchy import ProviderToolSpace @@ -16,11 +16,13 @@ from shared.services.retrieval.nav.nav_map_scores import ( build_score_units, compute_corpus_map_and_unit_scores, + compute_corpus_map_and_unit_scores_many, ) from shared.services.retrieval.nav.knowhere_hybrid import ( ScoreUnitRow, score_rows_hybrid_all, score_unit_stream_hybrid_all, + score_unit_stream_hybrid_many, ) @@ -28,6 +30,21 @@ class _FakeChunkStore: units_by_section: dict[str, list[UnitRow]] document_loads: int = 0 + batch_loads: int = 0 + + def load_documents_units( + self, + section_ids_by_document: Mapping[str, Sequence[str]], + ) -> dict[str, list[UnitRow]]: + self.batch_loads += 1 + return { + document_id: [ + unit + for section_id in section_ids + for unit in self.units_by_section.get(section_id, ()) + ] + for document_id, section_ids in section_ids_by_document.items() + } def load_document_units( self, @@ -118,6 +135,63 @@ def _providers() -> tuple[ProviderToolSpace, ProviderToolSpace, _FakeChunkStore] return ProviderToolSpace(eager), ProviderToolSpace(lazy), store +def _multi_document_providers() -> tuple[ + ProviderToolSpace, + ProviderToolSpace, + _FakeChunkStore, +]: + first_sections = [ + SectionRow("root-a", None, "Root A", "Root A", 0, "", 0), + SectionRow("leaf-a", "root-a", "Root A / Leaf A", "Leaf A", 1, "", 1), + ] + second_sections = [ + SectionRow("root-b", None, "Root B", "Root B", 0, "", 0), + SectionRow("leaf-b", "root-b", "Root B / Leaf B", "Leaf B", 1, "", 1), + ] + first_unit = UnitRow("chunk-a", "leaf-a", "text", "alpha evidence", 1) + second_unit = UnitRow("chunk-b", "leaf-b", "text", "beta evidence", 1) + eager = NamespaceKnowhereProvider( + [ + KnowhereProvider( + doc_id="doc-a", + sections=first_sections, + units=[first_unit], + ), + KnowhereProvider( + doc_id="doc-b", + sections=second_sections, + units=[second_unit], + ), + ], + titles={"doc-a": "Document A", "doc-b": "Document B"}, + ) + store = _FakeChunkStore( + { + "leaf-a": [first_unit], + "leaf-b": [second_unit], + } + ) + lazy = NamespaceKnowhereProvider( + [ + LazyKnowhereProvider( + doc_id="doc-a", + sections=first_sections, + chunk_store=store, + known_chunk_ids=[first_unit.chunk_id], + ), + LazyKnowhereProvider( + doc_id="doc-b", + sections=second_sections, + chunk_store=store, + known_chunk_ids=[second_unit.chunk_id], + ), + ], + titles={"doc-a": "Document A", "doc-b": "Document B"}, + chunk_owner_by_id={"chunk-a": "doc-a", "chunk-b": "doc-b"}, + ) + return ProviderToolSpace(eager), ProviderToolSpace(lazy), store + + def test_lazy_provider_preserves_score_units_and_scores() -> None: eager, lazy, store = _providers() @@ -207,6 +281,88 @@ def test_streaming_scorer_preserves_duplicate_id_eager_semantics() -> None: assert score_unit_stream_hybrid_all(lambda: rows, "alpha beta") == eager_scores +def test_streaming_scorer_scores_multiple_queries_with_one_corpus_read() -> None: + rows: list[ScoreUnitRow] = [ + { + "chunk_id": "unit-a", + "path_search_text": "root alpha", + "content_search_text": "alpha alpha evidence", + "term_search_text": "alpha alpha evidence root", + }, + { + "chunk_id": "unit-b", + "path_search_text": "root beta", + "content_search_text": "beta evidence", + "term_search_text": "beta evidence root", + }, + { + "chunk_id": "unit-c", + "path_search_text": "root common", + "content_search_text": "common evidence", + "term_search_text": "common evidence root", + }, + ] + queries: list[str] = ["alpha evidence", "beta evidence"] + expected = { + query: score_unit_stream_hybrid_all(lambda: rows, query) + for query in queries + } + read_count: int = 0 + + def unit_factory() -> Sequence[ScoreUnitRow]: + nonlocal read_count + read_count += 1 + return rows + + assert score_unit_stream_hybrid_many(unit_factory, queries) == expected + assert read_count == 1 + + +def test_corpus_map_scores_multiple_queries_with_one_lazy_load() -> None: + eager, lazy, store = _providers() + queries: list[str] = ["alpha retrieval", "supporting image"] + expected = { + query: compute_corpus_map_and_unit_scores( + eager, + doc_ids=["doc"], + query=query, + ) + for query in queries + } + + store.document_loads = 0 + store.batch_loads = 0 + actual = compute_corpus_map_and_unit_scores_many( + lazy, + doc_ids=["doc"], + queries=queries, + ) + + assert actual == expected + assert store.batch_loads == 1 + assert store.document_loads == 0 + + +def test_corpus_map_batches_multiple_documents_without_score_drift() -> None: + eager, lazy, store = _multi_document_providers() + queries: list[str] = ["alpha evidence", "beta evidence"] + expected = compute_corpus_map_and_unit_scores_many( + eager, + doc_ids=["doc-a", "doc-b"], + queries=queries, + ) + + actual = compute_corpus_map_and_unit_scores_many( + lazy, + doc_ids=["doc-a", "doc-b"], + queries=queries, + ) + + assert actual == expected + assert store.batch_loads == 1 + assert store.document_loads == 0 + + def test_native_chunk_store_strips_async_driver_from_database_url( monkeypatch: Any, ) -> None: diff --git a/packages/shared-python/shared/services/retrieval/nav/knowhere_hybrid.py b/packages/shared-python/shared/services/retrieval/nav/knowhere_hybrid.py index 7d8b195bc..a647566a5 100644 --- a/packages/shared-python/shared/services/retrieval/nav/knowhere_hybrid.py +++ b/packages/shared-python/shared/services/retrieval/nav/knowhere_hybrid.py @@ -619,12 +619,32 @@ def score_unit_stream_hybrid_all( weighted-RRF implementation, but keeps only token statistics, identifiers, and final scores between bounded provider reads. """ - query_tokens = tokenize_query_for_ranker(query) + return score_unit_stream_hybrid_many(unit_factory, [query]).get(query, {}) + + +def score_unit_stream_hybrid_many( + unit_factory: Callable[[], Iterable[ScoreUnitRow]], + queries: Sequence[str], +) -> Dict[str, Dict[str, float]]: + """Score several queries exactly while reading the corpus only once.""" + unique_queries = list(dict.fromkeys(str(query) for query in queries)) + if not unique_queries: + return {} + + query_tokens_by_query = { + query: tokenize_query_for_ranker(query) for query in unique_queries + } + query_token_set = { + token + for query_tokens in query_tokens_by_query.values() + for token in query_tokens + } + query_lower_by_query = { + query: query.lower().strip() for query in unique_queries + } path_stats = _StreamingBm25Stats.empty() content_stats = _StreamingBm25Stats.empty() - units: List[_StreamingUnit] = [] - query_token_set = set(query_tokens) - query_lower = query.lower().strip() + units: List[_StreamingManyUnit] = [] for row in unit_factory(): unit_id = str(row.get("chunk_id") or "").strip() if not unit_id: @@ -635,17 +655,24 @@ def score_unit_stream_hybrid_all( content_stats.observe(content_tokens) path_frequencies = Counter(path_tokens) content_frequencies = Counter(content_tokens) - term_score = 0.0 - if query_lower: - haystack = str(row.get("term_search_text") or "").lower() - if query_lower in haystack: - term_score = 100.0 - else: - hit_count = sum(1 for token in query_tokens if token in haystack) - if hit_count > 0: - term_score = float(hit_count) + haystack = str(row.get("term_search_text") or "").lower() + term_scores: List[float] = [] + for query in unique_queries: + query_lower = query_lower_by_query[query] + query_tokens = query_tokens_by_query[query] + term_score = 0.0 + if query_lower: + if query_lower in haystack: + term_score = 100.0 + else: + hit_count = sum( + 1 for token in query_tokens if token in haystack + ) + if hit_count > 0: + term_score = float(hit_count) + term_scores.append(term_score) units.append( - _StreamingUnit( + _StreamingManyUnit( unit_id=unit_id, path_length=len(path_tokens), content_length=len(content_tokens), @@ -659,11 +686,31 @@ def score_unit_stream_hybrid_all( for token in query_token_set if content_frequencies[token] }, - term_score=term_score, + term_scores=tuple(term_scores), ) ) path_stats.finalize() content_stats.finalize() + return { + query: _score_streaming_units( + units, + path_stats=path_stats, + content_stats=content_stats, + query_tokens=query_tokens_by_query[query], + query_index=index, + ) + for index, query in enumerate(unique_queries) + } + + +def _score_streaming_units( + units: Sequence["_StreamingManyUnit"], + *, + path_stats: "_StreamingBm25Stats", + content_stats: "_StreamingBm25Stats", + query_tokens: List[str], + query_index: int, +) -> Dict[str, float]: path_by_id: Dict[str, float] = {} content_by_id: Dict[str, float] = {} term_by_id: Dict[str, float] = {} @@ -677,7 +724,9 @@ def score_unit_stream_hybrid_all( ) path_by_id[unit.unit_id] = path_score content_by_id[unit.unit_id] = content_score - term_by_id[unit.unit_id] = unit.term_score + term_by_id[unit.unit_id] = float( + unit.term_scores[query_index] if query_index < len(unit.term_scores) else 0.0 + ) path_rows = [ (score, unit_id) @@ -717,13 +766,13 @@ def score_unit_stream_hybrid_all( @dataclass(frozen=True) -class _StreamingUnit: +class _StreamingManyUnit: unit_id: str path_length: int content_length: int path_frequencies: Mapping[str, int] content_frequencies: Mapping[str, int] - term_score: float + term_scores: Tuple[float, ...] class _StreamingBm25Stats: 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 0609cbc23..e6c8cafa3 100644 --- a/packages/shared-python/shared/services/retrieval/nav/nav_hierarchy.py +++ b/packages/shared-python/shared/services/retrieval/nav/nav_hierarchy.py @@ -249,6 +249,15 @@ def prefetch_document_units(self, doc_id: str) -> None: if str(getattr(provider, "doc_id", "")) == str(doc_id): fn() + def prefetch_document_units_batch(self, doc_ids: Sequence[str]) -> None: + """Forward a provider's bounded multi-document prefetch capability.""" + fn = getattr(self._provider, "prefetch_document_units_batch", None) + if callable(fn): + fn(doc_ids) + return + for doc_id in doc_ids: + self.prefetch_document_units(str(doc_id)) + def release_document_units(self, doc_id: str) -> None: """Forward release of one document's prefetched payloads.""" provider = self._provider 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 a1387a645..b3b39f2bc 100644 --- a/packages/shared-python/shared/services/retrieval/nav/nav_knowhere.py +++ b/packages/shared-python/shared/services/retrieval/nav/nav_knowhere.py @@ -22,7 +22,7 @@ import os from dataclasses import dataclass, field -from typing import Any, Callable, Dict, Iterable, List, Optional, Protocol, Sequence, Set, Tuple +from typing import Any, Callable, Dict, Iterable, List, Mapping, Optional, Protocol, Sequence, Set, Tuple from .nav_address import NavLevel from .nav_hierarchy import NodeMeta @@ -141,6 +141,12 @@ def knowhere_database_url() -> str: class ChunkStore(Protocol): + def load_documents_units( + self, + section_ids_by_document: Mapping[str, Sequence[str]], + ) -> Dict[str, List[UnitRow]]: + raise NotImplementedError + def load_document_units( self, document_id: str, @@ -268,6 +274,86 @@ def load_document_units( finally: cur.close() + def load_documents_units( + self, + section_ids_by_document: Mapping[str, Sequence[str]], + ) -> Dict[str, List[UnitRow]]: + """Load a bounded group of revisions with keyset-paged SQL queries.""" + requested = [ + (str(document_id).strip(), self._revisions.get(str(document_id).strip())) + for document_id in section_ids_by_document + if str(document_id).strip() + ] + revisions = [ + (document_id, str(job_result_id)) + for document_id, job_result_id in requested + if job_result_id + ] + if not revisions: + return {} + + values_sql = ", ".join(["(%s, %s)"] * len(revisions)) + params: List[object] = [ + value + for revision in revisions + for value in revision + ] + units_by_document: Dict[str, List[UnitRow]] = { + document_id: [] for document_id, _job_result_id in revisions + } + last_key: Optional[Tuple[str, str, int, str, str]] = None + while True: + page_params = list(params) + keyset_sql = "" + if last_key is not None: + keyset_sql = ( + " AND (chunks.document_id, chunks.job_result_id, " + "chunks.sort_order, chunks.chunk_id, chunks.id) > " + "(%s, %s, %s, %s, %s)" + ) + page_params.extend(last_key) + page_params.append(10_000) + cur = self._connection().cursor() + try: + cur.execute( + "SELECT chunks.document_id, chunks.job_result_id, chunks.chunk_id, " + "chunks.section_id, chunks.chunk_type, chunks.content, " + "chunks.sort_order, chunks.source_chunk_path, chunks.file_path, " + "chunks.chunk_metadata, chunks.id " + "FROM document_chunks AS chunks " + f"JOIN (VALUES {values_sql}) AS revisions(document_id, job_result_id) " + "ON chunks.document_id = revisions.document_id " + "AND chunks.job_result_id = revisions.job_result_id" + f"{keyset_sql} " + "ORDER BY chunks.document_id, chunks.job_result_id, " + "chunks.sort_order, chunks.chunk_id, chunks.id LIMIT %s", + page_params, + ) + rows = cur.fetchall() + finally: + cur.close() + if not rows: + break + for row in rows: + document_id = str(row[0]) + section_id = str(row[3]) if row[3] else None + if section_id and (document_id, section_id) in self._excluded_sections: + continue + units_by_document.setdefault(document_id, []).append( + _unit_from_row(row[2:10]) + ) + last = rows[-1] + last_key = ( + str(last[0]), + str(last[1]), + int(last[6] or 0), + str(last[2] or ""), + str(last[10] or ""), + ) + if len(rows) < 10_000: + break + return units_by_document + def close(self) -> None: if self._conn is not None: self._conn.close() @@ -608,6 +694,14 @@ def prefetch_document_units(self) -> None: section_ids, self._remounted_assets_by_section, ) + self.install_prefetched_document_units(loaded) + + def install_prefetched_document_units( + self, + loaded: Sequence[UnitRow], + ) -> None: + """Install rows fetched by either a document or corpus-group query.""" + section_ids = list(self._sections) by_section: Dict[str, List[UnitRow]] = {} for unit in loaded: sid = str(unit.section_id or "").strip() @@ -872,6 +966,41 @@ def prefetch_document_units(self, doc_id: str) -> None: if callable(prefetch): prefetch() + def prefetch_document_units_batch(self, doc_ids: Sequence[str]) -> None: + """Load a bounded document group with one query per shared chunk store.""" + providers = [ + self._docs[doc_id] + for raw_doc_id in doc_ids + if (doc_id := str(raw_doc_id).strip()) in self._docs + ] + providers_by_store: Dict[int, List[LazyKnowhereProvider]] = {} + stores_by_id: Dict[int, ChunkStore] = {} + for provider in providers: + if not isinstance(provider, LazyKnowhereProvider): + continue + store = provider._chunk_store + store_id = id(store) + stores_by_id[store_id] = store + providers_by_store.setdefault(store_id, []).append(provider) + + for store_id, lazy_providers in providers_by_store.items(): + store = stores_by_id[store_id] + batch_loader = getattr(store, "load_documents_units", None) + if not callable(batch_loader): + for provider in lazy_providers: + provider.prefetch_document_units() + continue + loaded_by_document = batch_loader( + { + provider.doc_id: list(provider._sections) + for provider in lazy_providers + } + ) + for provider in lazy_providers: + provider.install_prefetched_document_units( + loaded_by_document.get(provider.doc_id, ()) + ) + def release_document_units(self, doc_id: str) -> None: provider = self._docs.get(str(doc_id).strip()) release = getattr(provider, "release_document_units", None) 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 0223e2320..5909b83de 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 @@ -9,9 +9,13 @@ build_path_search_text, build_term_search_text, score_rows_hybrid_all, - score_unit_stream_hybrid_all, + score_unit_stream_hybrid_many, ) +# Keep one bulk read bounded when a namespace contains a very large document, +# while still replacing the per-document N+1 access pattern. +_CORPUS_PREFETCH_GROUP_SIZE = 8 + def _children_ids(ts: Any, section_id: str, doc_id: str) -> List[str]: children_fn = getattr(ts, "_children_for_section_path", None) @@ -362,6 +366,26 @@ def compute_corpus_map_and_unit_scores( All documents share one BM25 corpus, path/content normalization, channel ranking, and RRF pass. Document-level scores are keyed by bare ``document_id``. """ + return compute_corpus_map_and_unit_scores_many( + ts, + doc_ids=doc_ids, + queries=[query], + namespace=namespace, + ).get(query, ({}, {})) + + +def compute_corpus_map_and_unit_scores_many( + ts: Any, + *, + doc_ids: Sequence[str], + queries: Sequence[str], + namespace: Optional[str] = None, +) -> Dict[str, Tuple[Dict[str, float], Dict[str, float]]]: + """Globally score several queries with one replay of the corpus units.""" + unique_queries = list(dict.fromkeys(str(query) for query in queries)) + if not unique_queries: + return {} + valid_doc_ids: List[str] = [] seen_doc_ids: Set[str] = set() for raw in doc_ids: @@ -383,10 +407,37 @@ def compute_corpus_map_and_unit_scores( tree_by_doc[doc_id] = (children_map, leaves, titles) def unit_factory() -> Iterator[ScoreUnitRow]: + prefetch_batch = getattr(ts, "prefetch_document_units_batch", None) + release = getattr(ts, "release_document_units", None) + if callable(prefetch_batch): + for group_start in range( + 0, + len(valid_doc_ids), + _CORPUS_PREFETCH_GROUP_SIZE, + ): + document_group = valid_doc_ids[ + group_start : group_start + _CORPUS_PREFETCH_GROUP_SIZE + ] + prefetch_batch(document_group) + try: + for document_id in document_group: + 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, + ) + finally: + if callable(release): + for document_id in document_group: + release(document_id) + return + for document_id in valid_doc_ids: children_map, leaves, titles = tree_by_doc[document_id] prefetch = getattr(ts, "prefetch_document_units", None) - release = getattr(ts, "release_document_units", None) if callable(prefetch): prefetch(document_id) try: @@ -401,21 +452,27 @@ def unit_factory() -> Iterator[ScoreUnitRow]: if callable(release): release(document_id) - unit_scores = score_unit_stream_hybrid_all(unit_factory, query) - - map_scores: Dict[str, float] = {} - for doc_id in valid_doc_ids: - children_map, leaves, _titles = tree_by_doc[doc_id] - doc_map_scores = _pool_unit_scores_to_tree( - children_map, leaves, unit_scores - ) - map_scores.update(doc_map_scores) - doc_max = max( - (float(value) for value in doc_map_scores.values()), - default=0.0, - ) - map_scores[doc_id] = doc_max - return map_scores, unit_scores + unit_scores_by_query = score_unit_stream_hybrid_many( + unit_factory, + unique_queries, + ) + results: Dict[str, Tuple[Dict[str, float], Dict[str, float]]] = {} + for query in unique_queries: + unit_scores = unit_scores_by_query.get(query, {}) + map_scores: Dict[str, float] = {} + for doc_id in valid_doc_ids: + children_map, leaves, _titles = tree_by_doc[doc_id] + doc_map_scores = _pool_unit_scores_to_tree( + children_map, leaves, unit_scores + ) + map_scores.update(doc_map_scores) + doc_max = max( + (float(value) for value in doc_map_scores.values()), + default=0.0, + ) + map_scores[doc_id] = doc_max + results[query] = (map_scores, unit_scores) + return results def unit_id_to_section_id(unit_id: str) -> str: @@ -474,3 +531,44 @@ def relight_map_for_query( ts, doc_ids=doc_ids, query=query ) return map_scores, unit_scores, select_map_highlights(unit_scores, k=int(top_k)) + + +def relight_maps_for_queries( + ts: Any, + *, + doc_id: str, + queries: Sequence[str], + top_k: int = 6, +) -> Dict[str, Tuple[Dict[str, float], Dict[str, float], List[str]]]: + """Re-score a shared map for several queries with one corpus replay.""" + unique_queries = list(dict.fromkeys(str(query) for query in queries)) + if not unique_queries: + return {} + doc = str(doc_id or "").strip() + if doc: + return { + query: relight_map_for_query( + ts, + doc_id=doc, + query=query, + top_k=top_k, + ) + for query in unique_queries + } + + doc_ids = [str(value) for value in (ts.document_ids() or ()) if str(value).strip()] + if not doc_ids: + return {} + scored = compute_corpus_map_and_unit_scores_many( + ts, + doc_ids=doc_ids, + queries=unique_queries, + ) + return { + query: ( + map_scores, + unit_scores, + select_map_highlights(unit_scores, k=int(top_k)), + ) + for query, (map_scores, unit_scores) in scored.items() + } diff --git a/packages/shared-python/shared/services/retrieval/nav/nav_orchestrate.py b/packages/shared-python/shared/services/retrieval/nav/nav_orchestrate.py index 1e30b8f4c..91f45e1c5 100644 --- a/packages/shared-python/shared/services/retrieval/nav/nav_orchestrate.py +++ b/packages/shared-python/shared/services/retrieval/nav/nav_orchestrate.py @@ -126,6 +126,16 @@ def _unbound_retrieval_query(subgoal: Subgoal) -> str: return raw or (subgoal.need or "").strip() or subgoal.retrieval_query +def _resolve_subgoal_query(state: NavState, subgoal: Subgoal) -> str: + query = bind_slots(subgoal.retrieval_query, state.slot_bindings) + if unbound_slots(query): + query = _unbound_retrieval_query(subgoal) + refined = str( + (state.subgoal_refined_queries or {}).get(subgoal.id) or "" + ).strip() + return refined or query + + def _run_navigate_for_query( ts: Any, state: NavState, @@ -180,6 +190,7 @@ def _relit_map( config: NavConfig, *, query: str, + prepared: Optional[Tuple[Dict[str, float], Dict[str, float], List[str]]] = None, ) -> Iterator[None]: """Score the shared map against the harvest ``query`` for one call. @@ -189,9 +200,9 @@ def _relit_map( the query the policy is told to pursue. Scoring failures degrade to the episode lighting. """ - relit: Optional[Tuple[Dict[str, float], Dict[str, float], List[str]]] = None + relit = prepared q = (query or "").strip() - if q: + if relit is None and q: try: from .nav_map_scores import relight_map_for_query @@ -224,6 +235,10 @@ def _execute_subgoal_harvest_once( subgoal: Subgoal, *, steps_out: Optional[List[Any]], + retrieval_query: Optional[str] = None, + prepared_relight: Optional[ + Tuple[Dict[str, float], Dict[str, float], List[str]] + ] = None, ) -> Dict[str, Any]: """One harvest() call for this subgoal this wave — no internal retry loop. @@ -232,22 +247,21 @@ def _execute_subgoal_harvest_once( """ from .nav_harvest import harvest - rq = bind_slots(subgoal.retrieval_query, state.slot_bindings) - if unbound_slots(rq): - # F1: deps may be "settled" (satisfied or dropped) without ever - # producing this subgoal's referenced slot — degrade to a query with - # the unresolved {{...}} braces stripped rather than stalling. - rq = _unbound_retrieval_query(subgoal) + rq = retrieval_query or _resolve_subgoal_query(state, subgoal) refined = str((state.subgoal_refined_queries or {}).get(subgoal.id) or "").strip() - if refined: - rq = refined _set_focus(state, subgoal, rq) # Always enter at namespace/document root; prior dead-ends stay hidden via # subgoal_dismissed_section_ids so the next harvest sees siblings instead. before_sections = set(state.collected_section_ids) before_explicit = set(state.explicit_collect_ids) before_len = len(state.collected) - with _relit_map(ts, state, config, query=rq): + with _relit_map( + ts, + state, + config, + query=rq, + prepared=prepared_relight, + ): harvest_result = harvest( ts, state, @@ -433,10 +447,39 @@ def execute_plan( by_id = {s.id: s for s in plan.subgoals} outputs: List[Dict[str, Any]] = [] + query_by_subgoal = { + sid: _resolve_subgoal_query(state, by_id[sid]) for sid in ready + } + prepared_relights: Dict[ + str, + Tuple[Dict[str, float], Dict[str, float], List[str]], + ] = {} + try: + from .nav_map_scores import relight_maps_for_queries + + prepared_relights = relight_maps_for_queries( + ts, + doc_id=state.doc_id, + queries=list(query_by_subgoal.values()), + top_k=int(config.collect_top_k), + ) + except Exception: + prepared_relights = {} def _run_one(sid: str, working_state: NavState, out_steps: Optional[List[Any]]) -> Dict[str, Any]: + query = query_by_subgoal[sid] + prepared = prepared_relights.get(query) + if prepared is not None and not prepared[0]: + prepared = None return _execute_subgoal_harvest_once( - ts, working_state, config, plan, by_id[sid], steps_out=out_steps + ts, + working_state, + config, + plan, + by_id[sid], + steps_out=out_steps, + retrieval_query=query, + prepared_relight=prepared, ) # Serial wave execution (parallel fan-out retired with ThreadPoolExecutor). From bb846aefef92244f3e75c254562f52603afbb5bb Mon Sep 17 00:00:00 2001 From: suguanYang Date: Fri, 28 Aug 2026 18:05:57 +0800 Subject: [PATCH 2/3] perf: persist map-nav scoring index --- .../1d2e3f4a5b6c_add_document_map_units.py | 122 +++++ apps/api/scripts/backfill_map_unit_indexes.py | 97 ++++ .../test_retrieval_map_unit_index_contract.py | 423 ++++++++++++++++++ .../shared/models/database/__init__.py | 6 + .../shared/models/database/document.py | 157 ++++++- .../services/retrieval/map_unit_index.py | 163 +++++++ .../services/retrieval/nav/knowhere_hybrid.py | 220 ++++++--- .../services/retrieval/nav/nav_hierarchy.py | 40 +- .../services/retrieval/nav/nav_knowhere.py | 363 ++++++++++++++- .../services/retrieval/nav/nav_map_scores.py | 52 ++- .../services/retrieval/publication_content.py | 24 +- 11 files changed, 1560 insertions(+), 107 deletions(-) create mode 100644 apps/api/alembic/versions/1d2e3f4a5b6c_add_document_map_units.py create mode 100644 apps/api/scripts/backfill_map_unit_indexes.py create mode 100644 apps/api/tests/contract/test_retrieval_map_unit_index_contract.py create mode 100644 packages/shared-python/shared/services/retrieval/map_unit_index.py diff --git a/apps/api/alembic/versions/1d2e3f4a5b6c_add_document_map_units.py b/apps/api/alembic/versions/1d2e3f4a5b6c_add_document_map_units.py new file mode 100644 index 000000000..79aa7583c --- /dev/null +++ b/apps/api/alembic/versions/1d2e3f4a5b6c_add_document_map_units.py @@ -0,0 +1,122 @@ +"""Add revision-pinned map-nav score units.""" + +from __future__ import annotations + +import sqlalchemy as sa +from alembic import op + +revision = "1d2e3f4a5b6c" +down_revision = "0c1d2e3f4a5b" +branch_labels = None +depends_on = None + + +def upgrade() -> None: + bind = op.get_bind() + inspector = sa.inspect(bind) + if not inspector.has_table("document_map_unit_indexes"): + op.create_table( + "document_map_unit_indexes", + sa.Column("id", sa.String(length=100), nullable=False), + sa.Column("document_id", sa.String(length=36), nullable=False), + sa.Column("job_result_id", sa.String(length=36), nullable=False), + sa.Column("format_version", sa.Integer(), nullable=False), + sa.Column("unit_count", sa.Integer(), nullable=False), + sa.Column("token_count", sa.Integer(), nullable=False), + sa.Column("created_at", sa.DateTime(), nullable=False), + sa.ForeignKeyConstraint( + ["document_id"], ["documents.document_id"], ondelete="CASCADE" + ), + sa.ForeignKeyConstraint( + ["job_result_id"], ["job_results.id"], ondelete="CASCADE" + ), + sa.PrimaryKeyConstraint("id"), + sa.UniqueConstraint( + "document_id", + "job_result_id", + name="uq_document_map_unit_indexes_revision", + ), + ) + if not inspector.has_table("document_map_units"): + op.create_table( + "document_map_units", + sa.Column("id", sa.String(length=160), nullable=False), + sa.Column("document_id", sa.String(length=36), nullable=False), + sa.Column("job_result_id", sa.String(length=36), nullable=False), + sa.Column("unit_id", sa.String(length=128), nullable=False), + sa.Column("section_id", sa.String(length=36), nullable=False), + sa.Column("unit_kind", sa.String(length=32), nullable=False), + sa.Column("path_token_count", sa.Integer(), nullable=False), + sa.Column("content_token_count", sa.Integer(), nullable=False), + sa.Column("term_search_text_lower", sa.Text(), nullable=False), + sa.Column("sort_order", sa.Integer(), nullable=False), + sa.Column("created_at", sa.DateTime(), nullable=False), + sa.ForeignKeyConstraint( + ["document_id"], ["documents.document_id"], ondelete="CASCADE" + ), + sa.ForeignKeyConstraint( + ["job_result_id"], ["job_results.id"], ondelete="CASCADE" + ), + sa.PrimaryKeyConstraint("id"), + ) + if not inspector.has_table("document_map_unit_tokens"): + op.create_table( + "document_map_unit_tokens", + sa.Column("id", sa.String(length=36), nullable=False), + sa.Column("map_unit_id", sa.String(length=160), nullable=False), + sa.Column("channel", sa.String(length=16), nullable=False), + sa.Column("token", sa.Text(), nullable=False), + sa.Column("token_hash", sa.String(length=64), nullable=False), + sa.Column("frequency", sa.Integer(), nullable=False), + sa.ForeignKeyConstraint( + ["map_unit_id"], ["document_map_units.id"], ondelete="CASCADE" + ), + sa.PrimaryKeyConstraint("id"), + ) + inspector = sa.inspect(bind) + indexes = { + item["name"] for item in inspector.get_indexes("document_map_unit_indexes") + } + if "idx_document_map_unit_indexes_revision" not in indexes: + op.create_index( + "idx_document_map_unit_indexes_revision", + "document_map_unit_indexes", + ["document_id", "job_result_id"], + ) + indexes = {item["name"] for item in inspector.get_indexes("document_map_units")} + if "idx_document_map_units_revision_order" not in indexes: + op.create_index( + "idx_document_map_units_revision_order", + "document_map_units", + ["document_id", "job_result_id", "sort_order", "unit_id"], + ) + if "idx_document_map_units_section" not in indexes: + op.create_index( + "idx_document_map_units_section", "document_map_units", ["section_id"] + ) + indexes = { + item["name"] for item in inspector.get_indexes("document_map_unit_tokens") + } + if "idx_document_map_unit_tokens_lookup" not in indexes: + op.create_index( + "idx_document_map_unit_tokens_lookup", + "document_map_unit_tokens", + ["channel", "token_hash", "map_unit_id"], + ) + if "idx_document_map_unit_tokens_unit" not in indexes: + op.create_index( + "idx_document_map_unit_tokens_unit", + "document_map_unit_tokens", + ["map_unit_id", "channel"], + ) + + +def downgrade() -> None: + existing_tables = set(sa.inspect(op.get_bind()).get_table_names()) + for table_name in ( + "document_map_unit_tokens", + "document_map_units", + "document_map_unit_indexes", + ): + if table_name in existing_tables: + op.drop_table(table_name) diff --git a/apps/api/scripts/backfill_map_unit_indexes.py b/apps/api/scripts/backfill_map_unit_indexes.py new file mode 100644 index 000000000..78a88f6c0 --- /dev/null +++ b/apps/api/scripts/backfill_map_unit_indexes.py @@ -0,0 +1,97 @@ +"""Backfill persisted MAP-NAV lexical indexes for existing revisions. + +The migration creates empty derived tables intentionally. Run this command +after deployment with ``--apply`` so each revision is rebuilt and committed +independently; without ``--apply`` it is a read-only inventory. +""" + +# ruff: noqa: E402 + +from __future__ import annotations + +import argparse +import os +import sys +from pathlib import Path + + +def _bootstrap_python_path() -> None: + api_root = Path(__file__).resolve().parents[1] + repo_root = api_root.parents[1] + shared_root = repo_root / "packages" / "shared-python" + for path in (api_root, shared_root): + value = os.fspath(path) + if value not in sys.path: + sys.path.insert(0, value) + + +_bootstrap_python_path() + +from sqlalchemy import select + +from shared.core.database_sync import get_sync_session_factory +from shared.models.database.document import Document +from shared.services.retrieval.map_unit_index import replace_document_map_units +from shared.services.retrieval.publication_models import DocumentPublicationScope + + +def _build_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser( + description="Backfill MAP-NAV indexes for current document revisions." + ) + parser.add_argument( + "--apply", + action="store_true", + help="Build and commit each current revision index.", + ) + parser.add_argument("--document-id", default="", help="Limit the backfill to one document.") + return parser + + +def _load_documents(document_id: str) -> list[Document]: + session_factory = get_sync_session_factory() + with session_factory() as db: + statement = select(Document).where(Document.current_job_result_id.is_not(None)) + normalized_document_id = document_id.strip() + if normalized_document_id: + statement = statement.where(Document.document_id == normalized_document_id) + return list(db.scalars(statement).all()) + + +def backfill_map_unit_indexes(*, apply: bool, document_id: str = "") -> int: + documents = _load_documents(document_id) + if not apply: + for document in documents: + print(f"would backfill document={document.document_id} revision={document.current_job_result_id}") + return len(documents) + + session_factory = get_sync_session_factory() + for document in documents: + job_result_id = document.current_job_result_id + if not job_result_id: + continue + scope = DocumentPublicationScope( + user_id=document.user_id, + namespace=document.namespace, + document_id=document.document_id, + job_result_id=job_result_id, + source_file_name=str(document.source_file_name or ""), + ) + with session_factory() as db: + replace_document_map_units(db, scope=scope) + db.commit() + print(f"backfilled document={document.document_id} revision={job_result_id}") + return len(documents) + + +def main() -> None: + arguments = _build_parser().parse_args() + count = backfill_map_unit_indexes( + apply=bool(arguments.apply), document_id=str(arguments.document_id) + ) + action = "backfilled" if arguments.apply else "found" + print(f"{action} revisions={count}") + + +if __name__ == "__main__": + main() diff --git a/apps/api/tests/contract/test_retrieval_map_unit_index_contract.py b/apps/api/tests/contract/test_retrieval_map_unit_index_contract.py new file mode 100644 index 000000000..2f2027b01 --- /dev/null +++ b/apps/api/tests/contract/test_retrieval_map_unit_index_contract.py @@ -0,0 +1,423 @@ +from __future__ import annotations + +from collections.abc import Callable, Mapping, Sequence +from contextlib import AbstractAsyncContextManager +from uuid import uuid4 + +from httpx import AsyncClient +from sqlalchemy import delete, select + +from shared.models.database.document import ( + DocumentMapUnit, + DocumentMapUnitIndex, + DocumentMapUnitToken, +) +from shared.services.retrieval.nav.nav_hierarchy import ProviderToolSpace +from shared.services.retrieval.nav.nav_map_scores import ( + build_score_units, + compute_corpus_map_and_unit_scores, +) +from shared.services.retrieval.nav.nav_knowhere import ( + KnowhereProvider, + LazyKnowhereProvider, + NamespaceKnowhereProvider, + ReadOnlyChunkStore, + SectionRow, + UnitRow, +) +from shared.services.retrieval.nav_snapshot import load_nav_snapshot +from shared.services.retrieval.publication_content import ( + replace_document_revision_content, +) +from shared.services.retrieval.publication_models import DocumentPublicationScope +from tests.support.contract_database import ContractDatabase +from tests.support.retrieval_snapshot_support import contract_db_session + +_USER_ID = "local-dev-user" + + +class _IncompleteIndexStore: + """Minimal lazy store whose incomplete index forces legacy scoring.""" + + def __init__(self, units_by_section: Mapping[str, Sequence[UnitRow]]) -> None: + self.units_by_section = { + str(section_id): list(units) + for section_id, units in units_by_section.items() + } + self.persisted_loads = 0 + self.batch_loads = 0 + + def load_persisted_score_corpus( + self, + document_ids: Sequence[str], + allowed_section_ids_by_document: Mapping[str, Sequence[str]], + queries: Sequence[str], + ) -> None: + del document_ids, allowed_section_ids_by_document, queries + self.persisted_loads += 1 + return None + + def load_documents_units( + self, + section_ids_by_document: Mapping[str, Sequence[str]], + ) -> dict[str, list[UnitRow]]: + self.batch_loads += 1 + return { + str(document_id): [ + unit + for section_id in section_ids + for unit in self.units_by_section.get(str(section_id), ()) + ] + for document_id, section_ids in section_ids_by_document.items() + } + + def load_document_units( + self, + document_id: str, + section_ids: Sequence[str], + extra_chunk_ids_by_section: Mapping[str, Sequence[str]] | None = None, + ) -> list[UnitRow]: + del document_id, extra_chunk_ids_by_section + return [ + unit + for section_id in section_ids + for unit in self.units_by_section.get(str(section_id), ()) + ] + + def load_section_units( + self, + document_id: str, + section_id: str, + extra_chunk_ids: Sequence[str] = (), + ) -> list[UnitRow]: + del document_id, extra_chunk_ids + return list(self.units_by_section.get(str(section_id), ())) + + def close(self) -> None: + return None + + +async def test_published_map_units_preserve_scores_without_chunk_payload_reads( + developer_api_client_factory: Callable[ + [], AbstractAsyncContextManager[AsyncClient] + ], + monkeypatch, +) -> None: + identifier = uuid4().hex[:8] + namespace = f"map-unit-index-{identifier}" + document_id = f"doc_map_{identifier}" + job_id = f"job_map_{identifier}" + job_result_id = f"result_map_{identifier}" + async with developer_api_client_factory(): + await _seed_revision( + namespace=namespace, + document_id=document_id, + job_id=job_id, + job_result_id=job_result_id, + ) + scope = DocumentPublicationScope( + user_id=_USER_ID, + namespace=namespace, + document_id=document_id, + job_result_id=job_result_id, + source_file_name="indexed.pdf", + ) + chunks = [ + { + "chunk_id": "parent-a", + "type": "text", + "content": "common alpha parent evidence", + "path": "indexed.pdf/Root/Parent/intro-a", + "order": 1, + "metadata": {}, + }, + { + "chunk_id": "parent-b", + "type": "text", + "content": "common beta parent evidence", + "path": "indexed.pdf/Root/Parent/intro-b", + "order": 2, + "metadata": {}, + }, + { + "chunk_id": "leaf-a", + "type": "text", + "content": "common alpha leaf evidence", + "path": "indexed.pdf/Root/Parent/Leaf A/body", + "order": 3, + "metadata": {}, + }, + { + "chunk_id": "leaf-b", + "type": "text", + "content": "common beta leaf evidence", + "path": "indexed.pdf/Root/Parent/Leaf B/body", + "order": 4, + "metadata": {}, + }, + ] + async with contract_db_session() as db: + await db.run_sync( + lambda sync_db: replace_document_revision_content( + sync_db, + scope=scope, + chunks=chunks, + ) + ) + await db.commit() + + async with contract_db_session() as db: + eager_snapshot = await load_nav_snapshot( + db, + user_id=_USER_ID, + namespace=namespace, + ) + eager_toolspace = ProviderToolSpace(eager_snapshot.provider) + expected_units = build_score_units(eager_toolspace, document_id) + expected_scores = compute_corpus_map_and_unit_scores( + eager_toolspace, + doc_ids=[document_id], + query="common alpha", + ) + + async with contract_db_session() as db: + index = ( + await db.execute( + select(DocumentMapUnitIndex).where( + DocumentMapUnitIndex.document_id == document_id + ) + ) + ).scalar_one() + persisted_units = list( + ( + await db.execute( + select(DocumentMapUnit) + .where(DocumentMapUnit.document_id == document_id) + .order_by(DocumentMapUnit.sort_order) + ) + ).scalars() + ) + persisted_tokens = list( + ( + await db.execute( + select(DocumentMapUnitToken).where( + DocumentMapUnitToken.map_unit_id.in_( + [unit.id for unit in persisted_units] + ) + ) + ) + ).scalars() + ) + lazy_snapshot = await load_nav_snapshot( + db, + user_id=_USER_ID, + namespace=namespace, + lazy=True, + ) + + assert index.unit_count == len(expected_units) + assert [unit.unit_id for unit in persisted_units] == [ + str(unit["chunk_id"]) for unit in expected_units + ] + assert persisted_tokens + + def reject_payload_read( + _store: ReadOnlyChunkStore, + _section_ids_by_document: Mapping[str, Sequence[str]], + ) -> dict[str, list[UnitRow]]: + raise AssertionError("persisted map scoring loaded full chunk payloads") + + original_payload_loader = ReadOnlyChunkStore.load_documents_units + monkeypatch.setattr( + ReadOnlyChunkStore, + "load_documents_units", + reject_payload_read, + ) + actual_scores = compute_corpus_map_and_unit_scores( + ProviderToolSpace(lazy_snapshot.provider), + doc_ids=[document_id], + query="common alpha", + ) + monkeypatch.setattr( + ReadOnlyChunkStore, + "load_documents_units", + original_payload_loader, + ) + async with contract_db_session() as db: + token_id = ( + select(DocumentMapUnitToken.id) + .join( + DocumentMapUnit, + DocumentMapUnit.id == DocumentMapUnitToken.map_unit_id, + ) + .where(DocumentMapUnit.document_id == document_id) + .limit(1) + .scalar_subquery() + ) + await db.execute( + delete(DocumentMapUnitToken).where(DocumentMapUnitToken.id == token_id) + ) + await db.commit() + incomplete_snapshot = await load_nav_snapshot( + db, + user_id=_USER_ID, + namespace=namespace, + lazy=True, + ) + fallback_scores = compute_corpus_map_and_unit_scores( + ProviderToolSpace(incomplete_snapshot.provider), + doc_ids=[document_id], + query="common alpha", + ) + incomplete_snapshot.close() + lazy_snapshot.close() + eager_snapshot.close() + + assert actual_scores == expected_scores + assert fallback_scores == expected_scores + + +def test_incomplete_index_falls_back_for_duplicate_unit_ids() -> None: + first_sections = [ + SectionRow("root-a", None, "Root A", "Root A", 0, "", 0), + SectionRow("leaf-a", "root-a", "Root A / Leaf A", "Leaf A", 1, "", 1), + ] + second_sections = [ + SectionRow("root-b", None, "Root B", "Root B", 0, "", 0), + SectionRow("leaf-b", "root-b", "Root B / Leaf B", "Leaf B", 1, "", 1), + ] + first_unit = UnitRow("same-chunk", "leaf-a", "text", "alpha evidence", 1) + second_unit = UnitRow("same-chunk", "leaf-b", "text", "beta evidence", 1) + + eager = ProviderToolSpace( + NamespaceKnowhereProvider( + [ + KnowhereProvider( + doc_id="doc-a", sections=first_sections, units=[first_unit] + ), + KnowhereProvider( + doc_id="doc-b", sections=second_sections, units=[second_unit] + ), + ], + titles={"doc-a": "Document A", "doc-b": "Document B"}, + ) + ) + store = _IncompleteIndexStore( + {"leaf-a": [first_unit], "leaf-b": [second_unit]} + ) + lazy = ProviderToolSpace( + NamespaceKnowhereProvider( + [ + LazyKnowhereProvider( + doc_id="doc-a", + sections=first_sections, + chunk_store=store, + known_chunk_ids=[first_unit.chunk_id], + ), + LazyKnowhereProvider( + doc_id="doc-b", + sections=second_sections, + chunk_store=store, + known_chunk_ids=[second_unit.chunk_id], + ), + ], + titles={"doc-a": "Document A", "doc-b": "Document B"}, + chunk_owner_by_id={"same-chunk": "doc-a"}, + ) + ) + + expected = compute_corpus_map_and_unit_scores( + eager, doc_ids=["doc-a", "doc-b"], query="alpha beta" + ) + actual = compute_corpus_map_and_unit_scores( + lazy, doc_ids=["doc-a", "doc-b"], query="alpha beta" + ) + + assert actual == expected + assert store.persisted_loads == 1 + assert store.batch_loads == 1 + + +def test_titleless_leaf_has_identical_eager_and_lazy_path_scoring() -> None: + sections = [ + SectionRow("root", None, "Root", "Root", 0, "", 0), + SectionRow("leaf", "root", "Root / Leaf", "", 1, "", 1), + ] + unit = UnitRow("titleless-chunk", "leaf", "text", "alpha evidence", 1) + eager = ProviderToolSpace( + KnowhereProvider(doc_id="doc", sections=sections, units=[unit]) + ) + store = _IncompleteIndexStore({"leaf": [unit]}) + lazy = ProviderToolSpace( + LazyKnowhereProvider( + doc_id="doc", + sections=sections, + chunk_store=store, + known_chunk_ids=[unit.chunk_id], + ) + ) + + assert build_score_units(eager, "doc") == build_score_units(lazy, "doc") + assert compute_corpus_map_and_unit_scores( + eager, doc_ids=["doc"], query="alpha" + ) == compute_corpus_map_and_unit_scores( + lazy, doc_ids=["doc"], query="alpha" + ) + + +async def _seed_revision( + *, + namespace: str, + document_id: str, + job_id: str, + job_result_id: str, +) -> None: + await ContractDatabase.execute( + """ + INSERT INTO jobs ( + job_id, user_id, job_type, status, source_type, version, + webhook_enabled, created_at, updated_at, credits_charged, billing_status + ) VALUES ( + :job_id, :user_id, 'document_ingestion', 'done', 'file', 0, + false, NOW(), NOW(), 0, 'skipped' + ) + """, + {"job_id": job_id, "user_id": _USER_ID}, + ) + await ContractDatabase.execute( + """ + INSERT INTO documents ( + document_id, user_id, namespace, status, source_file_name, + parse_track, created_at, updated_at + ) VALUES ( + :document_id, :user_id, :namespace, 'active', + 'indexed.pdf', 'chunk', NOW(), NOW() + ) + """, + { + "document_id": document_id, + "user_id": _USER_ID, + "namespace": namespace, + }, + ) + await ContractDatabase.execute( + """ + INSERT INTO job_results ( + id, job_id, document_id, delivery_mode, created_at, updated_at + ) VALUES ( + :job_result_id, :job_id, :document_id, 'inline', NOW(), NOW() + ) + """, + { + "job_result_id": job_result_id, + "job_id": job_id, + "document_id": document_id, + }, + ) + await ContractDatabase.execute( + """ + UPDATE documents SET current_job_result_id = :job_result_id + WHERE document_id = :document_id + """, + {"job_result_id": job_result_id, "document_id": document_id}, + ) diff --git a/packages/shared-python/shared/models/database/__init__.py b/packages/shared-python/shared/models/database/__init__.py index 2d0e7d8a8..5535cd9bf 100644 --- a/packages/shared-python/shared/models/database/__init__.py +++ b/packages/shared-python/shared/models/database/__init__.py @@ -13,6 +13,9 @@ from .document import ( Document, DocumentChunk, + DocumentMapUnit, + DocumentMapUnitIndex, + DocumentMapUnitToken, DocumentSection, GraphEdge, GraphNode, @@ -56,6 +59,9 @@ "Document", "DocumentSection", "DocumentChunk", + "DocumentMapUnit", + "DocumentMapUnitIndex", + "DocumentMapUnitToken", "DocumentPagePlan", "DemoMaterialization", "GraphNode", diff --git a/packages/shared-python/shared/models/database/document.py b/packages/shared-python/shared/models/database/document.py index b23819964..33dda2f91 100644 --- a/packages/shared-python/shared/models/database/document.py +++ b/packages/shared-python/shared/models/database/document.py @@ -49,7 +49,9 @@ class Document(Base): document_metadata: Mapped[Optional[Dict[str, Any]]] = mapped_column( JSON, nullable=True ) - parse_track: Mapped[str] = mapped_column(String(32), nullable=False, default="chunk") + parse_track: Mapped[str] = mapped_column( + String(32), nullable=False, default="chunk" + ) created_at: Mapped[datetime] = mapped_column( DateTime, default=utc_now_naive, nullable=False ) @@ -139,6 +141,7 @@ class DocumentChunk(Base): id: Mapped[str] = mapped_column( String(36), primary_key=True, default=lambda: f"dchk_{uuid4().hex[:12]}" ) + chunk_id: Mapped[str] = mapped_column(String(64), nullable=False) user_id: Mapped[str] = mapped_column(Text, nullable=False) namespace: Mapped[str] = mapped_column( @@ -233,6 +236,110 @@ class DocumentChunk(Base): ) +class DocumentMapUnit(Base): + """Persisted lexical map unit for one document revision. + + These rows are a derived index of the exact leaf and interstitial units + used by map-nav. Full chunk payloads remain in ``document_chunks`` and are + loaded separately for evidence hydration. + """ + + __tablename__ = "document_map_units" + + id: Mapped[str] = mapped_column(String(160), primary_key=True) + document_id: Mapped[str] = mapped_column( + String(36), + ForeignKey("documents.document_id", ondelete="CASCADE"), + nullable=False, + ) + job_result_id: Mapped[str] = mapped_column( + String(36), ForeignKey("job_results.id", ondelete="CASCADE"), nullable=False + ) + unit_id: Mapped[str] = mapped_column(String(128), nullable=False) + section_id: Mapped[str] = mapped_column(String(36), nullable=False) + unit_kind: Mapped[str] = mapped_column(String(32), nullable=False) + path_token_count: Mapped[int] = mapped_column(Integer, nullable=False) + content_token_count: Mapped[int] = mapped_column(Integer, nullable=False) + term_search_text_lower: Mapped[str] = mapped_column(Text, nullable=False) + sort_order: Mapped[int] = mapped_column(Integer, nullable=False, default=0) + created_at: Mapped[datetime] = mapped_column( + DateTime, default=utc_now_naive, nullable=False + ) + + __table_args__ = ( + Index( + "idx_document_map_units_revision_order", + "document_id", + "job_result_id", + "sort_order", + "unit_id", + ), + Index("idx_document_map_units_section", "section_id"), + ) + + +class DocumentMapUnitToken(Base): + """One exact token frequency in a persisted map unit channel.""" + + __tablename__ = "document_map_unit_tokens" + + id: Mapped[str] = mapped_column(String(36), primary_key=True) + map_unit_id: Mapped[str] = mapped_column( + String(160), + ForeignKey("document_map_units.id", ondelete="CASCADE"), + nullable=False, + ) + channel: Mapped[str] = mapped_column(String(16), nullable=False) + token: Mapped[str] = mapped_column(Text, nullable=False) + token_hash: Mapped[str] = mapped_column(String(64), nullable=False) + frequency: Mapped[int] = mapped_column(Integer, nullable=False) + + __table_args__ = ( + Index( + "idx_document_map_unit_tokens_lookup", + "channel", + "token_hash", + "map_unit_id", + ), + Index("idx_document_map_unit_tokens_unit", "map_unit_id", "channel"), + ) + + +class DocumentMapUnitIndex(Base): + """Completeness marker for a revision's materialized map-unit index.""" + + __tablename__ = "document_map_unit_indexes" + + id: Mapped[str] = mapped_column(String(100), primary_key=True) + document_id: Mapped[str] = mapped_column( + String(36), + ForeignKey("documents.document_id", ondelete="CASCADE"), + nullable=False, + ) + job_result_id: Mapped[str] = mapped_column( + String(36), ForeignKey("job_results.id", ondelete="CASCADE"), nullable=False + ) + format_version: Mapped[int] = mapped_column(Integer, nullable=False) + unit_count: Mapped[int] = mapped_column(Integer, nullable=False) + token_count: Mapped[int] = mapped_column(Integer, nullable=False) + created_at: Mapped[datetime] = mapped_column( + DateTime, default=utc_now_naive, nullable=False + ) + + __table_args__ = ( + UniqueConstraint( + "document_id", + "job_result_id", + name="uq_document_map_unit_indexes_revision", + ), + Index( + "idx_document_map_unit_indexes_revision", + "document_id", + "job_result_id", + ), + ) + + class GraphNode(Base): """Persisted derived graph node used for routing and expansion.""" @@ -385,44 +492,58 @@ class RetrievalHitStat(Base): class RetrievalRun(Base): """One row per agentic retrieval query. Append-only analytics.""" - __tablename__ = 'retrieval_runs' + __tablename__ = "retrieval_runs" - run_id: Mapped[str] = mapped_column(String(36), primary_key=True, default=lambda: f'aret_{uuid4().hex[:12]}') + run_id: Mapped[str] = mapped_column( + String(36), primary_key=True, default=lambda: f"aret_{uuid4().hex[:12]}" + ) user_id: Mapped[str] = mapped_column(Text, nullable=False) - namespace: Mapped[str] = mapped_column(String(255), nullable=False, default='default') + namespace: Mapped[str] = mapped_column( + String(255), nullable=False, default="default" + ) query: Mapped[str] = mapped_column(Text, nullable=False) - query_hash: Mapped[str] = mapped_column(String(32), nullable=False, default='') + query_hash: Mapped[str] = mapped_column(String(32), nullable=False, default="") top_k: Mapped[int] = mapped_column(Integer, nullable=False, default=10) chunk_types: Mapped[Optional[List[str]]] = mapped_column(JSON, nullable=True) filters: Mapped[Optional[Dict[str, Any]]] = mapped_column(JSON, nullable=True) - policy_name: Mapped[str] = mapped_column(String(64), nullable=False, default='rule_based_v1') + policy_name: Mapped[str] = mapped_column( + String(64), nullable=False, default="rule_based_v1" + ) agentic_enabled: Mapped[bool] = mapped_column(nullable=False, default=True) cache_hit: Mapped[bool] = mapped_column(nullable=False, default=False) result_count: Mapped[int] = mapped_column(Integer, nullable=False, default=0) final_doc_ids: Mapped[Optional[List[str]]] = mapped_column(JSON, nullable=True) - result_provenance: Mapped[Optional[Dict[str, Any]]] = mapped_column(JSON, nullable=True) - parent_run_id: Mapped[Optional[str]] = mapped_column(String(36), nullable=True, index=True) + result_provenance: Mapped[Optional[Dict[str, Any]]] = mapped_column( + JSON, nullable=True + ) + parent_run_id: Mapped[Optional[str]] = mapped_column( + String(36), nullable=True, index=True + ) workflow_step_id: Mapped[Optional[str]] = mapped_column(String(64), nullable=True) workflow_plan: Mapped[Optional[Dict[str, Any]]] = mapped_column(JSON, nullable=True) latency_ms: Mapped[int] = mapped_column(Integer, nullable=False, default=0) token_count: Mapped[Optional[int]] = mapped_column(Integer, nullable=True) error: Mapped[Optional[str]] = mapped_column(Text, nullable=True) - created_at: Mapped[datetime] = mapped_column(DateTime, default=datetime.utcnow, nullable=False) + created_at: Mapped[datetime] = mapped_column( + DateTime, default=datetime.utcnow, nullable=False + ) completed_at: Mapped[Optional[datetime]] = mapped_column(DateTime, nullable=True) __table_args__ = ( - Index('idx_retrieval_runs_user_namespace', 'user_id', 'namespace'), - Index('idx_retrieval_runs_created', 'created_at'), - Index('idx_retrieval_runs_query_hash', 'query_hash'), + Index("idx_retrieval_runs_user_namespace", "user_id", "namespace"), + Index("idx_retrieval_runs_created", "created_at"), + Index("idx_retrieval_runs_query_hash", "query_hash"), ) class RetrievalStep(Base): """One row per agent step within a retrieval run. Append-only analytics.""" - __tablename__ = 'retrieval_steps' + __tablename__ = "retrieval_steps" - step_id: Mapped[str] = mapped_column(String(36), primary_key=True, default=lambda: f'arst_{uuid4().hex[:12]}') + step_id: Mapped[str] = mapped_column( + String(36), primary_key=True, default=lambda: f"arst_{uuid4().hex[:12]}" + ) run_id: Mapped[str] = mapped_column(String(36), nullable=False, index=True) step_index: Mapped[int] = mapped_column(Integer, nullable=False, default=0) action_type: Mapped[str] = mapped_column(String(64), nullable=False) @@ -434,9 +555,11 @@ class RetrievalStep(Base): token_count: Mapped[Optional[int]] = mapped_column(Integer, nullable=True) model_name: Mapped[Optional[str]] = mapped_column(String(128), nullable=True) error: Mapped[Optional[str]] = mapped_column(Text, nullable=True) - created_at: Mapped[datetime] = mapped_column(DateTime, default=datetime.utcnow, nullable=False) + created_at: Mapped[datetime] = mapped_column( + DateTime, default=datetime.utcnow, nullable=False + ) __table_args__ = ( - Index('idx_retrieval_steps_run', 'run_id', 'step_index'), - Index('idx_retrieval_steps_created', 'created_at'), + Index("idx_retrieval_steps_run", "run_id", "step_index"), + Index("idx_retrieval_steps_created", "created_at"), ) diff --git a/packages/shared-python/shared/services/retrieval/map_unit_index.py b/packages/shared-python/shared/services/retrieval/map_unit_index.py new file mode 100644 index 000000000..15454dc9d --- /dev/null +++ b/packages/shared-python/shared/services/retrieval/map_unit_index.py @@ -0,0 +1,163 @@ +"""Publication-time materialization of exact map-nav lexical units.""" + +from __future__ import annotations + +from collections import Counter +from hashlib import sha256 +from uuid import uuid4 + +from sqlalchemy import delete, select +from sqlalchemy.orm import Session + +from shared.models.database.document import ( + DocumentChunk, + DocumentMapUnit, + DocumentMapUnitIndex, + DocumentMapUnitToken, + DocumentSection, +) +from shared.services.retrieval.nav.nav_hierarchy import ProviderToolSpace +from shared.services.retrieval.nav.nav_knowhere import ( + KnowhereProvider, + SectionRow, + UnitRow, +) +from shared.services.retrieval.nav.nav_map_scores import build_score_units +from shared.services.retrieval.publication_models import DocumentPublicationScope + + +MAP_UNIT_INDEX_FORMAT_VERSION = 1 + + +def replace_document_map_units( + db: Session, + *, + scope: DocumentPublicationScope, +) -> None: + """Build the derived index through the authoritative map-unit constructor.""" + db.execute( + delete(DocumentMapUnitToken).where( + DocumentMapUnitToken.map_unit_id.in_( + select(DocumentMapUnit.id) + .where(DocumentMapUnit.document_id == scope.document_id) + .where(DocumentMapUnit.job_result_id == scope.job_result_id) + ) + ) + ) + db.execute( + delete(DocumentMapUnit) + .where(DocumentMapUnit.document_id == scope.document_id) + .where(DocumentMapUnit.job_result_id == scope.job_result_id) + ) + db.execute( + delete(DocumentMapUnitIndex) + .where(DocumentMapUnitIndex.document_id == scope.document_id) + .where(DocumentMapUnitIndex.job_result_id == scope.job_result_id) + ) + section_models = list( + db.scalars( + select(DocumentSection) + .where(DocumentSection.document_id == scope.document_id) + .where(DocumentSection.job_result_id == scope.job_result_id) + .order_by(DocumentSection.sort_order, DocumentSection.section_id) + ) + ) + chunk_models = list( + db.scalars( + select(DocumentChunk) + .where(DocumentChunk.document_id == scope.document_id) + .where(DocumentChunk.job_result_id == scope.job_result_id) + .order_by( + DocumentChunk.sort_order, + DocumentChunk.chunk_id, + DocumentChunk.id, + ) + ) + ) + provider = KnowhereProvider( + doc_id=scope.document_id, + sections=[_to_section_row(section) for section in section_models], + units=[_to_unit_row(chunk) for chunk in chunk_models], + ) + score_units = build_score_units( + ProviderToolSpace(provider), + scope.document_id, + ) + persisted_count = 0 + token_count = 0 + for sort_order, unit in enumerate(score_units): + unit_id = str(unit.get("chunk_id") or "").strip() + section_id = str(unit.get("section_id") or "").strip() + if not unit_id or not section_id: + continue + map_unit_id = f"dmu_{uuid4().hex}" + path_tokens = str(unit.get("path_search_text") or "").split() + content_tokens = str(unit.get("content_search_text") or "").split() + db.add( + DocumentMapUnit( + id=map_unit_id, + document_id=scope.document_id, + job_result_id=scope.job_result_id, + unit_id=unit_id, + section_id=section_id, + unit_kind=str(unit.get("kind") or "leaf"), + path_token_count=len(path_tokens), + content_token_count=len(content_tokens), + term_search_text_lower=str(unit.get("term_search_text") or "").lower(), + sort_order=sort_order, + ) + ) + for channel, frequencies in ( + ("path", Counter(path_tokens)), + ("content", Counter(content_tokens)), + ): + for token, frequency in frequencies.items(): + db.add( + DocumentMapUnitToken( + id=f"dmut_{uuid4().hex[:31]}", + map_unit_id=map_unit_id, + channel=channel, + token=token, + token_hash=sha256(token.encode("utf-8")).hexdigest(), + frequency=frequency, + ) + ) + token_count += len(frequencies) + persisted_count += 1 + db.add( + DocumentMapUnitIndex( + id=f"dmui_{uuid4().hex}", + document_id=scope.document_id, + job_result_id=scope.job_result_id, + format_version=MAP_UNIT_INDEX_FORMAT_VERSION, + unit_count=persisted_count, + token_count=token_count, + ) + ) + + +def _to_section_row(section: DocumentSection) -> SectionRow: + return SectionRow( + section_id=section.section_id, + parent_section_id=section.parent_section_id, + section_path=section.section_path, + section_title=str(section.section_title or ""), + section_level=section.section_level, + summary=str(section.summary or ""), + sort_order=section.sort_order, + ) + + +def _to_unit_row(chunk: DocumentChunk) -> UnitRow: + raw_metadata = chunk.chunk_metadata + metadata = dict(raw_metadata) if isinstance(raw_metadata, dict) else {} + return UnitRow( + chunk_id=chunk.chunk_id, + section_id=chunk.section_id, + chunk_type=chunk.chunk_type, + content=str(chunk.content or ""), + sort_order=chunk.sort_order, + source_chunk_path=str(chunk.source_chunk_path or ""), + file_path=str(chunk.file_path or ""), + metadata=metadata, + ) diff --git a/packages/shared-python/shared/services/retrieval/nav/knowhere_hybrid.py b/packages/shared-python/shared/services/retrieval/nav/knowhere_hybrid.py index a647566a5..369da962c 100644 --- a/packages/shared-python/shared/services/retrieval/nav/knowhere_hybrid.py +++ b/packages/shared-python/shared/services/retrieval/nav/knowhere_hybrid.py @@ -5,6 +5,7 @@ Reference: https://github.com/Ontos-AI/knowhere """ + from __future__ import annotations import os @@ -56,7 +57,9 @@ def _space_join_tokens(text: str) -> str: return " ".join(tokenize_for_retrieval(text, dedupe=False)) -def build_content_search_text(content: str, *, section_summary: Optional[str] = None) -> str: +def build_content_search_text( + content: str, *, section_summary: Optional[str] = None +) -> str: parts = [str(content or "").strip()] if section_summary and str(section_summary).strip(): parts.append(str(section_summary).strip()) @@ -116,7 +119,9 @@ def rank_rows_by_bm25( try: from rank_bm25 import BM25Okapi except ImportError: - return _rank_rows_by_token_overlap(rows, query_tokens, search_field=search_field) + return _rank_rows_by_token_overlap( + rows, query_tokens, search_field=search_field + ) corpus: List[List[str]] = [] ranked_rows: List[dict[str, Any]] = [] @@ -140,7 +145,9 @@ def rank_rows_by_bm25( return ranked_rows -def rank_rows_by_term_channel(rows: List[dict[str, Any]], query: str) -> List[dict[str, Any]]: +def rank_rows_by_term_channel( + rows: List[dict[str, Any]], query: str +) -> List[dict[str, Any]]: query_lower = query.lower().strip() query_tokens = tokenize_query_for_ranker(query) if not query_lower or not query_tokens: @@ -217,12 +224,24 @@ def normalize_row_scores( def _channel_weights() -> Tuple[float, float, float]: - path_w = float(os.environ.get("NAV_DISCOVERY_CHANNEL_WEIGHT_PATH", str(CHANNEL_WEIGHT_PATH)).strip() or CHANNEL_WEIGHT_PATH) + path_w = float( + os.environ.get( + "NAV_DISCOVERY_CHANNEL_WEIGHT_PATH", str(CHANNEL_WEIGHT_PATH) + ).strip() + or CHANNEL_WEIGHT_PATH + ) content_w = float( - os.environ.get("NAV_DISCOVERY_CHANNEL_WEIGHT_CONTENT", str(CHANNEL_WEIGHT_CONTENT)).strip() + os.environ.get( + "NAV_DISCOVERY_CHANNEL_WEIGHT_CONTENT", str(CHANNEL_WEIGHT_CONTENT) + ).strip() or CHANNEL_WEIGHT_CONTENT ) - term_w = float(os.environ.get("NAV_DISCOVERY_CHANNEL_WEIGHT_TERM", str(CHANNEL_WEIGHT_TERM)).strip() or CHANNEL_WEIGHT_TERM) + term_w = float( + os.environ.get( + "NAV_DISCOVERY_CHANNEL_WEIGHT_TERM", str(CHANNEL_WEIGHT_TERM) + ).strip() + or CHANNEL_WEIGHT_TERM + ) return path_w, content_w, term_w @@ -242,14 +261,23 @@ def hybrid_search_rows( recall_k = internal_recall_k if recall_k is None: - mult = int(os.environ.get("NAV_DISCOVERY_RECALL_MULT", str(INTERNAL_RECALL_K_MULTIPLIER)).strip() or INTERNAL_RECALL_K_MULTIPLIER) + mult = int( + os.environ.get( + "NAV_DISCOVERY_RECALL_MULT", str(INTERNAL_RECALL_K_MULTIPLIER) + ).strip() + or INTERNAL_RECALL_K_MULTIPLIER + ) recall_k = max(top_k, top_k * max(1, mult)) rrf_k = int(os.environ.get("NAV_DISCOVERY_RRF_K", str(RRF_K)).strip() or RRF_K) path_w, content_w, term_w = _channel_weights() - path_rows = rank_rows_by_bm25(list(rows), query_tokens, search_field="path_search_text")[:recall_k] - content_rows = rank_rows_by_bm25(list(rows), query_tokens, search_field="content_search_text")[:recall_k] + path_rows = rank_rows_by_bm25( + list(rows), query_tokens, search_field="path_search_text" + )[:recall_k] + content_rows = rank_rows_by_bm25( + list(rows), query_tokens, search_field="content_search_text" + )[:recall_k] term_rows = rank_rows_by_term_channel(list(rows), query)[:recall_k] fused = merge_channels_rrf( @@ -262,27 +290,32 @@ def hybrid_search_rows( return fused - def map_channel_weights() -> Tuple[float, float, float]: """Channel weights for map scoring (prefer NAV_MAP_* env, fall back to legacy names).""" path_w = float( os.environ.get( "NAV_MAP_CHANNEL_WEIGHT_PATH", - os.environ.get("NAV_DISCOVERY_CHANNEL_WEIGHT_PATH", str(CHANNEL_WEIGHT_PATH)), + os.environ.get( + "NAV_DISCOVERY_CHANNEL_WEIGHT_PATH", str(CHANNEL_WEIGHT_PATH) + ), ).strip() or CHANNEL_WEIGHT_PATH ) content_w = float( os.environ.get( "NAV_MAP_CHANNEL_WEIGHT_CONTENT", - os.environ.get("NAV_DISCOVERY_CHANNEL_WEIGHT_CONTENT", str(CHANNEL_WEIGHT_CONTENT)), + os.environ.get( + "NAV_DISCOVERY_CHANNEL_WEIGHT_CONTENT", str(CHANNEL_WEIGHT_CONTENT) + ), ).strip() or CHANNEL_WEIGHT_CONTENT ) term_w = float( os.environ.get( "NAV_MAP_CHANNEL_WEIGHT_TERM", - os.environ.get("NAV_DISCOVERY_CHANNEL_WEIGHT_TERM", str(CHANNEL_WEIGHT_TERM)), + os.environ.get( + "NAV_DISCOVERY_CHANNEL_WEIGHT_TERM", str(CHANNEL_WEIGHT_TERM) + ), ).strip() or CHANNEL_WEIGHT_TERM ) @@ -445,18 +478,23 @@ def fuse_channel_bm25_dense( dense_vals = [float(dense_by_id.get(uid, 0.0) or 0.0) for uid in unit_ids] bm25_n = _normalize_score_list(bm25_vals) dense_n = _normalize_score_list(dense_vals) - dense_w = float(os.environ.get("NAV_MAP_CHANNEL_DENSE_WEIGHT", "0.5").strip() or "0.5") + dense_w = float( + os.environ.get("NAV_MAP_CHANNEL_DENSE_WEIGHT", "0.5").strip() or "0.5" + ) dense_w = min(1.0, max(0.0, dense_w)) bm25_w = 1.0 - dense_w return { - uid: bm25_w * bm25_n[i] + dense_w * dense_n[i] - for i, uid in enumerate(unit_ids) + uid: bm25_w * bm25_n[i] + dense_w * dense_n[i] for i, uid in enumerate(unit_ids) } def _rank_ids_by_score(score_by_id: Dict[str, float]) -> List[str]: ranked = sorted( - ((sid, float(score)) for sid, score in score_by_id.items() if float(score) > 0.0), + ( + (sid, float(score)) + for sid, score in score_by_id.items() + if float(score) > 0.0 + ), key=lambda item: (-item[1], item[0]), ) return [sid for sid, _ in ranked] @@ -470,9 +508,7 @@ def score_rows_hybrid_all( content_texts: Optional[Dict[str, str]] = None, doc_id: Optional[str] = None, namespace: Optional[str] = None, - dense_scores_by_channel: Optional[ - Dict[str, Optional[Dict[str, float]]] - ] = None, + dense_scores_by_channel: Optional[Dict[str, Optional[Dict[str, float]]]] = None, ) -> List[dict[str, Any]]: """Score every row with path/content/term; optional within-channel dense fuse. @@ -490,7 +526,9 @@ def score_rows_hybrid_all( if not unit_ids: return [dict(row, score=0.0) for row in rows] - row_by_id = {str(row.get("chunk_id") or ""): dict(row) for row in rows if row.get("chunk_id")} + row_by_id = { + str(row.get("chunk_id") or ""): dict(row) for row in rows if row.get("chunk_id") + } path_w, content_w, term_w = map_channel_weights() rrf_k = int( os.environ.get( @@ -501,7 +539,9 @@ def score_rows_hybrid_all( ) if query_tokens: - path_ranked = rank_rows_by_bm25(list(rows), query_tokens, search_field="path_search_text") + path_ranked = rank_rows_by_bm25( + list(rows), query_tokens, search_field="path_search_text" + ) content_ranked = rank_rows_by_bm25( list(rows), query_tokens, search_field="content_search_text" ) @@ -513,7 +553,8 @@ def score_rows_hybrid_all( str(r.get("chunk_id") or ""): float(r.get("score") or 0.0) for r in path_ranked } content_bm25 = { - str(r.get("chunk_id") or ""): float(r.get("score") or 0.0) for r in content_ranked + str(r.get("chunk_id") or ""): float(r.get("score") or 0.0) + for r in content_ranked } term_bm25 = { str(r.get("chunk_id") or ""): float(r.get("score") or 0.0) for r in term_ranked @@ -553,19 +594,12 @@ def score_rows_hybrid_all( namespace=namespace, ) path_dense_by_id = ( - { - uid: float(path_dense_scores[i]) - for i, uid in enumerate(unit_ids) - } - if path_dense_scores is not None - and len(path_dense_scores) == len(unit_ids) + {uid: float(path_dense_scores[i]) for i, uid in enumerate(unit_ids)} + if path_dense_scores is not None and len(path_dense_scores) == len(unit_ids) else None ) content_dense_by_id = ( - { - uid: float(content_dense_scores[i]) - for i, uid in enumerate(unit_ids) - } + {uid: float(content_dense_scores[i]) for i, uid in enumerate(unit_ids)} if content_dense_scores is not None and len(content_dense_scores) == len(unit_ids) else None @@ -575,7 +609,9 @@ def score_rows_hybrid_all( content_dense_by_id = dense_scores_by_channel.get("content") path_channel = fuse_channel_bm25_dense(path_bm25, path_dense_by_id, unit_ids) - content_channel = fuse_channel_bm25_dense(content_bm25, content_dense_by_id, unit_ids) + content_channel = fuse_channel_bm25_dense( + content_bm25, content_dense_by_id, unit_ids + ) term_channel = {uid: float(term_bm25.get(uid, 0.0) or 0.0) for uid in unit_ids} # Convert channel scores to ranked lists for existing RRF merger. @@ -597,7 +633,9 @@ def _rows_from_scores(score_by_id: Dict[str, float]) -> List[dict[str, Any]]: top_k=len(unit_ids), k=rrf_k, ) - fused_by_id = {str(r.get("chunk_id") or ""): float(r.get("score") or 0.0) for r in fused} + fused_by_id = { + str(r.get("chunk_id") or ""): float(r.get("score") or 0.0) for r in fused + } out_rows: List[dict[str, Any]] = [] for uid in unit_ids: row = dict(row_by_id[uid]) @@ -639,18 +677,18 @@ def score_unit_stream_hybrid_many( for query_tokens in query_tokens_by_query.values() for token in query_tokens } - query_lower_by_query = { - query: query.lower().strip() for query in unique_queries - } + query_lower_by_query = {query: query.lower().strip() for query in unique_queries} + units: List[_StreamingManyUnit] = [] path_stats = _StreamingBm25Stats.empty() content_stats = _StreamingBm25Stats.empty() - units: List[_StreamingManyUnit] = [] for row in unit_factory(): unit_id = str(row.get("chunk_id") or "").strip() if not unit_id: continue path_tokens = _get_search_tokens(row, search_field="path_search_text") content_tokens = _get_search_tokens(row, search_field="content_search_text") + # BM25 corpus statistics are computed over the complete input rows by + # the eager scorer, including rows whose public IDs collide. path_stats.observe(path_tokens) content_stats.observe(content_tokens) path_frequencies = Counter(path_tokens) @@ -665,9 +703,7 @@ def score_unit_stream_hybrid_many( if query_lower in haystack: term_score = 100.0 else: - hit_count = sum( - 1 for token in query_tokens if token in haystack - ) + hit_count = sum(1 for token in query_tokens if token in haystack) if hit_count > 0: term_score = float(hit_count) term_scores.append(term_score) @@ -686,7 +722,7 @@ def score_unit_stream_hybrid_many( for token in query_token_set if content_frequencies[token] }, - term_scores=tuple(term_scores), + term_scores=tuple(term_scores), ) ) path_stats.finalize() @@ -725,25 +761,20 @@ def _score_streaming_units( path_by_id[unit.unit_id] = path_score content_by_id[unit.unit_id] = content_score term_by_id[unit.unit_id] = float( - unit.term_scores[query_index] if query_index < len(unit.term_scores) else 0.0 + unit.term_scores[query_index] + if query_index < len(unit.term_scores) + else 0.0 ) path_rows = [ - (score, unit_id) - for unit_id, score in path_by_id.items() - if score > 0.0 + (score, unit_id) for unit_id, score in path_by_id.items() if score > 0.0 ] content_rows = [ - (score, unit_id) - for unit_id, score in content_by_id.items() - if score > 0.0 + (score, unit_id) for unit_id, score in content_by_id.items() if score > 0.0 ] term_rows = [ - (score, unit_id) - for unit_id, score in term_by_id.items() - if score > 0.0 + (score, unit_id) for unit_id, score in term_by_id.items() if score > 0.0 ] - path_rows.sort(key=lambda item: (-item[0], item[1])) content_rows.sort(key=lambda item: (-item[0], item[1])) term_rows.sort(key=lambda item: (-item[0], item[1])) @@ -775,6 +806,87 @@ class _StreamingManyUnit: term_scores: Tuple[float, ...] +@dataclass(frozen=True) +class PersistedBm25Stats: + """Corpus statistics needed to reproduce ``BM25Okapi`` exactly.""" + + document_count: int + total_length: int + document_frequency: Mapping[str, int] + average_idf: float + + +@dataclass(frozen=True) +class PersistedScoreUnit: + """Query-specific frequencies for one persisted map unit.""" + + unit_id: str + path_length: int + content_length: int + path_frequencies: Mapping[str, int] + content_frequencies: Mapping[str, int] + term_scores: Tuple[float, ...] + + +@dataclass(frozen=True) +class PersistedScoreCorpus: + """Compact query projection loaded from the map-unit index.""" + + units: Sequence[PersistedScoreUnit] + path_stats: PersistedBm25Stats + content_stats: PersistedBm25Stats + + +def score_persisted_corpus_many( + corpus: PersistedScoreCorpus, + queries: Sequence[str], +) -> Dict[str, Dict[str, float]]: + """Apply the existing BM25/RRF scorer to persisted query projections.""" + unique_queries = list(dict.fromkeys(str(query) for query in queries)) + if not unique_queries: + return {} + path_stats = _restore_bm25_stats(corpus.path_stats) + content_stats = _restore_bm25_stats(corpus.content_stats) + units = [ + _StreamingManyUnit( + unit_id=unit.unit_id, + path_length=unit.path_length, + content_length=unit.content_length, + path_frequencies=unit.path_frequencies, + content_frequencies=unit.content_frequencies, + term_scores=unit.term_scores, + ) + for unit in corpus.units + ] + return { + query: _score_streaming_units( + units, + path_stats=path_stats, + content_stats=content_stats, + query_tokens=tokenize_query_for_ranker(query), + query_index=query_index, + ) + for query_index, query in enumerate(unique_queries) + } + + +def _restore_bm25_stats(source: PersistedBm25Stats) -> "_StreamingBm25Stats": + stats = _StreamingBm25Stats.empty() + stats.document_count = source.document_count + stats.total_length = source.total_length + stats.average_length = ( + source.total_length / source.document_count if source.document_count else 0.0 + ) + idf_by_token: Dict[str, float] = {} + for token, frequency in source.document_frequency.items(): + idf = math.log(source.document_count - frequency + 0.5) - math.log( + frequency + 0.5 + ) + idf_by_token[token] = 0.25 * source.average_idf if idf < 0.0 else idf + stats.idf_by_token = idf_by_token + return stats + + class _StreamingBm25Stats: """Exact BM25Okapi corpus statistics collected without row retention.""" 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 e6c8cafa3..d9e42b075 100644 --- a/packages/shared-python/shared/services/retrieval/nav/nav_hierarchy.py +++ b/packages/shared-python/shared/services/retrieval/nav/nav_hierarchy.py @@ -22,7 +22,22 @@ from __future__ import annotations from dataclasses import dataclass, field -from typing import Any, Dict, List, Optional, Protocol, Sequence, Set, Tuple, runtime_checkable +from typing import ( + Any, + Dict, + List, + Optional, + Protocol, + Sequence, + Set, + Tuple, + TYPE_CHECKING, + cast, + runtime_checkable, +) + +if TYPE_CHECKING: + from .knowhere_hybrid import PersistedScoreCorpus @dataclass @@ -170,7 +185,9 @@ def _node_unit_span(self, section_id: str) -> Tuple[str, int, int]: first_order = int(getattr(units[0], "sort_order", 0) or 0) return "\n".join(texts), first_order, len(units) - def _make_chunk(self, node_id: str, doc_id: str, text: str, order: int, section_id: str) -> Any: + def _make_chunk( + self, node_id: str, doc_id: str, text: str, order: int, section_id: str + ) -> Any: from ._compat import Chunk # type: ignore return Chunk( @@ -208,7 +225,9 @@ def _materialize_leaf_path_chunks(self, section_id: str, doc_id: str) -> List[An text = str(self._provider.content(section_id) or "") if not text.strip(): return [] - return [self._make_chunk(f"{section_id}__path", doc_id, text, 0, section_id)] + return [ + self._make_chunk(f"{section_id}__path", doc_id, text, 0, section_id) + ] # One unit per descendant leaf, plus one per interstitial parent, so # node ids line up with the keys nav_map_scores.build_score_units emits. @@ -227,7 +246,9 @@ def _materialize_leaf_path_chunks(self, section_id: str, doc_id: str) -> List[An out.sort(key=lambda c: (min(c.line_ids or (0,)), c.node_id)) return out - def read_chunks(self, section_id: str, query: str, *, doc_id: str, k: int) -> List[Any]: + def read_chunks( + self, section_id: str, query: str, *, doc_id: str, k: int + ) -> List[Any]: del section_id, query, doc_id, k return [] @@ -270,6 +291,17 @@ def release_document_units(self, doc_id: str) -> None: if str(getattr(provider, "doc_id", "")) == str(doc_id): fn() + def load_persisted_score_corpus( + self, + doc_ids: Sequence[str], + queries: Sequence[str], + ) -> Optional["PersistedScoreCorpus"]: + """Forward the optional revision-pinned map-unit index capability.""" + fn = getattr(self._provider, "load_persisted_score_corpus", None) + if not callable(fn): + return None + return cast(Optional["PersistedScoreCorpus"], fn(doc_ids, queries)) + @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 b3b39f2bc..cf1de7790 100644 --- a/packages/shared-python/shared/services/retrieval/nav/nav_knowhere.py +++ b/packages/shared-python/shared/services/retrieval/nav/nav_knowhere.py @@ -21,16 +21,36 @@ from __future__ import annotations import os +from hashlib import sha256 from dataclasses import dataclass, field -from typing import Any, Callable, Dict, Iterable, List, Mapping, Optional, Protocol, Sequence, Set, Tuple +from typing import ( + Any, + Callable, + Dict, + Iterable, + List, + Mapping, + Optional, + Protocol, + Sequence, + Set, + Tuple, +) from .nav_address import NavLevel from .nav_hierarchy import NodeMeta +from .knowhere_hybrid import ( + PersistedBm25Stats, + PersistedScoreCorpus, + PersistedScoreUnit, + tokenize_query_for_ranker, +) _ASSET_TYPES = ("table", "image") # Knowhere sentinel path for the virtual document container (not a collectable leaf). ROOT_SECTION_PATH = "Root" _DEFAULT_DSN = "postgresql://root:root123@127.0.0.1:5433/Knowhere" +_MAP_UNIT_INDEX_FORMAT_VERSION = 1 @dataclass(frozen=True) @@ -141,6 +161,14 @@ def knowhere_database_url() -> str: class ChunkStore(Protocol): + def load_persisted_score_corpus( + self, + document_ids: Sequence[str], + allowed_section_ids_by_document: Mapping[str, Sequence[str]], + queries: Sequence[str], + ) -> Optional[PersistedScoreCorpus]: + raise NotImplementedError + def load_documents_units( self, section_ids_by_document: Mapping[str, Sequence[str]], @@ -197,11 +225,20 @@ def load_section_units( doc_id = str(document_id).strip() sid = str(section_id).strip() job_result_id = self._revisions.get(doc_id) - if not doc_id or not sid or not job_result_id or (doc_id, sid) in self._excluded_sections: + if ( + not doc_id + or not sid + or not job_result_id + or (doc_id, sid) in self._excluded_sections + ): return [] cur = self._connection().cursor() try: - ids = [str(chunk_id).strip() for chunk_id in extra_chunk_ids if str(chunk_id).strip()] + ids = [ + str(chunk_id).strip() + for chunk_id in extra_chunk_ids + if str(chunk_id).strip() + ] if ids: cur.execute( "SELECT chunk_id, section_id, chunk_type, content, sort_order, " @@ -225,6 +262,259 @@ def load_section_units( finally: cur.close() + def load_persisted_score_corpus( + self, + document_ids: Sequence[str], + allowed_section_ids_by_document: Mapping[str, Sequence[str]], + queries: Sequence[str], + ) -> Optional[PersistedScoreCorpus]: + """Load query-relevant score inputs when every revision is indexed.""" + revisions = [ + (document_id, self._revisions[document_id]) + for raw_document_id in document_ids + if (document_id := str(raw_document_id).strip()) in self._revisions + ] + if not revisions or len(revisions) != len(document_ids): + return None + values_sql = ", ".join(["(%s, %s)"] * len(revisions)) + revision_params: List[object] = [ + value for revision in revisions for value in revision + ] + cur = self._connection().cursor() + try: + cur.execute( + "SELECT indexes.document_id, indexes.job_result_id, " + "indexes.format_version, indexes.unit_count, indexes.token_count " + "FROM document_map_unit_indexes AS indexes " + f"JOIN (VALUES {values_sql}) AS revisions(document_id, job_result_id) " + "ON indexes.document_id = revisions.document_id " + "AND indexes.job_result_id = revisions.job_result_id", + revision_params, + ) + manifests = list(cur.fetchall()) + if len(manifests) != len(revisions) or any( + int(row[2]) != _MAP_UNIT_INDEX_FORMAT_VERSION for row in manifests + ): + return None + + cur.execute( + "SELECT COUNT(*), COUNT(DISTINCT (units.document_id, units.unit_id)) " + "FROM document_map_units AS units " + f"JOIN (VALUES {values_sql}) AS revisions(document_id, job_result_id) " + "ON units.document_id = revisions.document_id " + "AND units.job_result_id = revisions.job_result_id", + revision_params, + ) + unit_count_row = cur.fetchone() + indexed_unit_count = int(unit_count_row[0]) if unit_count_row else 0 + distinct_unit_count = int(unit_count_row[1]) if unit_count_row else 0 + expected_count = sum(int(row[3]) for row in manifests) + if indexed_unit_count != expected_count or distinct_unit_count != indexed_unit_count: + return None + cur.execute( + "SELECT COUNT(*) FROM document_map_unit_tokens AS tokens " + "JOIN document_map_units AS units ON units.id = tokens.map_unit_id " + f"JOIN (VALUES {values_sql}) AS revisions(document_id, job_result_id) " + "ON units.document_id = revisions.document_id " + "AND units.job_result_id = revisions.job_result_id", + revision_params, + ) + token_row = cur.fetchone() + indexed_token_count = int(token_row[0]) if token_row else 0 + expected_token_count = sum(int(row[4]) for row in manifests) + if indexed_token_count != expected_token_count: + return None + # The public chunk id is content-derived and may repeat within a + # revision. The persisted scorer keys scores by that id, so use + # the legacy payload path whenever ambiguity would change results. + allowed_by_document = { + str(document_id): {str(section_id) for section_id in section_ids} + for document_id, section_ids in allowed_section_ids_by_document.items() + } + allowed_pairs = [ + (document_id, section_id) + for document_id, section_ids in allowed_by_document.items() + for section_id in section_ids + ] + unit_rows: list[Sequence[object]] = [] + if allowed_pairs: + allowed_document_ids = [pair[0] for pair in allowed_pairs] + allowed_section_ids = [pair[1] for pair in allowed_pairs] + cur.execute( + "SELECT units.id, units.document_id, units.unit_id, units.section_id, " + "units.path_token_count, units.content_token_count " + "FROM document_map_units AS units " + f"JOIN (VALUES {values_sql}) AS revisions(document_id, job_result_id) " + "ON units.document_id = revisions.document_id " + "AND units.job_result_id = revisions.job_result_id " + "JOIN UNNEST(%s::text[], %s::text[]) " + "AS allowed(document_id, section_id) " + "ON units.document_id = allowed.document_id " + "AND units.section_id = allowed.section_id " + "ORDER BY units.document_id, units.sort_order, units.unit_id", + [*revision_params, allowed_document_ids, allowed_section_ids], + ) + unit_rows = list(cur.fetchall()) + map_unit_ids = [str(row[0]) for row in unit_rows] + unique_queries = list(dict.fromkeys(str(query) for query in queries)) + query_tokens_by_query = { + query: tokenize_query_for_ranker(query) for query in unique_queries + } + query_tokens = list( + dict.fromkeys( + token + for query in unique_queries + for token in query_tokens_by_query[query] + ) + ) + frequencies: Dict[Tuple[str, str], Dict[str, int]] = {} + if map_unit_ids and query_tokens: + query_token_hashes = [ + sha256(token.encode("utf-8")).hexdigest() for token in query_tokens + ] + cur.execute( + "SELECT map_unit_id, channel, token, frequency " + "FROM document_map_unit_tokens " + "WHERE map_unit_id = ANY(%s) AND token_hash = ANY(%s) " + "AND token = ANY(%s)", + (map_unit_ids, query_token_hashes, query_tokens), + ) + for map_unit_id, channel, token, frequency in cur.fetchall(): + frequencies.setdefault((str(map_unit_id), str(channel)), {})[ + str(token) + ] = int(frequency) + + term_scores = self._load_term_scores( + cur, + map_unit_ids=map_unit_ids, + queries=unique_queries, + query_tokens_by_query=query_tokens_by_query, + ) + path_stats = self._load_persisted_bm25_stats( + cur, + unit_rows=unit_rows, + map_unit_ids=map_unit_ids, + channel="path", + query_tokens=query_tokens, + frequencies=frequencies, + length_index=4, + ) + content_stats = self._load_persisted_bm25_stats( + cur, + unit_rows=unit_rows, + map_unit_ids=map_unit_ids, + channel="content", + query_tokens=query_tokens, + frequencies=frequencies, + length_index=5, + ) + return PersistedScoreCorpus( + units=[ + PersistedScoreUnit( + unit_id=str(row[2]), + path_length=int(row[4]), + content_length=int(row[5]), + path_frequencies=frequencies.get((str(row[0]), "path"), {}), + content_frequencies=frequencies.get( + (str(row[0]), "content"), {} + ), + term_scores=term_scores.get( + str(row[0]), tuple(0.0 for _query in unique_queries) + ), + ) + for row in unit_rows + ], + path_stats=path_stats, + content_stats=content_stats, + ) + finally: + cur.close() + + def _load_term_scores( + self, + cur: "_SyncCursor", + *, + map_unit_ids: Sequence[str], + queries: Sequence[str], + query_tokens_by_query: Mapping[str, Sequence[str]], + ) -> Dict[str, Tuple[float, ...]]: + if not map_unit_ids or not queries: + return {} + expressions: List[str] = [] + params: List[object] = [] + for query in queries: + query_lower = query.lower().strip() + if not query_lower: + expressions.append("0.0") + continue + token_expressions = [ + "CASE WHEN POSITION(%s IN term_search_text_lower) > 0 THEN 1 ELSE 0 END" + for _token in query_tokens_by_query[query] + ] + token_sum = " + ".join(token_expressions) or "0" + expressions.append( + "CASE WHEN POSITION(%s IN term_search_text_lower) > 0 " + f"THEN 100.0 ELSE ({token_sum})::double precision END" + ) + params.append(query_lower) + params.extend(query_tokens_by_query[query]) + params.append(list(map_unit_ids)) + cur.execute( + "SELECT id, " + ", ".join(expressions) + " " + "FROM document_map_units WHERE id = ANY(%s)", + params, + ) + return { + str(row[0]): tuple(float(value) for value in row[1:]) + for row in cur.fetchall() + } + + def _load_persisted_bm25_stats( + self, + cur: "_SyncCursor", + *, + unit_rows: Sequence[Sequence[object]], + map_unit_ids: Sequence[str], + channel: str, + query_tokens: Sequence[str], + frequencies: Mapping[Tuple[str, str], Mapping[str, int]], + length_index: int, + ) -> PersistedBm25Stats: + lengths = [ + int(row[length_index]) for row in unit_rows if int(row[length_index]) > 0 + ] + document_count = len(lengths) + document_frequency = { + token: sum( + 1 + for row in unit_rows + if frequencies.get((str(row[0]), channel), {}).get(token, 0) > 0 + ) + for token in query_tokens + } + needs_average_idf = any( + frequency > document_count / 2 for frequency in document_frequency.values() + ) + average_idf = 0.0 + if needs_average_idf and map_unit_ids and document_count: + cur.execute( + "SELECT COALESCE(AVG(LN((%s - frequencies.document_frequency + 0.5) " + "/ (frequencies.document_frequency + 0.5))), 0.0) " + "FROM (SELECT token, COUNT(*) AS document_frequency " + "FROM document_map_unit_tokens " + "WHERE map_unit_id = ANY(%s) AND channel = %s " + "GROUP BY token) AS frequencies", + (document_count, list(map_unit_ids), channel), + ) + row = cur.fetchone() + average_idf = float(row[0]) if row else 0.0 + return PersistedBm25Stats( + document_count=document_count, + total_length=sum(lengths), + document_frequency=document_frequency, + average_idf=average_idf, + ) + def load_document_units( self, document_id: str, @@ -234,7 +524,11 @@ def load_document_units( """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()] + 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() @@ -293,11 +587,7 @@ def load_documents_units( return {} values_sql = ", ".join(["(%s, %s)"] * len(revisions)) - params: List[object] = [ - value - for revision in revisions - for value in revision - ] + params: List[object] = [value for revision in revisions for value in revision] units_by_document: Dict[str, List[UnitRow]] = { document_id: [] for document_id, _job_result_id in revisions } @@ -367,6 +657,9 @@ def execute(self, query: str, params: Sequence[object]) -> None: def fetchall(self) -> Sequence[Sequence[object]]: raise NotImplementedError + def fetchone(self) -> Optional[Sequence[object]]: + raise NotImplementedError + def close(self) -> None: raise NotImplementedError @@ -889,7 +1182,9 @@ def load_namespace_from_db( providers = [load_document_from_db(did, dsn=url) for did in wanted] merged_titles = dict(auto_titles) if titles: - merged_titles.update({str(k): str(v) for k, v in titles.items() if str(k).strip()}) + merged_titles.update( + {str(k): str(v) for k, v in titles.items() if str(k).strip()} + ) return NamespaceKnowhereProvider(providers, titles=merged_titles or None) @@ -1001,6 +1296,42 @@ def prefetch_document_units_batch(self, doc_ids: Sequence[str]) -> None: loaded_by_document.get(provider.doc_id, ()) ) + def load_persisted_score_corpus( + self, + doc_ids: Sequence[str], + queries: Sequence[str], + ) -> Optional[PersistedScoreCorpus]: + """Return the index projection only when all documents share one store.""" + providers = [ + self._docs[document_id] + for raw_document_id in doc_ids + if (document_id := str(raw_document_id).strip()) in self._docs + ] + if len(providers) != len(doc_ids) or not all( + isinstance(provider, LazyKnowhereProvider) for provider in providers + ): + return None + lazy_providers = [ + provider + for provider in providers + if isinstance(provider, LazyKnowhereProvider) + ] + stores = { + id(provider._chunk_store): provider._chunk_store + for provider in lazy_providers + } + if len(stores) != 1: + return None + store = next(iter(stores.values())) + loader = getattr(store, "load_persisted_score_corpus", None) + if not callable(loader): + return None + return loader( + [provider.doc_id for provider in lazy_providers], + {provider.doc_id: list(provider._sections) for provider in lazy_providers}, + queries, + ) + def release_document_units(self, doc_id: str) -> None: provider = self._docs.get(str(doc_id).strip()) release = getattr(provider, "release_document_units", None) @@ -1048,8 +1379,12 @@ def node_meta(self, section_id: str) -> NodeMeta: if sid in self._docs: provider = self._docs[sid] count_fn = getattr(provider, "chunk_count", None) - n_chunks = int(count_fn()) if callable(count_fn) else sum( - len(provider.self_units(sec)) for sec in provider.all_section_ids() + n_chunks = ( + int(count_fn()) + if callable(count_fn) + else sum( + len(provider.self_units(sec)) for sec in provider.all_section_ids() + ) ) return NodeMeta( title=self._titles.get(sid, sid), @@ -1081,7 +1416,9 @@ def content(self, section_id: str) -> str: if sid in self._docs: provider = self._docs[sid] return "\n".join( - provider.content(root) for root in provider.roots(sid) if provider.content(root) + provider.content(root) + for root in provider.roots(sid) + if provider.content(root) ) owner = self._section_owner.get(sid) if not owner: 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 5909b83de..6318f4903 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 @@ -9,6 +9,7 @@ build_path_search_text, build_term_search_text, score_rows_hybrid_all, + score_persisted_corpus_many, score_unit_stream_hybrid_many, ) @@ -22,7 +23,9 @@ def _children_ids(ts: Any, section_id: str, doc_id: str) -> List[str]: if not callable(children_fn): st = ts.get_structure(section_id) rows = st.get("children") or [] - return [str(r.get("section_id") or "").strip() for r in rows if r.get("section_id")] + return [ + str(r.get("section_id") or "").strip() for r in rows if r.get("section_id") + ] rows = children_fn(section_id, doc_id, limit=100000) return [str(r.get("section_id") or "").strip() for r in rows if r.get("section_id")] @@ -152,8 +155,7 @@ def _pool_unit_scores_to_tree( ) -> Dict[str, float]: """MAX-pool globally comparable unit scores onto one document tree.""" map_scores = { - leaf_id: float(unit_scores.get(leaf_id, 0.0) or 0.0) - for leaf_id in leaves + leaf_id: float(unit_scores.get(leaf_id, 0.0) or 0.0) for leaf_id in leaves } def score_node(section_id: str) -> float: @@ -164,12 +166,9 @@ def score_node(section_id: str) -> float: score = float(unit_scores.get(section_id, 0.0) or 0.0) map_scores[section_id] = score return score - descendant_leaves = _collect_descendant_leaves( - section_id, children_map, leaves - ) + descendant_leaves = _collect_descendant_leaves(section_id, children_map, leaves) parts = [ - float(unit_scores.get(leaf_id, 0.0) or 0.0) - for leaf_id in descendant_leaves + float(unit_scores.get(leaf_id, 0.0) or 0.0) for leaf_id in descendant_leaves ] self_key = f"{section_id}__self" if self_key in unit_scores: @@ -183,7 +182,9 @@ def score_node(section_id: str) -> float: return map_scores -def build_score_units(ts: Any, doc_id: str, root_ids: Optional[Sequence[str]] = None) -> List[dict]: +def build_score_units( + ts: Any, doc_id: str, root_ids: Optional[Sequence[str]] = None +) -> List[dict]: """Build leaf (+ interstitial self_only) units for hybrid scoring.""" if root_ids is None: root_ids = list(ts.sections_for_doc(doc_id)) @@ -212,7 +213,9 @@ def build_score_units(ts: Any, doc_id: str, root_ids: Optional[Sequence[str]] = section_path=path_text, section_title=title or content ), "content_search_text": build_content_search_text(content), - "term_search_text": build_term_search_text(content, path_text=path_text), + "term_search_text": build_term_search_text( + content, path_text=path_text + ), } ) @@ -239,7 +242,9 @@ def build_score_units(ts: Any, doc_id: str, root_ids: Optional[Sequence[str]] = section_path=path_text, section_title=titles.get(sid) or "" ), "content_search_text": build_content_search_text(self_text), - "term_search_text": build_term_search_text(self_text, path_text=path_text), + "term_search_text": build_term_search_text( + self_text, path_text=path_text + ), } ) return units @@ -274,7 +279,9 @@ def iter_score_units( section_path=path_text, section_title=title or content ), "content_search_text": build_content_search_text(content), - "term_search_text": build_term_search_text(content, path_text=path_text), + "term_search_text": build_term_search_text( + content, path_text=path_text + ), } finally: if callable(release): @@ -298,7 +305,9 @@ def iter_score_units( section_path=path_text, section_title=titles.get(sid) or "" ), "content_search_text": build_content_search_text(self_text), - "term_search_text": build_term_search_text(self_text, path_text=path_text), + "term_search_text": build_term_search_text( + self_text, path_text=path_text + ), } finally: if callable(release): @@ -349,7 +358,9 @@ def compute_map_and_unit_scores( doc_id=doc_id, namespace=ns, ) - unit_score = {str(r.get("chunk_id") or ""): float(r.get("score") or 0.0) for r in scored} + unit_score = { + str(r.get("chunk_id") or ""): float(r.get("score") or 0.0) for r in scored + } map_scores = _pool_unit_scores_to_tree(children_map, leaves, unit_score) return map_scores, unit_score @@ -452,9 +463,16 @@ def unit_factory() -> Iterator[ScoreUnitRow]: if callable(release): release(document_id) - unit_scores_by_query = score_unit_stream_hybrid_many( - unit_factory, - unique_queries, + persisted_loader = getattr(ts, "load_persisted_score_corpus", None) + persisted_corpus = ( + persisted_loader(valid_doc_ids, unique_queries) + if callable(persisted_loader) + else None + ) + unit_scores_by_query = ( + score_persisted_corpus_many(persisted_corpus, unique_queries) + if persisted_corpus is not None + else score_unit_stream_hybrid_many(unit_factory, unique_queries) ) results: Dict[str, Tuple[Dict[str, float], Dict[str, float]]] = {} for query in unique_queries: diff --git a/packages/shared-python/shared/services/retrieval/publication_content.py b/packages/shared-python/shared/services/retrieval/publication_content.py index 9ebbdd74a..19a7f2074 100644 --- a/packages/shared-python/shared/services/retrieval/publication_content.py +++ b/packages/shared-python/shared/services/retrieval/publication_content.py @@ -6,7 +6,13 @@ from sqlalchemy import delete from sqlalchemy.orm import Session -from shared.models.database.document import DocumentChunk, DocumentSection +from shared.models.database.document import ( + DocumentChunk, + DocumentMapUnit, + DocumentMapUnitIndex, + DocumentSection, +) +from shared.services.retrieval.map_unit_index import replace_document_map_units from shared.services.retrieval.publication_models import DocumentPublicationScope from shared.services.retrieval.search.lexical_text import ( build_content_lexical_text, @@ -57,7 +63,9 @@ def replace_document_revision_content( """Replace retrieval sections and chunks for one published document revision.""" _delete_existing_revision_content(db, scope=scope) section_publisher = DocumentSectionPublisher( - db=db, scope=scope, section_summaries=section_summaries, + db=db, + scope=scope, + section_summaries=section_summaries, ) for index, chunk in enumerate(chunks): safe_chunk = cast(dict[str, Any], remove_nul_characters(chunk)) @@ -81,6 +89,8 @@ def replace_document_revision_content( fallback_sort_order=index, ) ) + db.flush() + replace_document_map_units(db, scope=scope) class DocumentSectionPublisher: @@ -146,6 +156,16 @@ def _delete_existing_revision_content( *, scope: DocumentPublicationScope, ) -> None: + db.execute( + delete(DocumentMapUnitIndex) + .where(DocumentMapUnitIndex.document_id == scope.document_id) + .where(DocumentMapUnitIndex.job_result_id == scope.job_result_id) + ) + db.execute( + delete(DocumentMapUnit) + .where(DocumentMapUnit.document_id == scope.document_id) + .where(DocumentMapUnit.job_result_id == scope.job_result_id) + ) db.execute( delete(DocumentChunk) .where(DocumentChunk.document_id == scope.document_id) From 68970407605359d0726f55f561927ae5d780f583 Mon Sep 17 00:00:00 2001 From: suguanYang Date: Fri, 28 Aug 2026 18:43:05 +0800 Subject: [PATCH 3/3] docs: document map-nav index backfill operation --- deploy/ecs/README.md | 28 ++++++++++++++++++++++++++++ 1 file changed, 28 insertions(+) diff --git a/deploy/ecs/README.md b/deploy/ecs/README.md index e19200afe..dd0742b4d 100644 --- a/deploy/ecs/README.md +++ b/deploy/ecs/README.md @@ -69,6 +69,34 @@ workflow validates these resources, registers immutable image-digest task definitions, runs the production migration first, and then updates the ECS services. It does not create or delete AWS resources. +## Required post-deploy backfill + +The map-nav lexical-index migration creates the derived index tables, but it does +not rebuild indexes for revisions that already exist. Until those revisions are +backfilled, retrieval remains quality-preserving but uses the legacy scoring +path. Every release containing the map-nav index change must include the +following DevOps action in its release notification. + +Run the commands as a one-off container using the newly deployed API image and +the production database secret. Do not run them inside the long-lived API task. + +```bash +# Read-only inventory +python /app/scripts/backfill_map_unit_indexes.py + +# Optional canary: apply one affected document first +python /app/scripts/backfill_map_unit_indexes.py \ + --document-id \ + --apply + +# Apply to all current document revisions +python /app/scripts/backfill_map_unit_indexes.py --apply +``` + +The script commits each document revision independently and is safe to rerun. +Verify the canary retrieval before starting the full apply. New or republished +documents build their index automatically during publication. + ## Manual staging availability `.github/workflows/manage-staging.yml` exposes three manually dispatched