From 259065fec195db91d6abf6b572e479d2f26aafe5 Mon Sep 17 00:00:00 2001 From: suguanYang Date: Sat, 29 Aug 2026 14:41:02 +0800 Subject: [PATCH 1/5] perf: eliminate lazy mapnav tree n-plus-one loads --- ...etrieval_lazy_snapshot_quality_contract.py | 90 ++++++++++++++++++- .../test_retrieval_lazy_tree_contract.py | 33 +++++++ .../services/retrieval/nav/nav_hierarchy.py | 18 ++-- .../services/retrieval/nav/nav_knowhere.py | 54 ++++++++++- .../services/retrieval/nav/nav_map_scores.py | 31 +++++++ .../shared/services/retrieval/nav_snapshot.py | 34 ++++++- 6 files changed, 250 insertions(+), 10 deletions(-) create mode 100644 apps/api/tests/contract/test_retrieval_lazy_tree_contract.py 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 ae6d78ed..0752a3b9 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 @@ -2,6 +2,8 @@ from dataclasses import dataclass from collections.abc import Mapping, Sequence +from collections import Counter +import math from typing import Any from shared.services.retrieval.nav.nav_hierarchy import ProviderToolSpace @@ -20,7 +22,11 @@ ) from shared.services.retrieval.nav.knowhere_hybrid import ( ScoreUnitRow, + PersistedBm25Stats, + PersistedScoreCorpus, + PersistedScoreUnit, score_rows_hybrid_all, + score_persisted_corpus_many, score_unit_stream_hybrid_all, score_unit_stream_hybrid_many, ) @@ -304,8 +310,7 @@ def test_streaming_scorer_scores_multiple_queries_with_one_corpus_read() -> None ] queries: list[str] = ["alpha evidence", "beta evidence"] expected = { - query: score_unit_stream_hybrid_all(lambda: rows, query) - for query in queries + query: score_unit_stream_hybrid_all(lambda: rows, query) for query in queries } read_count: int = 0 @@ -318,6 +323,87 @@ def unit_factory() -> Sequence[ScoreUnitRow]: assert read_count == 1 +def test_persisted_score_projection_preserves_exact_eager_scores() -> None: + rows: list[ScoreUnitRow] = [ + { + "chunk_id": "unit-a", + "path_search_text": "root alpha", + "content_search_text": "common alpha alpha evidence", + "term_search_text": "common alpha alpha evidence root", + }, + { + "chunk_id": "unit-b", + "path_search_text": "root beta", + "content_search_text": "common beta evidence", + "term_search_text": "common beta evidence root", + }, + { + "chunk_id": "unit-c", + "path_search_text": "root common", + "content_search_text": "common evidence", + "term_search_text": "common evidence root", + }, + ] + queries = ["common alpha", "beta evidence"] + expected = { + query: score_unit_stream_hybrid_all(lambda: rows, query) for query in queries + } + query_tokens = {token for query in queries for token in query.split()} + + def build_stats(search_field: str) -> PersistedBm25Stats: + token_rows = [str(row[search_field]).split() for row in rows] + document_frequency = Counter( + token for tokens in token_rows for token in set(tokens) + ) + document_count = len(token_rows) + raw_idfs = [ + math.log(document_count - frequency + 0.5) - math.log(frequency + 0.5) + for frequency in document_frequency.values() + ] + return PersistedBm25Stats( + document_count=document_count, + total_length=sum(len(tokens) for tokens in token_rows), + document_frequency={ + token: document_frequency[token] for token in query_tokens + }, + average_idf=sum(raw_idfs) / len(raw_idfs), + ) + + corpus = PersistedScoreCorpus( + units=[ + PersistedScoreUnit( + unit_id=str(row["chunk_id"]), + path_length=len(str(row["path_search_text"]).split()), + content_length=len(str(row["content_search_text"]).split()), + path_frequencies={ + token: str(row["path_search_text"]).split().count(token) + for token in query_tokens + }, + content_frequencies={ + token: str(row["content_search_text"]).split().count(token) + for token in query_tokens + }, + term_scores=tuple( + 100.0 + if query in str(row["term_search_text"]) + else float( + sum( + token in str(row["term_search_text"]) + for token in query.split() + ) + ) + for query in queries + ), + ) + for row in rows + ], + path_stats=build_stats("path_search_text"), + content_stats=build_stats("content_search_text"), + ) + + assert score_persisted_corpus_many(corpus, queries) == expected + + def test_corpus_map_scores_multiple_queries_with_one_lazy_load() -> None: eager, lazy, store = _providers() queries: list[str] = ["alpha retrieval", "supporting image"] diff --git a/apps/api/tests/contract/test_retrieval_lazy_tree_contract.py b/apps/api/tests/contract/test_retrieval_lazy_tree_contract.py new file mode 100644 index 00000000..30994dad --- /dev/null +++ b/apps/api/tests/contract/test_retrieval_lazy_tree_contract.py @@ -0,0 +1,33 @@ +"""Contract coverage for lazy MAP-NAV tree traversal.""" + +from __future__ import annotations + +from shared.services.retrieval.nav.nav_hierarchy import NodeMeta, ProviderToolSpace +from shared.services.retrieval.nav.nav_knowhere import KnowhereProvider, SectionRow +from shared.services.retrieval.nav.nav_map_scores import _walk_tree + + +class _MetadataForbiddenProvider(KnowhereProvider): + def node_meta(self, section_id: str) -> NodeMeta: + raise AssertionError(f"tree traversal materialized metadata for {section_id}") + + +def test_tree_walk_reads_children_and_titles_without_materializing_metadata() -> None: + provider = _MetadataForbiddenProvider( + doc_id="doc", + sections=[ + SectionRow("root", None, "Root", "Root", 0, "", 0), + SectionRow("child", "root", "Root / Child", "Child", 1, "", 1), + ], + units=(), + ) + + children, leaves, titles = _walk_tree( + ProviderToolSpace(provider), + "doc", + ["root"], + ) + + assert children == {"root": ["child"], "child": []} + assert leaves == {"child"} + assert titles == {"root": "Root", "child": "Child"} 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 d9e42b07..e98a9d4c 100644 --- a/packages/shared-python/shared/services/retrieval/nav/nav_hierarchy.py +++ b/packages/shared-python/shared/services/retrieval/nav/nav_hierarchy.py @@ -145,14 +145,22 @@ def get_structure(self, section_id: str) -> dict: def _children_for_section_path( self, section_id: str, doc_id: str, limit: Optional[int] = None ) -> List[dict]: - del doc_id # section_id is globally addressable in this provider model. child_ids = [str(c) for c in self._provider.children(section_id)] if limit is not None: child_ids = child_ids[: max(0, int(limit))] - return [ - {"section_id": cid, "preview": self._provider.node_meta(cid).title} - for cid in child_ids - ] + # ``node_meta`` may materialize a lazy subtree to calculate chunk + # counts. Tree traversal needs only the child id/title; avoid an N+1 + # payload load while building the scoring tree. + out: List[dict] = [] + for cid in child_ids: + path = self.path_titles(cid, doc_id) + out.append( + { + "section_id": cid, + "preview": path.rsplit(" / ", 1)[-1] if path else "", + } + ) + return out def section_relation_ids( self, section_id: str, doc_id: str 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 cf1de779..c056a2d2 100644 --- a/packages/shared-python/shared/services/retrieval/nav/nav_knowhere.py +++ b/packages/shared-python/shared/services/retrieval/nav/nav_knowhere.py @@ -21,6 +21,8 @@ from __future__ import annotations import os +import logging +import time from hashlib import sha256 from dataclasses import dataclass, field from typing import ( @@ -51,6 +53,7 @@ ROOT_SECTION_PATH = "Root" _DEFAULT_DSN = "postgresql://root:root123@127.0.0.1:5433/Knowhere" _MAP_UNIT_INDEX_FORMAT_VERSION = 1 +_logger = logging.getLogger(__name__) @dataclass(frozen=True) @@ -282,6 +285,7 @@ def load_persisted_score_corpus( ] cur = self._connection().cursor() try: + stage_started = time.perf_counter() cur.execute( "SELECT indexes.document_id, indexes.job_result_id, " "indexes.format_version, indexes.unit_count, indexes.token_count " @@ -292,11 +296,17 @@ def load_persisted_score_corpus( revision_params, ) manifests = list(cur.fetchall()) + _logger.info( + "retrieval map-index load stage=manifests seconds=%.3f rows=%d", + time.perf_counter() - stage_started, + len(manifests), + ) if len(manifests) != len(revisions) or any( int(row[2]) != _MAP_UNIT_INDEX_FORMAT_VERSION for row in manifests ): return None + stage_started = time.perf_counter() cur.execute( "SELECT COUNT(*), COUNT(DISTINCT (units.document_id, units.unit_id)) " "FROM document_map_units AS units " @@ -308,9 +318,15 @@ def load_persisted_score_corpus( 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 + _logger.info( + "retrieval map-index load stage=unit_count seconds=%.3f rows=%d", + time.perf_counter() - stage_started, + indexed_unit_count, + ) 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 + stage_started = time.perf_counter() cur.execute( "SELECT COUNT(*) FROM document_map_unit_tokens AS tokens " "JOIN document_map_units AS units ON units.id = tokens.map_unit_id " @@ -321,6 +337,11 @@ def load_persisted_score_corpus( ) token_row = cur.fetchone() indexed_token_count = int(token_row[0]) if token_row else 0 + _logger.info( + "retrieval map-index load stage=token_count seconds=%.3f rows=%d", + time.perf_counter() - stage_started, + indexed_token_count, + ) expected_token_count = sum(int(row[4]) for row in manifests) if indexed_token_count != expected_token_count: return None @@ -340,6 +361,7 @@ def load_persisted_score_corpus( if allowed_pairs: allowed_document_ids = [pair[0] for pair in allowed_pairs] allowed_section_ids = [pair[1] for pair in allowed_pairs] + stage_started = time.perf_counter() cur.execute( "SELECT units.id, units.document_id, units.unit_id, units.section_id, " "units.path_token_count, units.content_token_count " @@ -355,6 +377,11 @@ def load_persisted_score_corpus( [*revision_params, allowed_document_ids, allowed_section_ids], ) unit_rows = list(cur.fetchall()) + _logger.info( + "retrieval map-index load stage=units seconds=%.3f rows=%d", + time.perf_counter() - stage_started, + len(unit_rows), + ) 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 = { @@ -372,6 +399,7 @@ def load_persisted_score_corpus( query_token_hashes = [ sha256(token.encode("utf-8")).hexdigest() for token in query_tokens ] + stage_started = time.perf_counter() cur.execute( "SELECT map_unit_id, channel, token, frequency " "FROM document_map_unit_tokens " @@ -383,6 +411,11 @@ def load_persisted_score_corpus( frequencies.setdefault((str(map_unit_id), str(channel)), {})[ str(token) ] = int(frequency) + _logger.info( + "retrieval map-index load stage=frequencies seconds=%.3f units=%d", + time.perf_counter() - stage_started, + len(map_unit_ids), + ) term_scores = self._load_term_scores( cur, @@ -390,6 +423,11 @@ def load_persisted_score_corpus( queries=unique_queries, query_tokens_by_query=query_tokens_by_query, ) + _logger.info( + "retrieval map-index load stage=complete units=%d queries=%d", + len(unit_rows), + len(unique_queries), + ) path_stats = self._load_persisted_bm25_stats( cur, unit_rows=unit_rows, @@ -459,14 +497,22 @@ def _load_term_scores( params.append(query_lower) params.extend(query_tokens_by_query[query]) params.append(list(map_unit_ids)) + stage_started = time.perf_counter() cur.execute( "SELECT id, " + ", ".join(expressions) + " " "FROM document_map_units WHERE id = ANY(%s)", params, ) + rows = cur.fetchall() + _logger.info( + "retrieval map-index load stage=term_scores units=%d queries=%d seconds=%.3f", + len(map_unit_ids), + len(queries), + time.perf_counter() - stage_started, + ) return { str(row[0]): tuple(float(value) for value in row[1:]) - for row in cur.fetchall() + for row in rows } def _load_persisted_bm25_stats( @@ -497,6 +543,7 @@ def _load_persisted_bm25_stats( ) average_idf = 0.0 if needs_average_idf and map_unit_ids and document_count: + stage_started = time.perf_counter() cur.execute( "SELECT COALESCE(AVG(LN((%s - frequencies.document_frequency + 0.5) " "/ (frequencies.document_frequency + 0.5))), 0.0) " @@ -508,6 +555,11 @@ def _load_persisted_bm25_stats( ) row = cur.fetchone() average_idf = float(row[0]) if row else 0.0 + _logger.info( + "retrieval map-index load stage=average_idf channel=%s seconds=%.3f", + channel, + time.perf_counter() - stage_started, + ) return PersistedBm25Stats( document_count=document_count, total_length=sum(lengths), 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 6318f490..18d793e9 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 @@ -1,6 +1,8 @@ from __future__ import annotations from collections.abc import Iterator +import logging +import time from typing import Any, Dict, List, Optional, Sequence, Set, Tuple from .knowhere_hybrid import ( @@ -16,6 +18,7 @@ # 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 +_logger = logging.getLogger(__name__) def _children_ids(ts: Any, section_id: str, doc_id: str) -> List[str]: @@ -412,10 +415,17 @@ def compute_corpus_map_and_unit_scores_many( str, Tuple[Dict[str, List[str]], Set[str], Dict[str, str]], ] = {} + tree_started = time.perf_counter() for doc_id in valid_doc_ids: root_ids = list(ts.sections_for_doc(doc_id)) children_map, leaves, titles = _walk_tree(ts, doc_id, root_ids) tree_by_doc[doc_id] = (children_map, leaves, titles) + _logger.info( + "retrieval mapnav phase=tree_build seconds=%.3f documents=%d sections=%d", + time.perf_counter() - tree_started, + len(valid_doc_ids), + sum(len(value[0]) for value in tree_by_doc.values()), + ) def unit_factory() -> Iterator[ScoreUnitRow]: prefetch_batch = getattr(ts, "prefetch_document_units_batch", None) @@ -464,17 +474,32 @@ def unit_factory() -> Iterator[ScoreUnitRow]: release(document_id) persisted_loader = getattr(ts, "load_persisted_score_corpus", None) + loader_started = time.perf_counter() persisted_corpus = ( persisted_loader(valid_doc_ids, unique_queries) if callable(persisted_loader) else None ) + _logger.info( + "retrieval mapnav phase=index_load seconds=%.3f persisted=%s", + time.perf_counter() - loader_started, + persisted_corpus is not None, + ) + score_started = time.perf_counter() 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) ) + _logger.info( + "retrieval mapnav phase=unit_scoring persisted=%s seconds=%.3f units=%d queries=%d", + persisted_corpus is not None, + time.perf_counter() - score_started, + sum(len(scores) for scores in unit_scores_by_query.values()), + len(unique_queries), + ) results: Dict[str, Tuple[Dict[str, float], Dict[str, float]]] = {} + pooling_started = time.perf_counter() for query in unique_queries: unit_scores = unit_scores_by_query.get(query, {}) map_scores: Dict[str, float] = {} @@ -490,6 +515,12 @@ def unit_factory() -> Iterator[ScoreUnitRow]: ) map_scores[doc_id] = doc_max results[query] = (map_scores, unit_scores) + _logger.info( + "retrieval mapnav phase=map_pooling seconds=%.3f documents=%d sections=%d", + time.perf_counter() - pooling_started, + len(valid_doc_ids), + sum(len(value[0]) for value in tree_by_doc.values()), + ) return results diff --git a/packages/shared-python/shared/services/retrieval/nav_snapshot.py b/packages/shared-python/shared/services/retrieval/nav_snapshot.py index 1edfd374..17668767 100644 --- a/packages/shared-python/shared/services/retrieval/nav_snapshot.py +++ b/packages/shared-python/shared/services/retrieval/nav_snapshot.py @@ -8,6 +8,8 @@ from __future__ import annotations import json +import logging +import time from dataclasses import dataclass from typing import Any, Protocol @@ -33,7 +35,10 @@ # statement bounded. The contract benchmark verifies this page size against # the full 2 KiB content and metadata payload. _CHUNK_BATCH_SIZE = 10_000 -_REVISION_GROUP_SIZE = 32 +# Keep revision predicates bounded while reducing round trips for large +# namespaces. Keyset paging still caps each payload query at 10,000 rows. +_REVISION_GROUP_SIZE = 64 +_logger = logging.getLogger(__name__) class SnapshotSession(Protocol): @@ -252,6 +257,9 @@ async def _load_chunk_index( ref_index: dict[str, dict[str, Any]] = {} root_assets_by_doc: dict[str, set[str]] = {} text_connections_by_doc: dict[str, list[tuple[str, str]]] = {} + query_seconds = 0.0 + assembly_seconds = 0.0 + query_count = 0 for group_start in range(0, len(document_revisions), _REVISION_GROUP_SIZE): revision_group = document_revisions[group_start : group_start + _REVISION_GROUP_SIZE] last_key: tuple[str, str, int, str, str] | None = None @@ -289,9 +297,13 @@ async def _load_chunk_index( ) > tuple_(*[literal(value) for value in last_key]) ) + query_started = time.perf_counter() rows = (await db.execute(stmt)).all() + query_seconds += time.perf_counter() - query_started + query_count += 1 if not rows: break + assembly_started = time.perf_counter() for row in rows: document_id = str(row[0]) job_result_id = str(row[1]) @@ -339,6 +351,7 @@ async def _load_chunk_index( text_connections_by_doc.setdefault(document_id, []).append( (section_id or "", target) ) + assembly_seconds += time.perf_counter() - assembly_started last = rows[-1] last_key = ( str(last[0]), @@ -356,6 +369,13 @@ async def _load_chunk_index( if target in asset_ids: owners.setdefault(section_id, []).append(target) remounted[document_id] = {"root": sorted(asset_ids), "owners": owners} + _logger.info( + "retrieval snapshot phase=chunk_index rows=%d queries=%d query_seconds=%.3f assembly_seconds=%.3f", + sum(len(chunk_ids) for chunk_ids in ids_by_doc.values()), + query_count, + query_seconds, + assembly_seconds, + ) return ids_by_doc, ref_index, remounted @@ -392,7 +412,11 @@ async def _load_sections( by_doc: dict[str, list[SectionRow]] = {} path_by_id: dict[str, str] = {} - for row in (await db.execute(stmt)).all(): + query_started = time.perf_counter() + rows = (await db.execute(stmt)).all() + query_seconds = time.perf_counter() - query_started + assembly_started = time.perf_counter() + for row in rows: document_id = str(row[0]) section_path = str(row[3] or "") if is_excluded_section( @@ -413,6 +437,12 @@ async def _load_sections( ) by_doc.setdefault(document_id, []).append(section) path_by_id[section_id] = section_path + _logger.info( + "retrieval snapshot phase=sections rows=%d query_seconds=%.3f assembly_seconds=%.3f", + len(rows), + query_seconds, + time.perf_counter() - assembly_started, + ) return by_doc, path_by_id From 0d859b93ca1dc10d68e7598c405084be25b2b5b4 Mon Sep 17 00:00:00 2001 From: suguanYang Date: Sat, 29 Aug 2026 16:08:54 +0800 Subject: [PATCH 2/5] perf: defer lazy reference metadata and tune token lookup --- ...b6c7d_add_map_unit_token_covering_index.py | 40 +++++ .../test_retrieval_map_unit_index_contract.py | 147 +++++++++++++++++- .../shared/models/database/document.py | 7 + .../services/retrieval/nav/nav_knowhere.py | 41 ++++- .../shared/services/retrieval/nav_snapshot.py | 71 +++++++-- 5 files changed, 290 insertions(+), 16 deletions(-) create mode 100644 apps/api/alembic/versions/2e3f4a5b6c7d_add_map_unit_token_covering_index.py diff --git a/apps/api/alembic/versions/2e3f4a5b6c7d_add_map_unit_token_covering_index.py b/apps/api/alembic/versions/2e3f4a5b6c7d_add_map_unit_token_covering_index.py new file mode 100644 index 00000000..61250639 --- /dev/null +++ b/apps/api/alembic/versions/2e3f4a5b6c7d_add_map_unit_token_covering_index.py @@ -0,0 +1,40 @@ +"""Add a covering index for persisted map-nav token lookups.""" + +from __future__ import annotations + +from alembic import op + + +revision = "2e3f4a5b6c7d" +down_revision = "1d2e3f4a5b6c" +branch_labels = None +depends_on = None + +_INDEX_NAME = "idx_document_map_unit_tokens_unit_lookup" + + +def upgrade() -> None: + external_transaction = bool( + op.get_context().opts.get("knowhere_external_transaction", False) + ) + statement = ( + f"CREATE INDEX {{concurrently}}IF NOT EXISTS {_INDEX_NAME} " + "ON document_map_unit_tokens (map_unit_id, channel, token_hash) " + "INCLUDE (token, frequency)" + ) + if external_transaction: + op.execute(statement.format(concurrently="")) + return + with op.get_context().autocommit_block(): + op.execute(statement.format(concurrently="CONCURRENTLY ")) + + +def downgrade() -> None: + external_transaction = bool( + op.get_context().opts.get("knowhere_external_transaction", False) + ) + if external_transaction: + op.execute(f"DROP INDEX IF EXISTS {_INDEX_NAME}") + return + with op.get_context().autocommit_block(): + op.execute(f"DROP INDEX CONCURRENTLY IF EXISTS {_INDEX_NAME}") 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 index 2f2027b0..56a2dced 100644 --- a/apps/api/tests/contract/test_retrieval_map_unit_index_contract.py +++ b/apps/api/tests/contract/test_retrieval_map_unit_index_contract.py @@ -5,7 +5,7 @@ from uuid import uuid4 from httpx import AsyncClient -from sqlalchemy import delete, select +from sqlalchemy import delete, select, text from shared.models.database.document import ( DocumentMapUnit, @@ -13,9 +13,11 @@ DocumentMapUnitToken, ) from shared.services.retrieval.nav.nav_hierarchy import ProviderToolSpace +from shared.services.retrieval.nav._compat import Chunk, EpisodeResult from shared.services.retrieval.nav.nav_map_scores import ( build_score_units, compute_corpus_map_and_unit_scores, + select_map_highlights, ) from shared.services.retrieval.nav.nav_knowhere import ( KnowhereProvider, @@ -29,6 +31,7 @@ from shared.services.retrieval.publication_content import ( replace_document_revision_content, ) +from shared.services.retrieval.nav_bridge import build_referenced_chunks 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 @@ -179,6 +182,7 @@ async def test_published_map_units_preserve_scores_without_chunk_payload_reads( doc_ids=[document_id], query="common alpha", ) + expected_highlights = select_map_highlights(expected_scores[1], k=3) async with contract_db_session() as db: index = ( @@ -208,6 +212,18 @@ async def test_published_map_units_preserve_scores_without_chunk_payload_reads( ) ).scalars() ) + index_names = { + str(row[0]) + for row in ( + await db.execute( + text( + "SELECT indexname FROM pg_indexes " + "WHERE schemaname = current_schema() " + "AND tablename = 'document_map_unit_tokens'" + ) + ) + ).all() + } lazy_snapshot = await load_nav_snapshot( db, user_id=_USER_ID, @@ -220,6 +236,7 @@ async def test_published_map_units_preserve_scores_without_chunk_payload_reads( str(unit["chunk_id"]) for unit in expected_units ] assert persisted_tokens + assert "idx_document_map_unit_tokens_lookup" in index_names def reject_payload_read( _store: ReadOnlyChunkStore, @@ -274,9 +291,137 @@ def reject_payload_read( eager_snapshot.close() assert actual_scores == expected_scores + assert select_map_highlights(actual_scores[1], k=3) == expected_highlights assert fallback_scores == expected_scores +async def test_lazy_snapshot_defers_selected_asset_reference_metadata( + developer_api_client_factory: Callable[ + [], AbstractAsyncContextManager[AsyncClient] + ], + monkeypatch, +) -> None: + identifier = uuid4().hex[:8] + namespace = f"lazy-ref-{identifier}" + document_id = f"doc_ref_{identifier}" + job_id = f"job_ref_{identifier}" + job_result_id = f"result_ref_{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="refs.pdf", + ) + chunks = [ + { + "chunk_id": "body", + "type": "text", + "content": "body evidence", + "path": "refs.pdf/Root/Section/body", + "order": 1, + "metadata": {"connect_to": [{"target": "asset"}]}, + }, + { + "chunk_id": "asset", + "type": "image", + "content": "image description", + "path": "refs.pdf/Root/image", + "order": 2, + "file_path": "images/asset.png", + "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() + + calls: list[tuple[str, str]] = [] + original = ReadOnlyChunkStore.load_chunk_reference_metadata + + def record_reference_load( + store: ReadOnlyChunkStore, + document: str, + chunk: str, + ) -> Mapping[str, object] | None: + calls.append((document, chunk)) + return original(store, document, chunk) + + monkeypatch.setattr( + ReadOnlyChunkStore, + "load_chunk_reference_metadata", + record_reference_load, + ) + async with contract_db_session() as db: + snapshot = await load_nav_snapshot( + db, + user_id=_USER_ID, + namespace=namespace, + lazy=True, + ) + + assert calls == [] + assert snapshot.chunk_ref_index[f"{document_id}:asset"]["file_path"] == ( + "images/asset.png" + ) + assert calls == [(document_id, "asset")] + episode = EpisodeResult( + representation="", + steps=[], + scored_chunks=[ + ( + Chunk( + node_id="asset", + doc_id=document_id, + text="image description", + line_ids=(2,), + section_id="root", + ), + 0.75, + ) + ], + kept_chunks=[ + Chunk( + node_id="asset", + doc_id=document_id, + text="image description", + line_ids=(2,), + section_id="root", + ) + ], + evidence_text="image description", + evidence_chars_actual=17, + retrieved_nodes=["asset"], + ) + refs, scores = build_referenced_chunks(episode, snapshot) + assert refs == [ + { + "chunk_id": "asset", + "document_id": document_id, + "chunk_type": "image", + "section_path": "Root / image", + "file_path": "images/asset.png", + "job_id": job_id, + "score": 0.75, + } + ] + assert scores == {"asset": 0.75} + snapshot.close() + + def test_incomplete_index_falls_back_for_duplicate_unit_ids() -> None: first_sections = [ SectionRow("root-a", None, "Root A", "Root A", 0, "", 0), diff --git a/packages/shared-python/shared/models/database/document.py b/packages/shared-python/shared/models/database/document.py index 33dda2f9..ad59857e 100644 --- a/packages/shared-python/shared/models/database/document.py +++ b/packages/shared-python/shared/models/database/document.py @@ -301,6 +301,13 @@ class DocumentMapUnitToken(Base): "token_hash", "map_unit_id", ), + Index( + "idx_document_map_unit_tokens_unit_lookup", + "map_unit_id", + "channel", + "token_hash", + postgresql_include=["token", "frequency"], + ), Index("idx_document_map_unit_tokens_unit", "map_unit_id", "channel"), ) 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 c056a2d2..b7f3f119 100644 --- a/packages/shared-python/shared/services/retrieval/nav/nav_knowhere.py +++ b/packages/shared-python/shared/services/retrieval/nav/nav_knowhere.py @@ -53,6 +53,7 @@ ROOT_SECTION_PATH = "Root" _DEFAULT_DSN = "postgresql://root:root123@127.0.0.1:5433/Knowhere" _MAP_UNIT_INDEX_FORMAT_VERSION = 1 +_MAP_SCORE_CHANNELS: Tuple[str, str] = ("path", "content") _logger = logging.getLogger(__name__) @@ -164,6 +165,13 @@ def knowhere_database_url() -> str: class ChunkStore(Protocol): + def load_chunk_reference_metadata( + self, + document_id: str, + chunk_id: str, + ) -> Optional[Mapping[str, Any]]: + raise NotImplementedError + def load_persisted_score_corpus( self, document_ids: Sequence[str], @@ -265,6 +273,30 @@ def load_section_units( finally: cur.close() + def load_chunk_reference_metadata( + self, + document_id: str, + chunk_id: str, + ) -> Optional[Mapping[str, Any]]: + """Resolve deferred reference fields for one selected chunk.""" + doc_id = str(document_id).strip() + cid = str(chunk_id).strip() + job_result_id = self._revisions.get(doc_id) + if not doc_id or not cid or not job_result_id: + return None + cur = self._connection().cursor() + try: + cur.execute( + "SELECT file_path FROM document_chunks " + "WHERE document_id = %s AND job_result_id = %s AND chunk_id = %s " + "ORDER BY sort_order DESC, id DESC LIMIT 1", + (doc_id, job_result_id, cid), + ) + row = cur.fetchone() + return {"file_path": str(row[0] or "") or None} if row else None + finally: + cur.close() + def load_persisted_score_corpus( self, document_ids: Sequence[str], @@ -404,8 +436,13 @@ def load_persisted_score_corpus( "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), + "AND token = ANY(%s) AND channel = ANY(%s)", + ( + map_unit_ids, + query_token_hashes, + query_tokens, + list(_MAP_SCORE_CHANNELS), + ), ) for map_unit_id, channel, token, frequency in cur.fetchall(): frequencies.setdefault((str(map_unit_id), str(channel)), {})[ diff --git a/packages/shared-python/shared/services/retrieval/nav_snapshot.py b/packages/shared-python/shared/services/retrieval/nav_snapshot.py index 17668767..4eec0404 100644 --- a/packages/shared-python/shared/services/retrieval/nav_snapshot.py +++ b/packages/shared-python/shared/services/retrieval/nav_snapshot.py @@ -11,6 +11,7 @@ import logging import time from dataclasses import dataclass +from collections.abc import Callable, Iterator, Mapping from typing import Any, Protocol from sqlalchemy import Executable, literal, select, tuple_ @@ -53,7 +54,7 @@ class NavSnapshot: """In-memory corpus for one map-nav episode.""" provider: NamespaceKnowhereProvider - chunk_ref_index: dict[str, dict[str, Any]] + chunk_ref_index: Mapping[str, dict[str, Any]] document_ids: list[str] document_titles: dict[str, str] @@ -199,11 +200,9 @@ async def load_nav_snapshot( providers, titles=kept_titles, chunk_owner_by_id={ - chunk_id: meta["document_id"] - for chunk_id, meta in chunk_ref_index.items() - if ":" not in chunk_id - and isinstance(meta, dict) - and meta.get("document_id") + chunk_id: document_id + for document_id, chunk_ids in chunk_ids_by_doc.items() + for chunk_id in chunk_ids }, ) except Exception: @@ -211,7 +210,10 @@ async def load_nav_snapshot( raise return NavSnapshot( provider=provider, - chunk_ref_index=dict(chunk_ref_index), + chunk_ref_index=LazyChunkRefIndex( + chunk_ref_index, + resolver=store.load_chunk_reference_metadata, + ), document_ids=list(provider.document_ids()), document_titles={did: kept_titles.get(did, did) for did in provider.document_ids()}, ) @@ -252,7 +254,7 @@ async def _load_chunk_index( section_path_by_id: dict[str, str], job_id_by_result_id: dict[str, str], ) -> tuple[dict[str, list[str]], dict[str, dict[str, Any]], dict[str, dict[str, Any]]]: - """Load only IDs/reference metadata; content remains lazy.""" + """Load only IDs/reference metadata; content and asset paths remain lazy.""" ids_by_doc: dict[str, list[str]] = {} ref_index: dict[str, dict[str, Any]] = {} root_assets_by_doc: dict[str, set[str]] = {} @@ -271,7 +273,8 @@ async def _load_chunk_index( DocumentChunk.chunk_id, DocumentChunk.section_id, DocumentChunk.chunk_type, - DocumentChunk.file_path, + # Asset paths are only needed for selected references and + # are resolved by ``LazyChunkRefIndex`` at bridge time. DocumentChunk.chunk_metadata["connect_to"].label("connect_to"), DocumentChunk.sort_order, DocumentChunk.id, @@ -323,7 +326,7 @@ async def _load_chunk_index( "document_id": document_id, "section_path": section_path, "chunk_type": chunk_type, - "file_path": str(row[5] or "") or None, + "file_path": None, "job_id": job_id_by_result_id.get(job_result_id), } ids_by_doc.setdefault(document_id, []).append(chunk_id) @@ -335,7 +338,7 @@ async def _load_chunk_index( and section_path == "Root" ): root_assets_by_doc.setdefault(document_id, set()).add(chunk_id) - connections = row[6] + connections = row[5] if isinstance(connections, str) and connections.strip(): try: connections = json.loads(connections) @@ -356,9 +359,9 @@ async def _load_chunk_index( last_key = ( str(last[0]), str(last[1]), - int(last[7] or 0), + int(last[6] or 0), str(last[2] or ""), - str(last[8]), + str(last[7]), ) if len(rows) < _CHUNK_BATCH_SIZE: break @@ -379,6 +382,48 @@ async def _load_chunk_index( return ids_by_doc, ref_index, remounted +class LazyChunkRefIndex(Mapping[str, dict[str, Any]]): + """Reference metadata map that resolves selected asset paths on demand. + + Snapshot construction still records every chunk identity, section path, + type, and job id so navigation ownership is unchanged. ``file_path`` is + fetched only when the exit bridge asks for a selected reference. + """ + + def __init__( + self, + base: Mapping[str, dict[str, Any]], + *, + resolver: Callable[[str, str], Mapping[str, Any] | None], + ) -> None: + self._base = {str(key): dict(value) for key, value in base.items()} + self._resolver = resolver + + def __getitem__(self, key: str) -> dict[str, Any]: + value = self.get(key) + if value is None: + raise KeyError(key) + return value + + def __iter__(self) -> Iterator[str]: + return iter(self._base) + + def __len__(self) -> int: + return len(self._base) + + def get(self, key: str, default: Any = None) -> dict[str, Any] | Any: + value = self._base.get(str(key)) + if value is None: + return default + if value.get("file_path") is None: + document_id = str(value.get("document_id") or "").strip() + chunk_id = str(key).split(":", 1)[-1].strip() + resolved = self._resolver(document_id, chunk_id) + if resolved is not None: + value.update({"file_path": resolved.get("file_path") or None}) + return value + + async def _load_sections( db: SnapshotSession, *, From 00fc7e2981a51ca34a8b5331b7f483efe4a4b9fb Mon Sep 17 00:00:00 2001 From: suguanYang Date: Sat, 29 Aug 2026 17:35:37 +0800 Subject: [PATCH 3/5] perf: trust transactional map index marker --- .../test_retrieval_map_unit_index_contract.py | 5 +++ .../services/retrieval/nav/nav_knowhere.py | 44 +++---------------- 2 files changed, 10 insertions(+), 39 deletions(-) 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 index 56a2dced..eb124fb2 100644 --- a/apps/api/tests/contract/test_retrieval_map_unit_index_contract.py +++ b/apps/api/tests/contract/test_retrieval_map_unit_index_contract.py @@ -274,6 +274,11 @@ def reject_payload_read( await db.execute( delete(DocumentMapUnitToken).where(DocumentMapUnitToken.id == token_id) ) + await db.execute( + delete(DocumentMapUnitIndex).where( + DocumentMapUnitIndex.document_id == document_id + ) + ) await db.commit() incomplete_snapshot = await load_nav_snapshot( db, 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 b7f3f119..2aad0b29 100644 --- a/packages/shared-python/shared/services/retrieval/nav/nav_knowhere.py +++ b/packages/shared-python/shared/services/retrieval/nav/nav_knowhere.py @@ -338,45 +338,11 @@ def load_persisted_score_corpus( ): return None - stage_started = time.perf_counter() - 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 - _logger.info( - "retrieval map-index load stage=unit_count seconds=%.3f rows=%d", - time.perf_counter() - stage_started, - indexed_unit_count, - ) - 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 - stage_started = time.perf_counter() - 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 - _logger.info( - "retrieval map-index load stage=token_count seconds=%.3f rows=%d", - time.perf_counter() - stage_started, - indexed_token_count, - ) - expected_token_count = sum(int(row[4]) for row in manifests) - if indexed_token_count != expected_token_count: - return None + # The marker is written last in the same transaction that inserts + # all units and token rows. A committed marker therefore denotes + # one complete revision snapshot; avoid recounting millions of + # token rows on every request. Any rebuild deletes the marker + # first, so readers fall back to legacy scoring until completion. # 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. From 1f335316e1edb1ff262138504272332ffad38a28 Mon Sep 17 00:00:00 2001 From: suguanYang Date: Sun, 30 Aug 2026 12:40:39 +0800 Subject: [PATCH 4/5] perf: add revision-pinned retrieval serving index --- CONTEXT.md | 70 +++ ...6c7d8e_add_section_snapshot_order_index.py | 41 ++ ...d8e9f_add_retrieval_serving_generations.py | 171 +++++++ .../5b6c7d8e9f0a_add_term_trigram_indexes.py | 35 ++ .../services/documents/lifecycle_service.py | 47 +- apps/api/scripts/backfill_map_unit_indexes.py | 70 ++- .../test_bm25_fts_prefilter_contract.py | 38 +- .../tests/contract/test_documents_contract.py | 86 +++- .../test_retrieval_manifest_cache_contract.py | 56 +++ .../test_retrieval_mapnav_session_contract.py | 4 + ...test_retrieval_relit_map_cache_contract.py | 27 ++ .../test_retrieval_revision_races_contract.py | 102 ++++ .../test_retrieval_rrf_duplicate_contract.py | 19 + ...est_retrieval_serving_manifest_contract.py | 51 ++ ...retrieval_snapshot_consistency_contract.py | 54 +++ .../test_retrieval_term_score_contract.py | 75 +++ ...0005-stream-retrieval-progress-over-sse.md | 47 ++ ...mically-publish-retrieval-serving-index.md | 6 + ...-coherent-retrieval-serving-generations.md | 6 + ...enance-window-for-serving-index-rollout.md | 6 + docs/adr/README.md | 5 +- docs/design/retrieval-serving-index-plan.md | 425 ++++++++++++++++ docs/design/retrieval-streaming-sse.md | 207 ++++++++ .../shared/models/database/document.py | 171 +++++++ .../services/retrieval/execution/plan.py | 26 +- .../retrieval/execution/reference_resolver.py | 3 + .../retrieval/execution/revision_pins.py | 92 ++++ .../retrieval/execution/route_types.py | 3 + .../services/retrieval/execution/routes.py | 81 +++- .../services/retrieval/hydration/connected.py | 22 +- .../services/retrieval/hydration/reference.py | 44 +- .../retrieval/hydration/result_assembly.py | 3 + .../services/retrieval/manifest_cache.py | 47 ++ .../services/retrieval/nav/nav_knowhere.py | 308 +++++++++--- .../services/retrieval/nav/nav_orchestrate.py | 3 + .../services/retrieval/nav/nav_types.py | 6 + .../shared/services/retrieval/nav_snapshot.py | 453 ++++++++++++++---- .../services/retrieval/publication_content.py | 3 + .../services/retrieval/publication_service.py | 41 ++ .../services/retrieval/search/channels.py | 53 +- .../services/retrieval/search/discovery.py | 56 ++- .../services/retrieval/search/ranking.py | 109 +++-- .../retrieval/search/scoped_corpus.py | 119 ++++- .../services/retrieval/search/scoring.py | 10 +- .../services/retrieval/serving_generation.py | 60 +++ .../services/retrieval/serving_manifest.py | 405 ++++++++++++++++ 46 files changed, 3500 insertions(+), 266 deletions(-) create mode 100644 apps/api/alembic/versions/3f4a5b6c7d8e_add_section_snapshot_order_index.py create mode 100644 apps/api/alembic/versions/4a5b6c7d8e9f_add_retrieval_serving_generations.py create mode 100644 apps/api/alembic/versions/5b6c7d8e9f0a_add_term_trigram_indexes.py create mode 100644 apps/api/tests/contract/test_retrieval_manifest_cache_contract.py create mode 100644 apps/api/tests/contract/test_retrieval_relit_map_cache_contract.py create mode 100644 apps/api/tests/contract/test_retrieval_revision_races_contract.py create mode 100644 apps/api/tests/contract/test_retrieval_rrf_duplicate_contract.py create mode 100644 apps/api/tests/contract/test_retrieval_serving_manifest_contract.py create mode 100644 apps/api/tests/contract/test_retrieval_term_score_contract.py create mode 100644 docs/adr/0005-stream-retrieval-progress-over-sse.md create mode 100644 docs/adr/0006-atomically-publish-retrieval-serving-index.md create mode 100644 docs/adr/0007-use-coherent-retrieval-serving-generations.md create mode 100644 docs/adr/0008-use-a-maintenance-window-for-serving-index-rollout.md create mode 100644 docs/design/retrieval-serving-index-plan.md create mode 100644 docs/design/retrieval-streaming-sse.md create mode 100644 packages/shared-python/shared/services/retrieval/execution/revision_pins.py create mode 100644 packages/shared-python/shared/services/retrieval/manifest_cache.py create mode 100644 packages/shared-python/shared/services/retrieval/serving_generation.py create mode 100644 packages/shared-python/shared/services/retrieval/serving_manifest.py diff --git a/CONTEXT.md b/CONTEXT.md index c8903a58..b67e16d0 100644 --- a/CONTEXT.md +++ b/CONTEXT.md @@ -173,6 +173,76 @@ The query workflow that returns cited evidence from published documents. The typed retrieval request that owns cache-shaping fields and route policy: scope, filters, data type, channels, ranking options, and agentic toggle. +### Retrieval Run + +One execution of a Retrieval Query, recorded across classic, map-nav, +small-corpus, cache-hit, failed, and cancelled outcomes with route, timing, +and terminal-status metadata. + +### Retrieval Progress Event + +A safe, user-facing update about the current phase of a Retrieval run. It uses +the fixed phases `started`, `planning`, `searching`, `reviewing_sources`, and +`finalizing`; it never contains chain-of-thought or raw planner output. + +### Retrieval Stream + +The server-to-client event stream for a Retrieval Query. It carries +Retrieval Progress Events during execution and one authoritative final result +or terminal failure, while leaving answer generation to downstream clients. + +### Retrieval Duration + +The end-to-end server time for a Retrieval Query, measured from retrieval +execution start through final public-result assembly. It includes cache lookup +and excludes authentication, network transfer, SSE delivery time, and +downstream answer generation. + +### Retrieval Non-LLM Work + +The database and retrieval-engine work for a Retrieval Query: snapshot or +serving-index loading, lexical scoring, ranking, result hydration, citation +assembly, and asset-reference resolution. It excludes planner, harvest, +control, and answer-generation model time, which are measured separately. + +### Retrieval Serving Index + +The publication-derived read model used to load retrieval structure and +scoring inputs without rebuilding them from the full document corpus for each +query. It is revision-pinned and complete before its document revision becomes +active. + +### Retrieval Serving Fallback + +The exact legacy retrieval path used when a serving index is missing, +incomplete, or inconsistent. It preserves retrieval quality while sacrificing +the serving-index latency target until the derived data is repaired. + +### Retrieval Serving Generation + +The namespace-scoped version that identifies one coherent set of active +document revisions and their serving-index statistics. Retrieval captures one +generation and retries or falls back if publication changes it during capture. + +### Retrieval Semantic Parity + +The compatibility requirement that a serving-index retrieval returns the same +selected chunk IDs, ordering, rounded scores, citations, and asset references +as the legacy retrieval path for the same request. + +### Retrieval Revision Pin + +The set of document revision IDs captured at retrieval start and used for the +entire retrieval run, including lazy content and asset resolution. A later +publication affects subsequent runs, not the run already in progress. + +### Online Retrieval Serving Rollout + +The additive rollout of retrieval-serving schema and derived data while +retrieval and document publication remain available. Incomplete or +inconsistent revisions use the exact legacy retrieval path until backfill and +validation finish. + ### Workflow Run Request The agentic Retrieval request passed through planning and step execution. It diff --git a/apps/api/alembic/versions/3f4a5b6c7d8e_add_section_snapshot_order_index.py b/apps/api/alembic/versions/3f4a5b6c7d8e_add_section_snapshot_order_index.py new file mode 100644 index 00000000..4dcc5011 --- /dev/null +++ b/apps/api/alembic/versions/3f4a5b6c7d8e_add_section_snapshot_order_index.py @@ -0,0 +1,41 @@ +"""Add the index used by lazy map-nav section pagination.""" + +from __future__ import annotations + +from alembic import op + + +revision = "3f4a5b6c7d8e" +down_revision = "2e3f4a5b6c7d" +branch_labels = None +depends_on = None + +_INDEX_NAME = "idx_document_sections_revision_snapshot_order" + + +def upgrade() -> None: + external_transaction = bool( + op.get_context().opts.get("knowhere_external_transaction", False) + ) + statement = ( + f"CREATE INDEX {{concurrently}}IF NOT EXISTS {_INDEX_NAME} " + "ON document_sections " + "(document_id, job_result_id, sort_order, section_id)" + ) + if external_transaction: + op.execute(statement.format(concurrently="")) + return + with op.get_context().autocommit_block(): + op.execute(statement.format(concurrently="CONCURRENTLY ")) + + +def downgrade() -> None: + external_transaction = bool( + op.get_context().opts.get("knowhere_external_transaction", False) + ) + statement = f"DROP INDEX {{concurrently}}IF EXISTS {_INDEX_NAME}" + if external_transaction: + op.execute(statement.format(concurrently="")) + return + with op.get_context().autocommit_block(): + op.execute(statement.format(concurrently="CONCURRENTLY ")) diff --git a/apps/api/alembic/versions/4a5b6c7d8e9f_add_retrieval_serving_generations.py b/apps/api/alembic/versions/4a5b6c7d8e9f_add_retrieval_serving_generations.py new file mode 100644 index 00000000..426ff6e0 --- /dev/null +++ b/apps/api/alembic/versions/4a5b6c7d8e9f_add_retrieval_serving_generations.py @@ -0,0 +1,171 @@ +"""Add revision-pinned serving manifests and namespace statistics.""" + +from __future__ import annotations + +from collections.abc import Sequence + +from alembic import op +import sqlalchemy as sa + + +revision: str = "4a5b6c7d8e9f" +down_revision: str | None = "3f4a5b6c7d8e" +branch_labels: Sequence[str] | None = None +depends_on: Sequence[str] | None = None + + +def upgrade() -> None: + if not sa.inspect(op.get_bind()).has_table("retrieval_namespace_generations"): + op.create_table( + "retrieval_namespace_generations", + sa.Column("id", sa.String(length=100), nullable=False), + sa.Column("user_id", sa.Text(), nullable=False), + sa.Column("namespace", sa.String(length=255), nullable=False), + sa.Column( + "generation", sa.BigInteger(), nullable=False, server_default="0" + ), + sa.Column("updated_at", sa.DateTime(), nullable=False), + sa.PrimaryKeyConstraint("id"), + sa.UniqueConstraint( + "user_id", + "namespace", + name="uq_retrieval_namespace_generations_scope", + ), + ) + if not sa.inspect(op.get_bind()).has_table("retrieval_serving_revision_manifests"): + op.create_table( + "retrieval_serving_revision_manifests", + sa.Column("id", sa.String(length=100), nullable=False), + sa.Column("user_id", sa.Text(), nullable=False), + sa.Column("namespace", sa.String(length=255), 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("payload_zlib", sa.LargeBinary(), nullable=False), + sa.Column("checksum", sa.String(length=64), 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_retrieval_serving_revision_manifests_revision", + ), + ) + if not _index_exists("idx_retrieval_serving_revision_manifests_scope"): + op.create_index( + "idx_retrieval_serving_revision_manifests_scope", + "retrieval_serving_revision_manifests", + ["user_id", "namespace", "document_id", "job_result_id"], + ) + if not sa.inspect(op.get_bind()).has_table("retrieval_serving_revision_stats"): + op.create_table( + "retrieval_serving_revision_stats", + sa.Column("id", sa.String(length=100), nullable=False), + sa.Column("user_id", sa.Text(), nullable=False), + sa.Column("namespace", sa.String(length=255), 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("payload_zlib", sa.LargeBinary(), nullable=False), + sa.Column("checksum", sa.String(length=64), 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_retrieval_serving_revision_stats_revision", + ), + ) + if not _index_exists("idx_retrieval_serving_revision_stats_scope"): + op.create_index( + "idx_retrieval_serving_revision_stats_scope", + "retrieval_serving_revision_stats", + ["user_id", "namespace", "document_id", "job_result_id"], + ) + if not sa.inspect(op.get_bind()).has_table("retrieval_namespace_stats"): + op.create_table( + "retrieval_namespace_stats", + sa.Column("id", sa.String(length=100), nullable=False), + sa.Column("user_id", sa.Text(), nullable=False), + sa.Column("namespace", sa.String(length=255), nullable=False), + sa.Column("generation", sa.BigInteger(), nullable=False), + sa.Column("payload_zlib", sa.LargeBinary(), nullable=False), + sa.Column("checksum", sa.String(length=64), nullable=False), + sa.Column("updated_at", sa.DateTime(), nullable=False), + sa.PrimaryKeyConstraint("id"), + sa.UniqueConstraint( + "user_id", "namespace", name="uq_retrieval_namespace_stats_scope" + ), + ) + if not sa.inspect(op.get_bind()).has_table("retrieval_namespace_token_stats"): + op.create_table( + "retrieval_namespace_token_stats", + sa.Column("id", sa.String(length=100), nullable=False), + sa.Column("user_id", sa.Text(), nullable=False), + sa.Column("namespace", sa.String(length=255), nullable=False), + sa.Column("generation", sa.BigInteger(), nullable=False), + sa.Column("channel", sa.String(length=32), nullable=False), + sa.Column("token_hash", sa.String(length=64), nullable=False), + sa.Column("document_frequency", sa.Integer(), nullable=False), + sa.PrimaryKeyConstraint("id"), + sa.UniqueConstraint( + "user_id", + "namespace", + "channel", + "token_hash", + name="uq_retrieval_namespace_token_stats_key", + ), + ) + if not _index_exists("idx_retrieval_namespace_token_stats_lookup"): + op.create_index( + "idx_retrieval_namespace_token_stats_lookup", + "retrieval_namespace_token_stats", + ["user_id", "namespace", "generation", "channel", "token_hash"], + ) + + +def _index_exists(index_name: str) -> bool: + bind = op.get_bind() + inspector = sa.inspect(bind) + for table_name in inspector.get_table_names(): + if any( + index.get("name") == index_name + for index in inspector.get_indexes(table_name) + ): + return True + return False + + +def downgrade() -> None: + op.drop_index( + "idx_retrieval_namespace_token_stats_lookup", + table_name="retrieval_namespace_token_stats", + if_exists=True, + ) + op.drop_table("retrieval_namespace_token_stats", if_exists=True) + op.drop_table("retrieval_namespace_stats", if_exists=True) + op.drop_index( + "idx_retrieval_serving_revision_stats_scope", + table_name="retrieval_serving_revision_stats", + if_exists=True, + ) + op.drop_table("retrieval_serving_revision_stats", if_exists=True) + op.drop_index( + "idx_retrieval_serving_revision_manifests_scope", + table_name="retrieval_serving_revision_manifests", + if_exists=True, + ) + op.drop_table("retrieval_serving_revision_manifests", if_exists=True) + op.drop_table("retrieval_namespace_generations", if_exists=True) diff --git a/apps/api/alembic/versions/5b6c7d8e9f0a_add_term_trigram_indexes.py b/apps/api/alembic/versions/5b6c7d8e9f0a_add_term_trigram_indexes.py new file mode 100644 index 00000000..71720ff0 --- /dev/null +++ b/apps/api/alembic/versions/5b6c7d8e9f0a_add_term_trigram_indexes.py @@ -0,0 +1,35 @@ +"""Add trigram acceleration for exact term-channel candidate discovery.""" + +from __future__ import annotations + +from collections.abc import Sequence + +from alembic import op + + +revision: str = "5b6c7d8e9f0a" +down_revision: str | None = "4a5b6c7d8e9f" +branch_labels: Sequence[str] | None = None +depends_on: Sequence[str] | None = None + +_MAP_UNIT_INDEX = "idx_document_map_units_term_trgm" +_CHUNK_INDEX = "idx_document_chunks_term_trgm" + + +def upgrade() -> None: + op.execute("CREATE EXTENSION IF NOT EXISTS pg_trgm") + op.execute( + f"CREATE INDEX IF NOT EXISTS {_MAP_UNIT_INDEX} " + "ON document_map_units USING gin " + "(term_search_text_lower gin_trgm_ops)" + ) + op.execute( + f"CREATE INDEX IF NOT EXISTS {_CHUNK_INDEX} " + "ON document_chunks USING gin " + "(lower(COALESCE(term_search_text, '')) gin_trgm_ops)" + ) + + +def downgrade() -> None: + op.execute(f"DROP INDEX IF EXISTS {_CHUNK_INDEX}") + op.execute(f"DROP INDEX IF EXISTS {_MAP_UNIT_INDEX}") diff --git a/apps/api/app/services/documents/lifecycle_service.py b/apps/api/app/services/documents/lifecycle_service.py index 1bd5ae3a..6b39507d 100644 --- a/apps/api/app/services/documents/lifecycle_service.py +++ b/apps/api/app/services/documents/lifecycle_service.py @@ -8,13 +8,25 @@ from app.repositories.document_repository import DocumentRepository from loguru import logger +from sqlalchemy import delete from sqlalchemy.ext.asyncio import AsyncSession -from shared.models.database.document import DocumentChunk, DocumentSection +from shared.models.database.document import ( + DocumentChunk, + DocumentSection, + RetrievalServingRevisionStat, +) from shared.services.retrieval.cache_service import ( invalidate_retrieval_cache_namespaces, ) from shared.services.retrieval.graph.service import DocumentGraphService, GraphScope +from shared.services.retrieval.serving_generation import ( + advance_namespace_generation, + lock_namespace_generation, +) +from shared.services.retrieval.serving_manifest import ( + rebuild_namespace_serving_statistics, +) from shared.services.storage.result_storage import ResultStorage, get_result_storage _DOCUMENT_CHUNK_ASSET_URL_EXPIRES_SECONDS = 7 * 24 * 60 * 60 @@ -103,7 +115,7 @@ def _normalize_page_asset(raw_asset: dict[str, Any]) -> dict[str, Any] | None: "content_type": content_type, "source": source, } - if (asset_url := str(raw_asset.get("asset_url") or "").strip()): + if asset_url := str(raw_asset.get("asset_url") or "").strip(): asset["asset_url"] = asset_url if (width := _positive_int(raw_asset.get("width"))) is not None: asset["width"] = width @@ -446,7 +458,38 @@ async def archive_document( return document_payload(document) previous_namespace = document.namespace + await db.run_sync( + lambda sync_db: lock_namespace_generation( + sync_db, + user_id=user_id, + namespace=previous_namespace, + ) + ) await self._repository.archive_document(db, document=document) + current_revision = document.current_job_result_id + if current_revision: + await db.run_sync( + lambda sync_db: sync_db.execute( + delete(RetrievalServingRevisionStat).where( + RetrievalServingRevisionStat.document_id == document_id, + RetrievalServingRevisionStat.job_result_id == current_revision, + ) + ) + ) + await db.run_sync( + lambda sync_db: rebuild_namespace_serving_statistics( + sync_db, + user_id=user_id, + namespace=previous_namespace, + ) + ) + await db.run_sync( + lambda sync_db: advance_namespace_generation( + sync_db, + user_id=user_id, + namespace=previous_namespace, + ) + ) await db.run_sync( lambda sync_db: self._graph_service.remove_document_graph( sync_db, diff --git a/apps/api/scripts/backfill_map_unit_indexes.py b/apps/api/scripts/backfill_map_unit_indexes.py index 78a88f6c..6a8a79db 100644 --- a/apps/api/scripts/backfill_map_unit_indexes.py +++ b/apps/api/scripts/backfill_map_unit_indexes.py @@ -33,6 +33,14 @@ def _bootstrap_python_path() -> None: 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 +from shared.services.retrieval.serving_generation import ( + advance_namespace_generation, + lock_namespace_generation, +) +from shared.services.retrieval.serving_manifest import ( + persist_revision_serving_state, + rebuild_namespace_serving_statistics, +) def _build_parser() -> argparse.ArgumentParser: @@ -44,14 +52,20 @@ def _build_parser() -> argparse.ArgumentParser: action="store_true", help="Build and commit each current revision index.", ) - parser.add_argument("--document-id", default="", help="Limit the backfill to one document.") + 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)) + statement = ( + select(Document) + .where(Document.status == "active") + .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) @@ -62,7 +76,9 @@ 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}") + print( + f"would backfill document={document.document_id} revision={document.current_job_result_id}" + ) return len(documents) session_factory = get_sync_session_factory() @@ -70,15 +86,49 @@ def backfill_map_unit_indexes(*, apply: bool, document_id: str = "") -> int: 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: + lock_namespace_generation( + db, + user_id=document.user_id, + namespace=document.namespace, + ) + locked_document = db.execute( + select(Document) + .where(Document.document_id == document.document_id) + .with_for_update() + ).scalar_one_or_none() + if ( + locked_document is None + or locked_document.status != "active" + or locked_document.current_job_result_id != job_result_id + or locked_document.user_id != document.user_id + or locked_document.namespace != document.namespace + ): + db.rollback() + print( + f"skipped stale or inactive document={document.document_id} " + f"revision={job_result_id}" + ) + continue + scope = DocumentPublicationScope( + user_id=locked_document.user_id, + namespace=locked_document.namespace, + document_id=locked_document.document_id, + job_result_id=job_result_id, + source_file_name=str(locked_document.source_file_name or ""), + ) replace_document_map_units(db, scope=scope) + persist_revision_serving_state(db, scope=scope) + rebuild_namespace_serving_statistics( + db, + user_id=scope.user_id, + namespace=scope.namespace, + ) + advance_namespace_generation( + db, + user_id=scope.user_id, + namespace=scope.namespace, + ) db.commit() print(f"backfilled document={document.document_id} revision={job_result_id}") return len(documents) diff --git a/apps/api/tests/contract/test_bm25_fts_prefilter_contract.py b/apps/api/tests/contract/test_bm25_fts_prefilter_contract.py index 6b958329..56add8e5 100644 --- a/apps/api/tests/contract/test_bm25_fts_prefilter_contract.py +++ b/apps/api/tests/contract/test_bm25_fts_prefilter_contract.py @@ -23,10 +23,10 @@ user_id TEXT, namespace TEXT, status TEXT, - current_job_result_id INTEGER, + current_job_result_id TEXT, source_file_name TEXT ); -CREATE TABLE job_results (id INTEGER PRIMARY KEY, job_id TEXT); +CREATE TABLE job_results (id TEXT PRIMARY KEY, job_id TEXT); CREATE TABLE document_sections (section_id TEXT PRIMARY KEY, section_path TEXT); CREATE TABLE document_chunks ( id SERIAL PRIMARY KEY, @@ -38,7 +38,7 @@ source_chunk_path TEXT, file_path TEXT, chunk_metadata JSONB, - job_result_id INTEGER, + job_result_id TEXT, sort_order INTEGER, content_search_text TEXT, content_search_tsv TSVECTOR GENERATED ALWAYS AS @@ -197,3 +197,35 @@ async def test_exclusions_still_apply_under_the_prefilter( exclude_sections=[], ) assert rows == [] + + +@pytest.mark.asyncio +async def test_content_channel_uses_the_requested_revision_pin( + seeded_session: AsyncSession, +) -> None: + await seeded_session.execute(text("INSERT INTO job_results VALUES (2, 'job2')")) + await seeded_session.execute( + text( + "INSERT INTO document_chunks " + "(chunk_id, document_id, section_id, chunk_type, content, " + " job_result_id, sort_order, content_search_text, path_search_text) " + "VALUES ('new-hit', 'd1', 's1', 'text', 'new body', 2, 1, " + " 'alpha replacement', 'new path')" + ) + ) + await seeded_session.execute( + text("UPDATE documents SET current_job_result_id = 2 WHERE document_id = 'd1'") + ) + + rows = await content_channel( + seeded_session, + user_id="u1", + namespace="ns1", + query="alpha", + top_k=50, + exclude_document_ids=[], + exclude_sections=[], + revision_pins={"d1": "1"}, + ) + + assert [str(row["chunk_id"]) for row in rows] == ["hit-en"] diff --git a/apps/api/tests/contract/test_documents_contract.py b/apps/api/tests/contract/test_documents_contract.py index 6b8e33b0..a71d4029 100644 --- a/apps/api/tests/contract/test_documents_contract.py +++ b/apps/api/tests/contract/test_documents_contract.py @@ -13,6 +13,7 @@ from tests.support.contract_database import ContractDatabase from shared.testing.contract_runtime import get_contract_database_url +from shared.services.retrieval.serving_manifest import encode_serving_manifest async def _create_contract_engine() -> AsyncEngine: @@ -1019,9 +1020,7 @@ async def test_should_include_media_asset_urls_in_document_chunk_list_when_reque assert chunks[1]["asset_url"] == expected_asset_url assert default_response.status_code == 200 - default_chunks = cast( - list[dict[str, object]], default_response.json()["chunks"] - ) + default_chunks = cast(list[dict[str, object]], default_response.json()["chunks"]) assert default_chunks[1]["asset_url"] is None @@ -1334,3 +1333,84 @@ async def test_should_archive_a_document_via_the_legacy_archive_route( assert response_json["archived_at"] assert persisted_document["status"] == "archived" assert persisted_document["archived_at"] is not None + + +@pytest.mark.asyncio +async def test_archive_removes_revision_serving_stats_and_advances_generation( + developer_api_client_factory: Callable[ + [], AbstractAsyncContextManager[AsyncClient] + ], +) -> None: + document_id = f"doc_{uuid4().hex[:12]}" + namespace = f"archive-serving-{uuid4().hex[:8]}" + async with developer_api_client_factory() as api_client: + revision = await _insert_document_revision_with_chunks( + document_id=document_id, + namespace=namespace, + chunks=[ + { + "id": f"dchk_{uuid4().hex[:12]}", + "chunk_id": "archive-serving-chunk", + "chunk_type": "text", + "content": "serving contribution", + "source_chunk_path": "Archive/Serving", + "metadata": {}, + } + ], + ) + payload_bytes, checksum, version = encode_serving_manifest( + { + "document_id": document_id, + "job_result_id": revision["job_result_id"], + "unit_count": 1, + "path_token_count": 1, + "content_token_count": 2, + "token_frequencies": { + "path": {"archive": 1}, + "content": {"serving": 1}, + }, + } + ) + await ContractDatabase.execute( + """ + INSERT INTO retrieval_serving_revision_stats ( + id, user_id, namespace, document_id, job_result_id, + format_version, payload_zlib, checksum, created_at + ) VALUES ( + :id, :user_id, :namespace, :document_id, :job_result_id, + :format_version, :payload_zlib, :checksum, NOW() + ) + """, + { + "id": f"rss_{uuid4().hex[:12]}", + "user_id": "local-dev-user", + "namespace": namespace, + "document_id": document_id, + "job_result_id": revision["job_result_id"], + "format_version": version, + "payload_zlib": payload_bytes, + "checksum": checksum, + }, + ) + response = await api_client.post(f"/api/v1/documents/{document_id}/archive") + + assert response.status_code == 200 + remaining_revision_stats = await ContractDatabase.fetch_one( + """ + SELECT id + FROM retrieval_serving_revision_stats + WHERE document_id = :document_id AND job_result_id = :job_result_id + """, + {"document_id": document_id, "job_result_id": revision["job_result_id"]}, + ) + namespace_stats = await ContractDatabase.fetch_one( + """ + SELECT generation + FROM retrieval_namespace_stats + WHERE user_id = :user_id AND namespace = :namespace + """, + {"user_id": "local-dev-user", "namespace": namespace}, + ) + assert remaining_revision_stats is None + assert namespace_stats is not None + assert int(namespace_stats["generation"]) >= 1 diff --git a/apps/api/tests/contract/test_retrieval_manifest_cache_contract.py b/apps/api/tests/contract/test_retrieval_manifest_cache_contract.py new file mode 100644 index 00000000..e2f73e47 --- /dev/null +++ b/apps/api/tests/contract/test_retrieval_manifest_cache_contract.py @@ -0,0 +1,56 @@ +from __future__ import annotations + +import pytest +from sqlalchemy.exc import SQLAlchemyError + +from shared.services.retrieval.manifest_cache import ( + cache_manifest_payloads, + get_cached_manifest_payloads, +) +from shared.services.retrieval.search.scoped_corpus import count_manifest_chunks + + +class _SessionWithInfo: + def __init__(self) -> None: + self.info: dict[str, object] = {} + + +def test_manifest_payload_cache_is_scoped_to_revision_pin_set() -> None: + session = _SessionWithInfo() + revisions = {"doc-a": "result-a", "doc-b": "result-b"} + payloads = { + ("doc-a", "result-a"): {"chunks": [{"chunk_id": "chunk-a"}]}, + ("doc-b", "result-b"): {"chunks": [{"chunk_id": "chunk-b"}]}, + } + + cache_manifest_payloads(session, revisions=revisions, payloads=payloads) + + assert get_cached_manifest_payloads(session, revisions=revisions) == payloads + assert get_cached_manifest_payloads( + session, + revisions={"doc-a": "different-result"}, + ) is None + + +class _UnavailableManifestSession: + def __init__(self) -> None: + self.rollback_count = 0 + + async def execute(self, _statement: object) -> object: + raise SQLAlchemyError("serving manifest table is unavailable") + + async def rollback(self) -> None: + self.rollback_count += 1 + + +@pytest.mark.asyncio +async def test_manifest_count_falls_back_after_derived_table_error() -> None: + session = _UnavailableManifestSession() + + result = await count_manifest_chunks( + session, # type: ignore[arg-type] + revision_pins={"doc-a": "result-a"}, + ) + + assert result is None + assert session.rollback_count == 1 diff --git a/apps/api/tests/contract/test_retrieval_mapnav_session_contract.py b/apps/api/tests/contract/test_retrieval_mapnav_session_contract.py index a277c4c0..d267941b 100644 --- a/apps/api/tests/contract/test_retrieval_mapnav_session_contract.py +++ b/apps/api/tests/contract/test_retrieval_mapnav_session_contract.py @@ -135,12 +135,14 @@ async def fake_resolve_workflow_references( namespace: str, refs: list[RouteRow], score_by_chunk_id: dict[str, float] | None = None, + revision_pins: dict[str, str] | None = None, ) -> ResolvedWorkflowReferences: assert db is fresh_db assert user_id == "contract-user" assert namespace == "contract-namespace" assert refs assert score_by_chunk_id is not None + assert revision_pins is None events.append("resolve_references") row = { "document_id": "doc_contract", @@ -159,11 +161,13 @@ async def fake_assemble_retrieval_results( exclude_document_ids: list[str], exclude_sections: list[dict[str, str]], allowed_chunk_types: set[str] | None, + revision_pins: dict[str, str] | None = None, ) -> list[RouteRow]: assert db is fresh_db assert exclude_document_ids == [] assert exclude_sections == [] assert allowed_chunk_types is None + assert revision_pins is None events.append("assemble_results") return rows diff --git a/apps/api/tests/contract/test_retrieval_relit_map_cache_contract.py b/apps/api/tests/contract/test_retrieval_relit_map_cache_contract.py new file mode 100644 index 00000000..b8257c9c --- /dev/null +++ b/apps/api/tests/contract/test_retrieval_relit_map_cache_contract.py @@ -0,0 +1,27 @@ +from __future__ import annotations + +from shared.services.retrieval.nav.nav_orchestrate import _relit_map +from shared.services.retrieval.nav.nav_types import NavState + + +def test_relit_map_reuses_same_query_within_episode(monkeypatch) -> None: + calls: list[str] = [] + + def fake_relight_map_for_query(*_args, **kwargs): + calls.append(str(kwargs["query"])) + return {"section": 1.0}, {"unit": 2.0}, ["section"] + + monkeypatch.setattr( + "shared.services.retrieval.nav.nav_map_scores.relight_map_for_query", + fake_relight_map_for_query, + ) + state = NavState(doc_id="", query="retrieval") + config = type("Config", (), {"collect_top_k": 6})() + + with _relit_map(None, state, config, query="retrieval"): + pass + with _relit_map(None, state, config, query="retrieval"): + pass + + assert calls == ["retrieval"] + assert state.relit_map_cache["retrieval"][0] == {"section": 1.0} diff --git a/apps/api/tests/contract/test_retrieval_revision_races_contract.py b/apps/api/tests/contract/test_retrieval_revision_races_contract.py new file mode 100644 index 00000000..dd73d966 --- /dev/null +++ b/apps/api/tests/contract/test_retrieval_revision_races_contract.py @@ -0,0 +1,102 @@ +"""Deterministic contracts for revision and channel-session coherence.""" + +from __future__ import annotations + +from collections.abc import AsyncGenerator +from contextlib import asynccontextmanager +from typing import Any, cast + +import pytest +from sqlalchemy.ext.asyncio import AsyncSession + +from shared.services.retrieval.execution.revision_pins import ( + RetrievalRevisionPins, + is_revision_generation_stable, +) +from shared.services.retrieval.search import discovery + + +@pytest.mark.asyncio +async def test_classic_channels_share_pins_but_use_distinct_sessions( + monkeypatch: pytest.MonkeyPatch, +) -> None: + pins = RetrievalRevisionPins( + revisions={"doc-1": "revision-1"}, + generation=7, + ) + sessions: list[object] = [] + observed: list[tuple[object, object]] = [] + + @asynccontextmanager + async def fake_context() -> AsyncGenerator[object, None]: + session = object() + sessions.append(session) + yield session + + async def fake_channel( + db: AsyncSession, + **kwargs: Any, + ) -> list[dict[str, Any]]: + observed.append((db, kwargs["revision_pins"])) + return [] + + monkeypatch.setattr("shared.core.database.get_db_context", fake_context) + monkeypatch.setattr(discovery, "path_channel", fake_channel) + monkeypatch.setattr(discovery, "content_channel", fake_channel) + monkeypatch.setattr(discovery, "term_channel", fake_channel) + + result = await discovery.bottom_discovery( + cast(AsyncSession, object()), + user_id="user-1", + namespace="namespace-1", + query="coherent query", + top_k=3, + exclude_document_ids=[], + exclude_sections=[], + revision_pins=pins, + ) + + assert result.status == "discovery_done" + assert len(sessions) == 3 + assert len({id(session) for session in sessions}) == 3 + assert len(observed) == 3 + assert {id(session) for session, _pins in observed} == { + id(session) for session in sessions + } + assert all(observed_pins is pins for _session, observed_pins in observed) + + +class _GenerationResult: + def __init__(self, value: int | None) -> None: + self._value = value + + def scalar_one_or_none(self) -> int | None: + return self._value + + +class _GenerationSession: + def __init__(self, values: list[int | None]) -> None: + self._values = iter(values) + + async def execute(self, _statement: object) -> _GenerationResult: + return _GenerationResult(next(self._values)) + + +@pytest.mark.asyncio +async def test_generation_change_is_detected_before_scoring() -> None: + pins = RetrievalRevisionPins(revisions={"doc-1": "revision-1"}, generation=7) + stable = await is_revision_generation_stable( + cast(AsyncSession, _GenerationSession([7])), + user_id="user-1", + namespace="namespace-1", + pins=pins, + ) + changed = await is_revision_generation_stable( + cast(AsyncSession, _GenerationSession([8])), + user_id="user-1", + namespace="namespace-1", + pins=pins, + ) + + assert stable is True + assert changed is False diff --git a/apps/api/tests/contract/test_retrieval_rrf_duplicate_contract.py b/apps/api/tests/contract/test_retrieval_rrf_duplicate_contract.py new file mode 100644 index 00000000..e056af6c --- /dev/null +++ b/apps/api/tests/contract/test_retrieval_rrf_duplicate_contract.py @@ -0,0 +1,19 @@ +"""Contracts for duplicate chunk handling in reciprocal-rank fusion.""" + +from __future__ import annotations + +from shared.services.retrieval.search.scoring import merge_channels_rrf + + +def test_rrf_counts_each_chunk_once_per_channel() -> None: + rows = [ + {"chunk_id": "shared", "document_id": "doc-a"}, + {"chunk_id": "shared", "document_id": "doc-b"}, + {"chunk_id": "other", "document_id": "doc-c"}, + ] + + result = merge_channels_rrf([rows], [1.0], top_k=3) + + assert [row["chunk_id"] for row in result] == ["shared", "other"] + assert result[0]["score"] == round(1.0 / 61, 6) + assert result[1]["score"] == round(1.0 / 62, 6) diff --git a/apps/api/tests/contract/test_retrieval_serving_manifest_contract.py b/apps/api/tests/contract/test_retrieval_serving_manifest_contract.py new file mode 100644 index 00000000..5bb68c06 --- /dev/null +++ b/apps/api/tests/contract/test_retrieval_serving_manifest_contract.py @@ -0,0 +1,51 @@ +"""Contract tests for serving-manifest integrity and version handling.""" + +from __future__ import annotations + +import pytest + +from shared.services.retrieval.serving_manifest import ( + SERVING_MANIFEST_FORMAT_VERSION, + decode_serving_manifest, + encode_serving_manifest, +) + + +def test_serving_manifest_round_trip_preserves_canonical_payload() -> None: + payload = { + "document_id": "doc_contract", + "job_result_id": "result_contract", + "sections": [{"section_id": "sec_1", "sort_order": 0}], + "chunks": [{"chunk_id": "chunk_1", "connect_to": []}], + } + + compressed, checksum, version = encode_serving_manifest(payload) + + assert version == SERVING_MANIFEST_FORMAT_VERSION + assert decode_serving_manifest( + compressed, + checksum=checksum, + format_version=version, + ) == payload + + +def test_serving_manifest_rejects_checksum_mismatch() -> None: + compressed, _, version = encode_serving_manifest({"document_id": "doc"}) + + with pytest.raises(ValueError, match="checksum mismatch"): + decode_serving_manifest( + compressed, + checksum="0" * 64, + format_version=version, + ) + + +def test_serving_manifest_rejects_unknown_version() -> None: + compressed, checksum, _ = encode_serving_manifest({"document_id": "doc"}) + + with pytest.raises(ValueError, match="unsupported serving manifest version"): + decode_serving_manifest( + compressed, + checksum=checksum, + format_version=SERVING_MANIFEST_FORMAT_VERSION + 1, + ) diff --git a/apps/api/tests/contract/test_retrieval_snapshot_consistency_contract.py b/apps/api/tests/contract/test_retrieval_snapshot_consistency_contract.py index 64aaf4c0..bbe03e05 100644 --- a/apps/api/tests/contract/test_retrieval_snapshot_consistency_contract.py +++ b/apps/api/tests/contract/test_retrieval_snapshot_consistency_contract.py @@ -6,8 +6,12 @@ from uuid import uuid4 from httpx import AsyncClient +import pytest from sqlalchemy import Executable, Result +from shared.services.retrieval.execution.reference_resolver import ( + resolve_workflow_references, +) from shared.services.retrieval.nav_snapshot import SnapshotSession, load_nav_snapshot from tests.support.retrieval_snapshot_support import contract_db_session from tests.support.contract_database import ContractDatabase @@ -180,3 +184,53 @@ async def publish_new_revision() -> None: assert snapshot.chunk_ref_index[chunks[0].chunk_id]["section_path"] == ( "republished.pdf/old" ) + + +@pytest.mark.asyncio +async def test_reference_hydration_keeps_the_captured_revision_after_republish( + developer_api_client_factory: Callable[ + [], AbstractAsyncContextManager[AsyncClient] + ], +) -> None: + namespace = f"revision-hydration-race-{uuid4().hex[:8]}" + async with developer_api_client_factory(): + document_id, new_result_id = await _seed_republished_document(namespace) + revision_rows = await ContractDatabase.fetch_all( + """ + SELECT job_result_id, chunk_id + FROM document_chunks + WHERE document_id = :document_id + ORDER BY job_result_id + """, + {"document_id": document_id}, + ) + old_revision = next( + row for row in revision_rows if row["job_result_id"] != new_result_id + ) + await ContractDatabase.execute( + """ + UPDATE documents + SET current_job_result_id = :new_result_id + WHERE document_id = :document_id + """, + {"new_result_id": new_result_id, "document_id": document_id}, + ) + + async with contract_db_session() as db: + resolved = await resolve_workflow_references( + db=db, + user_id=_USER_ID, + namespace=namespace, + refs=[ + { + "document_id": document_id, + "chunk_id": old_revision["chunk_id"], + } + ], + revision_pins={document_id: old_revision["job_result_id"]}, + ) + + assert [row["content"] for row in resolved.rows] == ["old content"] + assert [row["job_result_id"] for row in resolved.rows] == [ + old_revision["job_result_id"] + ] diff --git a/apps/api/tests/contract/test_retrieval_term_score_contract.py b/apps/api/tests/contract/test_retrieval_term_score_contract.py new file mode 100644 index 00000000..425f4de0 --- /dev/null +++ b/apps/api/tests/contract/test_retrieval_term_score_contract.py @@ -0,0 +1,75 @@ +"""Contracts for persisted map-unit term scoring.""" + +from __future__ import annotations + +from shared.services.retrieval.nav.nav_knowhere import ReadOnlyChunkStore + + +class _FakeCursor: + def __init__(self, rows: list[tuple[str, str]]) -> None: + self.rows = rows + self.statement = "" + self.parameters: object = None + + def execute(self, statement: str, parameters: object) -> None: + self.statement = statement + self.parameters = parameters + + def fetchall(self) -> list[tuple[str, str]]: + return self.rows + + +def _store() -> ReadOnlyChunkStore: + return ReadOnlyChunkStore.__new__(ReadOnlyChunkStore) + + +def test_term_scores_keep_literal_substring_and_token_hit_semantics() -> None: + cursor = _FakeCursor( + [ + ("unit-full", "prefix alpha beta suffix"), + ("unit-token", "prefix alpha gamma suffix"), + ] + ) + queries = ["alpha beta", "", "alpha"] + query_tokens = { + "alpha beta": ["alpha", "beta"], + "": [], + "alpha": ["alpha"], + } + + scores = _store()._load_term_scores( + cursor, # type: ignore[arg-type] + map_unit_ids=["unit-full", "unit-token", "unit-miss"], + queries=queries, + query_tokens_by_query=query_tokens, + ) + + assert scores == { + "unit-full": (100.0, 0.0, 100.0), + "unit-token": (1.0, 0.0, 100.0), + } + assert "LIKE ANY" in cursor.statement + assert "POSITION" not in cursor.statement + parameters = cursor.parameters + assert isinstance(parameters, tuple) + assert parameters[0] == ["unit-full", "unit-token", "unit-miss"] + assert parameters[1] == ["%alpha beta%", "%alpha%", "%beta%"] + + +def test_long_query_uses_constant_shape_candidate_sql() -> None: + cursor = _FakeCursor([]) + tokens = [f"token-{index}" for index in range(300)] + query = " ".join(tokens) + + _store()._load_term_scores( + cursor, # type: ignore[arg-type] + map_unit_ids=["unit-1"], + queries=[query], + query_tokens_by_query={query: tokens}, + ) + + assert cursor.statement.count("POSITION") == 0 + assert "LIKE ANY" not in cursor.statement + parameters = cursor.parameters + assert isinstance(parameters, tuple) + assert parameters == (["unit-1"],) diff --git a/docs/adr/0005-stream-retrieval-progress-over-sse.md b/docs/adr/0005-stream-retrieval-progress-over-sse.md new file mode 100644 index 00000000..dfaafbd8 --- /dev/null +++ b/docs/adr/0005-stream-retrieval-progress-over-sse.md @@ -0,0 +1,47 @@ +# Stream Retrieval Progress Over SSE + +## Status + +Accepted + +## Context + +Online Brain users currently wait for the complete Retrieval response while +map-nav planning, searching, and source review run. Knowhere owns retrieval and +evidence, while answer generation belongs to downstream clients. A streaming +contract must improve perceived latency without exposing chain-of-thought or +making partial citations authoritative. + +## Decision + +Add a v2-only `POST /v2/retrieval/query/stream` endpoint using Server-Sent +Events. The endpoint accepts the full `RetrievalQueryRequest` shape and is +live-only: disconnecting cooperatively cancels the retrieval run, and retries +start a new run. In-progress events use a fixed, route-aware phase vocabulary +(`started`, `planning`, `searching`, `reviewing_sources`, `finalizing`) and may +include only safe aggregate counts. The stream terminates with a versioned +envelope containing either the existing retrieval response as the authoritative +result or a typed, user-safe failure (`failed`, `cancelled`, or `no_results`). + +The synchronous map-nav engine remains intact and publishes sanitized progress +through an optional step callback bridged to the SSE route by an async queue. +Answer-token streaming remains downstream. The endpoint sends heartbeats, +disables proxy buffering, and uses per-connection event IDs without promising +replay in the first version. + +Retrieval duration is measured from the execution plan's start through final +public-result assembly. The same duration definition is used for persisted +`retrieval_runs.latency_ms` and latency aggregates; cache hits are recorded too. +Trace persistence must receive the execution start timestamp rather than +starting its own timer after navigation and hydration. Timing and terminal +outcomes must cover classic, map-nav, small-corpus, cache-hit, failed, and +cancelled routes. `retrieval_runs` is the ledger for all of those routes and +stores explicit route, cache, latency, and terminal-status fields. + +## Consequences + +The existing JSON retrieval endpoint remains backward compatible, while SDKs +and Online Brain clients need a new streaming adapter and UI state model. A +future resumable stream would require durable event replay and is deliberately +out of scope. Partial evidence and provisional citations are also deferred +until their grounding and revision semantics are defined. diff --git a/docs/adr/0006-atomically-publish-retrieval-serving-index.md b/docs/adr/0006-atomically-publish-retrieval-serving-index.md new file mode 100644 index 00000000..6f0b0800 --- /dev/null +++ b/docs/adr/0006-atomically-publish-retrieval-serving-index.md @@ -0,0 +1,6 @@ +# Atomically publish the retrieval-serving index + +- Status: Accepted +- Context: Retrieval will use a persistent derived serving index to avoid rebuilding a large namespace on every first request. A document revision without a complete index would have unpredictable latency and could produce inconsistent scoring metadata. +- Decision: Build the serving manifest and scoring statistics in the same database transaction as the document revision. Write the completeness marker last. If serving-index construction fails, roll back the publication and retry the job; do not expose an active revision with a partial serving index. +- Consequences: Active revisions have a simple completeness invariant and predictable first-request behavior. Publication takes more work and storage, and an index failure can delay publication, but retrieval can retain a guarded legacy fallback for migrations or already-existing incomplete revisions. diff --git a/docs/adr/0007-use-coherent-retrieval-serving-generations.md b/docs/adr/0007-use-coherent-retrieval-serving-generations.md new file mode 100644 index 00000000..89d2ca05 --- /dev/null +++ b/docs/adr/0007-use-coherent-retrieval-serving-generations.md @@ -0,0 +1,6 @@ +# Use coherent retrieval-serving generations + +- Status: Accepted +- Context: A namespace can contain many active document revisions, and publication can replace them while a retrieval request is loading serving metadata and scoring statistics. +- Decision: Assign each namespace a serving generation. Retrieval captures one generation and verifies it across serving reads; if it changes, retry once and use the exact legacy path if consistency cannot be established. +- Consequences: Retrieval never combines incompatible revision metadata and scoring statistics. Publication and retrieval need a small amount of generation bookkeeping, and rare concurrent updates may cause a retry or slower fallback. diff --git a/docs/adr/0008-use-a-maintenance-window-for-serving-index-rollout.md b/docs/adr/0008-use-a-maintenance-window-for-serving-index-rollout.md new file mode 100644 index 00000000..e7db06bf --- /dev/null +++ b/docs/adr/0008-use-a-maintenance-window-for-serving-index-rollout.md @@ -0,0 +1,6 @@ +# Roll out the serving index online + +- Status: Accepted +- Context: The server and document publication must remain available while serving-index schema changes and backfill are introduced. +- Decision: Use additive online migrations and bounded idempotent backfill. The retrieval reader automatically uses the serving index only when a revision is complete and consistent; otherwise it uses the exact legacy reader. New publication continues online and builds serving data atomically before activating a revision. Do not expose partial serving data. +- Consequences: There is no planned retrieval or publication downtime. Backfill consumes bounded database resources and some revisions remain on the slower legacy path until complete. Generation checks and stale-revision guards are required while backfill and publication run concurrently. diff --git a/docs/adr/README.md b/docs/adr/README.md index 68e15b22..00d7b9da 100644 --- a/docs/adr/README.md +++ b/docs/adr/README.md @@ -19,4 +19,7 @@ Use this shape: | [0002](0002-use-typed-workflow-outcomes.md) | Use typed workflow outcomes | | [0003](0003-keep-retrieval-workflow-policy-explicit.md) | Keep retrieval workflow policy explicit | | [0004](0004-anonymous-self-hosted-telemetry.md) | Anonymous self-hosted telemetry | -| \ No newline at end of file +| [0005](0005-stream-retrieval-progress-over-sse.md) | Stream retrieval progress over SSE | +| [0006](0006-atomically-publish-retrieval-serving-index.md) | Atomically publish the retrieval-serving index | +| [0007](0007-use-coherent-retrieval-serving-generations.md) | Use coherent retrieval-serving generations | +| [0008](0008-use-a-maintenance-window-for-serving-index-rollout.md) | Roll out the serving index online | diff --git a/docs/design/retrieval-serving-index-plan.md b/docs/design/retrieval-serving-index-plan.md new file mode 100644 index 00000000..5e66f6d6 --- /dev/null +++ b/docs/design/retrieval-serving-index-plan.md @@ -0,0 +1,425 @@ +# Retrieval-serving index: online performance plan + +**Status:** Proposed for review +**Reviewed against:** current `knowhere` retrieval and publication code, 2026-08-29 +**Scope:** first-request retrieval performance; LLM planner/harvest/control time is excluded + +## 1. Goal and non-goals + +The target is predictable, bounded **Retrieval Non-LLM Work** on the current +production-sized namespace (about 643 active documents, 50k sections, and 60k +chunks). The measured baseline is roughly 160-170 seconds before map-nav's LLM +episode begins. We do not require a one-second absolute target in this phase; +we require that the request path avoid repeated full-corpus work and have a +clear linear/bounded complexity profile. + +The target includes: + +- snapshot or serving-index loading; +- classic or map-nav lexical scoring; +- ranking; +- selected-result hydration; +- citation and asset-reference assembly. + +It excludes planner, harvest, control, and answer-generation model time. Those +stages must continue to have separate timings. + +The plan does not change prompts, models, tokenization, BM25 formulas, RRF +weights, cache semantics, citation rules, or public HTTP response shapes. +It does not add LLM-response caching, process-wide serving caches, or startup +prewarming. The acceptance benchmark is a cold, uncached retrieval request, +and planner/harvest/control model time remains a separately reported +dependency. Episode-local reuse is allowed only while one request is active. + +## 2. Verified current behavior + +The current code has two retrieval routes in +`shared/services/retrieval/execution/routes.py`: + +- `use_agentic=false` runs `bottom_discovery()` and then ranking/hydration. +- The default map-nav route calls `load_nav_snapshot(..., lazy=True)`, runs the + synchronous navigation episode, then opens a fresh database context for + reference resolution and result assembly. + +`load_nav_snapshot()` currently: + +1. Reads active documents and their current job-result IDs. +2. Loads all matching sections into memory. +3. Loads all chunk identities and `connect_to` metadata into memory. +4. Uses a lazy store for selected chunk content and asset paths. + +The current persisted map scorer in +`shared/services/retrieval/nav/nav_knowhere.py` still loads all eligible map +units, then reads query frequencies and term scores from the persisted tables. +The scorer itself is fast; the broad database projection is not. + +`bottom_discovery()` currently executes path, content, and term channels +sequentially. `term_channel()` uses substring predicates over lowercased text, +without a trigram index. + +Publication currently writes sections/chunks and then calls +`replace_document_map_units()` in the same SQLAlchemy transaction. The existing +`document_map_unit_indexes` row is a per-document-revision completeness marker. +The existing backfill script only rebuilds that map-unit index; it does not yet +build the proposed serving manifest or namespace statistics. + +The working tree already contains an uncommitted keyset-pagination change and +the `3f4a5b6c7d8e` section-order migration. Keep those changes separate from +the serving-index implementation and from unrelated documentation edits. + +## 3. Architecture decision + +Use PostgreSQL as the source of truth and add a persistent, revision-pinned +serving read model. Do not add OpenSearch, Elasticsearch, Tantivy, or another +search service in this version. PostgreSQL's existing FTS plus `pg_trgm` keeps +the current scoring and tie-breaking behavior directly testable. + +### 3.1 Revision serving manifest + +Add one compressed manifest row per `(document_id, job_result_id)`. The payload +contains ordered metadata only: + +- document, revision, source filename, and job identity; +- section IDs, parent IDs, paths, titles, levels, summaries, and sort order; +- chunk IDs, section IDs, types, sort order, and `connect_to` target IDs; +- map-unit row IDs, unit IDs, unit kinds, token lengths, and sort order; +- root-asset IDs and remounted asset owners. + +It must not contain full chunk content or asset file paths. Those remain in the +canonical tables and are loaded lazily for selected evidence. + +Store the payload as canonical JSON compressed with the standard-library zlib +implementation, with a format version and checksum. The serving loader must +reject an unknown version, checksum mismatch, or incomplete payload. + +### 3.2 Serving generations and statistics + +Add namespace-scoped generation metadata and persistent scoring statistics: + +- `retrieval_namespace_generations`: current generation per user/namespace; +- `retrieval_serving_revision_stats`: compressed per-revision contributions; +- `retrieval_namespace_stats`: aggregate unit counts, total lengths, vocabulary + frequency histograms, and generation; +- `retrieval_namespace_token_stats`: queryable document frequency per channel + and token hash. + +The generation is a consistency marker, not a replacement for revision IDs. +Retrieval captures active revision IDs and one generation. If generation changes +while the snapshot is being captured or before scoring starts, retry once; if +consistency still cannot be proven, use the exact legacy reader. + +Every retrieval route must carry that capture as an immutable revision pin set: +`{document_id, namespace, job_result_id}` plus the captured generation. The pin +set is the source of truth for the request after capture. Downstream queries +must constrain sections, map units, chunks, connected assets, ranking lookups, +reference resolution, and result assembly by the pinned `job_result_id`; they +must not re-join through the live `Document.current_job_result_id`. A generation +change after scoring has started must never cause a mix of old and new rows: +finish against the captured pins (or return an exact legacy result), and only +retry before work that depends on the snapshot begins. +Snapshot admission is therefore decided at capture time. If a later archive +must suppress an in-flight result, discard/retry the whole request; do not +replace its pinned revision with the document's new current revision. + +Cache hits occur before route execution, so every operation that changes the +serving generation (publication, republish, archive, or namespace move) must +also advance the namespace retrieval-cache version, or store the generation in +the cache entry and reject mismatches. This keeps cached responses from +outliving the generation they represent without changing the public response +shape. + +### 3.3 Indexes + +Additive migrations should provide: + +- a token-first covering index for map-unit token candidates. The existing + `idx_document_map_unit_tokens_lookup` is token-first but does not cover the + selected columns; the pending `2e3f4a5b6c7d` migration adds a unit-first + covering index for a different access pattern, so the serving reader may + need one additional token-first covering index; +- a revision/section lookup index for map units; +- a trigram GIN index on `document_map_units.term_search_text_lower`; +- a generated lowercased term field and trigram GIN index for + `document_chunks.term_search_text`; +- the existing chunk ordering index plus the pending token-covering and + section-order migrations (`2e3f4a5b6c7d` and `3f4a5b6c7d8e`). + +Enable PostgreSQL's built-in `pg_trgm` extension. No separate search service is +required. + +## 4. Retrieval changes + +Capture the revision pin set and generation at retrieval-route entry, before +the small-corpus count or route selection. Pass that capture into whichever +route is selected; a route-local capture is allowed only when it is performed +as the same snapshot transaction. This prevents the count/load pair in the +small-corpus optimization from straddling a publication. + +### 4.1 Fast map-nav snapshot + +Extend `load_nav_snapshot()` to try the serving manifest first: + +1. Capture active documents, current revision IDs, and namespace generation in a + short read-only transaction, returning the immutable revision pin set with + the snapshot. +2. Fetch one manifest row per active revision. +3. Decode and validate manifests. +4. Apply the existing document and section exclusion predicates. +5. Build the current `LazyKnowhereProvider` and `LazyChunkRefIndex` from the + decoded metadata. +6. Pin the lazy chunk store to the captured revision IDs. +7. Verify generation stability before returning the snapshot. + +For an unfiltered map-nav request, route selection may count chunks from these +same validated manifests instead of scanning `document_chunks`; filtered and +classic requests retain the exact SQL counter. The count shortcut must fall +back when any manifest is missing or invalid. + +If any manifest is absent or invalid, use the existing legacy snapshot loader. +The legacy loader must return the same revision pin set and apply the same +downstream predicates. This fallback is automatic and exact; it is not a public +feature flag. + +The serving path must preserve current ordering, duplicate bare/document-scoped +reference keys, root-asset remounting, section filtering, and revision pinning. +Keep the pin set available through the complete map-nav request. After the LLM +episode, either materialize selected rows (including connected assets) from the +pinned lazy store before closing it, or pass the pin set to +`resolve_workflow_references()` and `assemble_retrieval_results()`. Their SQL +must select the captured `(document_id, job_result_id)` rows directly, so a +republish during the episode cannot make final citations resolve against the +new current revision. + +### 4.2 Exact persisted map scoring + +Keep `PersistedScoreCorpus` and the existing scorer unchanged wherever possible. +Replace only the data-loading strategy: + +- Prepare the immutable, revision-pinned unit projection and namespace scoring + statistics once per retrieval episode. Checklist relight waves must reuse that + projection; they may fetch or compute only query-specific postings/scores. + A wave must not issue another full-namespace unit/statistics load for the same + pin set. Instrument the loader call count and include it in the benchmark + report so repeated projection loads cannot hide behind separate wave timings. +- use manifest map-unit metadata to represent all units, including zero-score + units; the serving reader should not re-query `document_map_units` for these + IDs, lengths, or section membership; +- query token postings only for tokens in the request; +- filter postings by captured revisions and allowed sections; +- discover term-channel candidates through the trigram index while retaining the + current exact substring/token-hit scoring. The candidate predicate must use + the trigram-indexed `LIKE '%term%'` form (with the same lowercased query and + tokens), then apply the existing exact score expression; do not scan every + unit's term text in Python. +- obtain normal-corpus lengths, document frequencies, and IDF-flooring data from + persistent statistics; +- preserve the existing lexical sort key and RRF ranking. + +Queries with document or section exclusions must remain exact. If adjusted +statistics cannot be calculated with certainty, use the legacy scorer for that +request rather than approximating them. + +### 4.3 Classic retrieval — one pinned revision snapshot + +Keep the existing channel implementations and result projection. In +`bottom_discovery()`: + +- capture one revision pin set and generation before starting any channel; +- execute enabled channels concurrently; +- give each channel its own short-lived database session; +- pass the same pin set to every channel and constrain every channel query to + those revisions; +- preserve channel limits, Python BM25, term scoring, RRF merge, score + normalization, and all-or-error behavior; +- use the new trigram index only to narrow term candidates. + +Do not share one `AsyncSession` across concurrent channel tasks. +Ranking lookups, duplicate suppression, connected-target hydration, and final +assembly must receive the same pin set as discovery. The classic result must +therefore contain rows from one revision per document even if publication +replaces a document while one of the channel sessions is running. The +small-corpus optimization must use this same captured snapshot/pin contract (or +the exact legacy equivalent), rather than loading all rows through live current +revision joins. + +## 5. Publication and lifecycle behavior + +Refactor publication so the same build pass produces: + +- canonical sections/chunks; +- existing map-unit rows and completeness marker; +- the revision serving manifest; +- revision statistics and namespace-statistics deltas. + +All of this happens synchronously in the existing publication transaction. The +completeness marker and generation update are written last. If serving-index +construction fails, the publication transaction rolls back. + +New publication remains online during backfill. First publication, republish, +archive, and namespace-move paths must update statistics under the same +namespace generation row lock. The lock covers the active revision set, +namespace membership, revision contributions, and the generation increment, so +readers and writers have one lifecycle ordering. + +Backfill must rebuild the complete derived serving state (map units, manifest, +revision contribution, and namespace-statistics delta), not only the existing +map-unit index. It must select only documents with `status = 'active'`, a non-null +`current_job_result_id`, and the intended user/namespace. Immediately before +writing a contribution, it must hold the namespace lock and re-read the +document, then require all of the following to remain true: active status, +unchanged user/namespace, and `current_job_result_id` equal to the captured +revision. Otherwise it skips that revision without adding statistics. This +active-status predicate is required in the selector as well as in the +commit-time guard; update `apps/api/scripts/backfill_map_unit_indexes.py` to +include it in the existing selector. + +Archiving must atomically remove or invalidate that document revision's serving +statistics contribution while holding the same lock and advance the namespace +generation. `archive` currently changes `status` without clearing +`current_job_result_id`, so checking the revision pointer alone is insufficient +and would allow an in-flight backfill to re-add an archived revision. + +## 6. Online rollout + +There is no planned downtime, runtime feature flag, or production shadow-read +mode. + +1. Deploy additive schema/index migrations, beginning with the pending + `2e3f4a5b6c7d` and `3f4a5b6c7d8e` migrations. +2. Deploy code that automatically uses the serving reader only for complete, + valid revisions and otherwise uses the legacy reader. +3. Run an explicit, idempotent, bounded backfill for existing active revisions. +4. Keep retrieval and publication online while backfill runs. +5. Verify manifest checksums, revision coverage, namespace statistics, and + generation consistency. +6. Run strict legacy-versus-serving differential checks before considering the + rollout complete. + +If backfill is incomplete, affected revisions continue on the exact legacy +path. If online serving data is corrupted, reject it, alert, repair it with the +backfill/rebuild script, and do not serve partial data. + +Before enabling the serving reader for a namespace, record an inventory of +active `(document_id, current_job_result_id)` pairs, manifest completeness, and +expected per-revision and aggregate unit counts. After backfill, reconcile those +same values and verify that every aggregate includes only active, namespace- +member revisions. Abort the fast-path rollout on any missing/extra revision, +checksum failure, count mismatch, archived contribution, or generation +discontinuity. + +Any migration, backfill, or other database write—especially against +production—requires explicit approval immediately before execution. Read-only +inspection and benchmarking may proceed without that approval. + +## 6.1 DevOps operations runbook + +DevOps owns the production rollout mechanics; application code does not run a +startup backfill or create serving tables implicitly. Execute the following in +order: + +1. **Preflight (read-only):** confirm the target account, database, migration + head, available disk, connection headroom, and a recent rollback point. Record + the active `(document_id, current_job_result_id)` inventory for each namespace + that will be backfilled. +2. **Schema rollout:** with explicit approval immediately beforehand, apply the + additive migrations in dependency order: `2e3f4a5b6c7d`, + `3f4a5b6c7d8e`, `4a5b6c7d8e9f`, then `5b6c7d8e9f0a`. Run the trigram-index + migration during a low-traffic window and monitor for blocking locks. +3. **Application rollout:** deploy the API and worker versions containing the + serving reader and atomic publication changes. Verify health, error rate, and + legacy fallback before starting the backfill. +4. **Bounded backfill:** with separate approval, run + `uv run python apps/api/scripts/backfill_map_unit_indexes.py --apply` from a + controlled operator environment. Limit concurrency, pause on database + saturation, and resume safely; the operation is idempotent and stale or + inactive revisions must be skipped. +5. **Reconciliation:** compare the preflight inventory with serving manifests, + checksums, per-revision unit counts, namespace aggregates, and generation + values. Confirm aggregates contain only active documents still belonging to + the namespace. Investigate every missing, extra, stale, or invalid revision. +6. **Acceptance:** run the production read-only legacy-versus-serving + differential harness and record latency, selected IDs, order, scores, + citations, section paths, asset references, and fallback behavior. Declare + the rollout complete only after zero semantic mismatches. + +If migration or backfill must be stopped, leave the serving tables in place and +stop the operator job. The reader will continue using the exact legacy path for +incomplete revisions. Roll back application code first if necessary; do not +drop serving tables or indexes as an emergency rollback action. Repair a failed +revision by rerunning the bounded backfill after the cause is understood. + +## 7. Contract tests and benchmarks + +Use contract tests only. Add contracts for: + +- manifest round-trip, checksum, version, and revision pinning; +- eager versus serving snapshot equivalence; +- exclusions, duplicate chunk IDs, document-scoped references, and root assets; +- exact Latin/CJK, empty, no-hit, phrase, token-only, and negative-IDF cases; +- incomplete serving data falling back to legacy; +- publication replacement, archive deltas, concurrent generation changes, and + stale backfill protection; +- cache invalidation racing with a generation change, proving an old cached + response is not returned for a newer serving generation; +- map-nav republish during the LLM episode, proving final hydration and + connected-asset resolution stay on the captured revisions; +- classic publication replacement during concurrent channels, ranking, and + final assembly, proving every returned row shares the channel's pin set; +- archive/backfill races proving archived or namespace-moved revisions never + contribute to serving statistics; +- concurrent classic channels preserving IDs, order, scores, citations, and + fallback behavior. + +Race tests must use barriers or an equivalent deterministic hook to force a +republish during the map-nav episode, a publication between classic channel +sessions, an archive during backfill, and a namespace move during backfill. +Each test must assert both the returned evidence and the persisted statistics, +not merely that the request completed. + +The validation harness must also inspect the generated SQL/query plans (or an +equivalent query-boundary assertion) to prove pinned reads do not use live +`Document.current_job_result_id` joins. Run cache-version/generation races and +verify an old cached response is rejected after a lifecycle change. + +Run a differential harness against the production read-only database and +compare selected IDs, ordering, rounded scores, citations, section paths, +asset references, and fallback behavior. + +Benchmark fresh processes and uncached queries. Report separately: + +- serving capture/decode; +- map index projection and scoring; +- episode-local projection reuse (number of full projection loads and per-wave + query-only scoring time); +- classic discovery; +- ranking; +- hydration/assembly; +- total Retrieval Non-LLM Work; +- planner/harvest/control LLM time. + +The complexity check is explicit: one request may perform one full pinned +snapshot/projection pass, relight work should scale with query postings rather +than reloading the corpus, and hydration should scale with selected evidence +(`top_k`/references), not namespace size. Navigation wave count must not +multiply full-corpus database loads. + +Record peak resident memory for a fresh worker during the same benchmark and +repeat it with the expected concurrent-request level. Memory is reported as an +operational trade-off rather than a latency acceptance gate for this phase; +before production rollout, any episode-local or process-local reuse still needs +an explicit byte/item budget and an agreed worker ceiling. + +The fast path is accepted only after zero semantic mismatches and evidence that +the cold request performs one bounded serving projection, does not repeat +full-corpus loads per navigation wave, and meets an agreed latency budget for +the current production-sized corpus. + +## 8. Main tradeoffs and risks + +- Publication becomes slower and uses more storage because derived data is built + synchronously. +- Existing documents need an explicit backfill before they use the fast path. +- A serving-index inconsistency causes a slower legacy request, not approximate + evidence. +- PostgreSQL remains a scaling dependency; a future search-engine migration + would require a new semantic-parity review. diff --git a/docs/design/retrieval-streaming-sse.md b/docs/design/retrieval-streaming-sse.md new file mode 100644 index 00000000..b0b5a73e --- /dev/null +++ b/docs/design/retrieval-streaming-sse.md @@ -0,0 +1,207 @@ +# Retrieval Streaming over SSE + +**Status:** Accepted design +**Related issue:** [#330](https://github.com/Ontos-AI/knowhere/issues/330) +**Related ADR:** [0005](../adr/0005-stream-retrieval-progress-over-sse.md) + +## Purpose + +Online Brain users currently wait for a complete retrieval response while +map-nav planning, searching, source review, and final hydration run. This +design makes that work visible without exposing chain-of-thought or changing +who owns answer generation. + +## Ownership boundary + +Knowhere owns retrieval, safe progress, evidence, and authoritative citations. +The downstream Online Brain client owns answer synthesis and answer-token +streaming. Knowhere does not begin generating the final answer as part of this +feature. + +## Public API + +Add `POST /v2/retrieval/query/stream` with `Content-Type: text/event-stream`. +The request body is the full existing +`RetrievalQueryRequest`; streaming changes delivery, not retrieval semantics. +The existing JSON endpoints remain unchanged as the fallback. + +The stream is live-only. A disconnect cooperatively cancels the run; retrying +starts a new run. Event IDs are monotonic per connection but do not imply +replay or resumability. + +The route remains behind the normal authenticated-user dependency and v2 +route-admission policy. The guest-key allowlist, system-limit configuration, +OpenAPI registration, and any CORS policy must explicitly include the stream +route. v2 BYOK `llm_config` is accepted exactly as on the JSON route and is +never copied into an event payload or log message. + +## Event contract + +Each SSE frame has an event name (`progress`, `heartbeat`, or `terminal`) and a +JSON payload with `schema_version`, `stream_id`, `sequence`, and server +`elapsed_ms`. Progress payloads use the fixed, route-aware phases: + +```text +started → planning? → searching → reviewing_sources → finalizing +``` + +Classic and small-corpus routes omit `planning`; they must not emit phases that +did not occur. In-progress events contain only safe aggregate counts such as +`candidate_source_count` and `reviewed_source_count`. They do not contain +document names, chunk content, query rewrites, raw planner output, citations, +or chain-of-thought. + +The terminal event is a versioned envelope containing the existing retrieval +response as the authoritative result: + +```json +{ + "schema_version": 1, + "stream_id": "rst_...", + "sequence": 7, + "elapsed_ms": 2410, + "status": "completed", + "response": { "namespace": "default", "query": "...", "router_used": "mapnav", "evidence_text": "...", "referenced_chunks": [], "results": [] } +} +``` + +Failure terminals use `failed`, `cancelled`, or `no_results`. Failures expose a +stable user-safe `code` and `message`; detailed provider, database, +authentication, and planner errors remain in server logs. + +Wire requirements are part of the contract: frames use UTF-8 SSE `id`, +`event`, and one `data` line followed by a blank line; a heartbeat is an SSE +comment or named heartbeat event and carries no retrieval data; exactly one +terminal event is sent before the connection closes; no `retry` directive is +promised because the stream is not resumable. Progress events may be +coalesced when a bounded queue is full, but terminal events must never be +dropped. + +Cache hits still emit `started` and a terminal `completed` event, with a safe +`cache_hit: true` indicator. They omit phases for work that did not run. +`no_results` is reserved for a successful retrieval execution that produces no +evidence; provider, timeout, validation, and internal failures remain +`failed`. + +HTTP authentication, request validation, and route-admission failures happen +before the stream starts and use ordinary HTTP error responses. Once a `200` +SSE response has been opened, execution failures must be represented by a +terminal SSE event because the HTTP status can no longer be changed. + +## Internal implementation seam + +Keep the synchronous map-nav implementation. Add an optional callback that +receives a sanitized progress projection after each completed planner, +search, or review step. The SSE route bridges this callback to an +`asyncio.Queue` using a thread-safe loop handoff while retrieval continues in +its existing worker thread. + +The queue is bounded and owns cleanup of the worker task, callback, heartbeat +task, and database session. The callback must not touch an `AsyncSession` from +the worker thread. The route polls request disconnect state and propagates a +cancellation token; all producer tasks are joined or cancelled in a `finally` +block so abandoned streams cannot leak threads or connections. + +Add cooperative cancellation checks between steps. An in-flight synchronous +provider call may finish before cancellation takes effect. Cancelled runs do +not perform final hydration when cancellation is observed in time. + +Phase ownership is explicit: the route emits `started`; the map-nav adapter +emits `planning` before `plan_query` and `searching` before navigation or +classic discovery; the route emits `reviewing_sources` after retrieval +selection and before reference hydration; and it emits `finalizing` before +public projection. Counts are sourced from existing snapshot, reference, and +assembled-result counts and are omitted when not yet known. + +## Correct duration accounting + +The execution plan already starts a monotonic timer before cache lookup and +logs elapsed time after the route. However, `TraceRecorder` currently starts a +second timer in its constructor, and the map-nav route constructs it only +after navigation, reference resolution, and result assembly. Persisted +`retrieval_runs.latency_ms` and its aggregates therefore under-report retrieval +latency and mostly measure trace flush time. + +The canonical `Retrieval Duration` is the execution-plan timer from retrieval +start through final public-result assembly. It includes cache lookup and +applies to cache hits and misses. It excludes authentication, network/SSE +delivery, and downstream answer generation. + +Required changes: + +- pass the execution start timestamp into `TraceRecorder`; +- set `retrieval_runs.latency_ms` from that timestamp; +- record cache-hit runs with the same definition; +- ensure classic, map-nav, small-corpus, cache-hit, failed, and cancelled + retrievals all have an explicit timing/observability outcome; +- expose separate `time_to_first_event_ms`, `retrieval_latency_ms`, and + downstream `time_to_first_token_ms` measurements; +- retain per-step `elapsed_ms` as step latency, not total request latency. + +`retrieval_runs` is the ledger for every retrieval execution, not only +map-nav. Each row records the route type, `agentic_enabled`, `cache_hit`, +canonical latency, and terminal status for classic, map-nav, small-corpus, +cache-hit, failed, and cancelled runs. Add a backward-compatible status field +and migration rather than overloading free-form error text. + +The execution timer must end after public response projection, not merely when +the internal route outcome is assembled. If trace persistence is best-effort, +its latency update must still use the captured execution timestamps and must +not extend the user-visible retrieval duration with an unbounded database +flush. + +Because the map-nav route deliberately rolls back its request session before +the synchronous LLM episode, a trace row must not be created in that session +before the rollback. Capture the execution start immediately, then create or +update the trace record at the terminal persistence point with the captured +start and an explicit end/duration supplied by the execution plan (or use a +separate trace session). `TraceRecorder.complete()` must not silently choose a +later local constructor time or include its own flush duration. + +## Operational requirements + +- heartbeat every 15 seconds while active; +- `Cache-Control: no-cache` and `X-Accel-Buffering: no`; +- flush after every event; +- use `fetch`-style clients where POST authorization headers are required; +- propagate disconnects to the cancellation token. +- bound maximum stream lifetime and enforce the same request/rate-limit policy + as the JSON endpoint; +- document worker-thread, database-connection, and concurrent-stream limits. +- count one stream request as one retrieval request under the existing user and + system limits; retries count as new requests and cannot bypass quota. + +## Delivery slices + +1. **Knowhere contract and timing:** typed event models, stream route, callback + bridge, cancellation, route admission, cache semantics, corrected + `RetrievalRun` timing across all route types, and API contract tests. +2. **SDK adapters:** typed Python and Node stream consumers with fallback to + the existing JSON query; parse named SSE events, expose abort/error + handling, and preserve the full terminal response shape. +3. **Online Brain UX:** phase state model, progress view, safe error states, + and integration with existing downstream answer-token streaming. +4. **Production verification:** proxy-path e2e tests and latency/buffering + instrumentation before setting hard p50/p95 targets. + +## Verification gates + +- correct phase order for map-nav, classic, and small-corpus routes; +- cache-hit streams emit only applicable phases and identify the cache hit; +- no sensitive planner or evidence data before the terminal event; +- terminal citations match the existing JSON endpoint; +- cancellation, timeout, no-results, provider failure, and disconnect are + distinguishable; +- persisted latency includes the full retrieval path for every supported route + and matches API timing within an expected tolerance; +- cache-hit latency and route type are visible in observability data; +- heartbeats pass through the deployed proxy without buffering; +- first-event and first-token times are measured separately. + +## Explicit non-goals + +- v1 API streaming endpoint; +- resumable/replayed streams; +- partial authoritative evidence or provisional citation revision; +- moving answer generation into Knowhere; +- invented latency SLAs before a production baseline exists. diff --git a/packages/shared-python/shared/models/database/document.py b/packages/shared-python/shared/models/database/document.py index ad59857e..c396283d 100644 --- a/packages/shared-python/shared/models/database/document.py +++ b/packages/shared-python/shared/models/database/document.py @@ -12,8 +12,10 @@ DateTime, Float, ForeignKey, + BigInteger, Index, Integer, + LargeBinary, String, Text, UniqueConstraint, @@ -130,6 +132,13 @@ class DocumentSection(Base): ), Index("idx_document_sections_scope", "user_id", "namespace"), Index("idx_document_sections_doc_revision", "document_id", "job_result_id"), + Index( + "idx_document_sections_revision_snapshot_order", + "document_id", + "job_result_id", + "sort_order", + "section_id", + ), ) @@ -347,6 +356,168 @@ class DocumentMapUnitIndex(Base): ) +class RetrievalNamespaceGeneration(Base): + """Monotonic serving generation for one user-owned namespace.""" + + __tablename__ = "retrieval_namespace_generations" + + id: Mapped[str] = mapped_column( + String(100), primary_key=True, default=lambda: f"rng_{uuid4().hex}" + ) + user_id: Mapped[str] = mapped_column(Text, nullable=False) + namespace: Mapped[str] = mapped_column(String(255), nullable=False) + generation: Mapped[int] = mapped_column(BigInteger, nullable=False, default=0) + updated_at: Mapped[datetime] = mapped_column( + DateTime, default=utc_now_naive, onupdate=utc_now_naive, nullable=False + ) + + __table_args__ = ( + UniqueConstraint( + "user_id", + "namespace", + name="uq_retrieval_namespace_generations_scope", + ), + ) + + +class RetrievalServingRevisionManifest(Base): + """Compressed ordered metadata for one document revision.""" + + __tablename__ = "retrieval_serving_revision_manifests" + + id: Mapped[str] = mapped_column( + String(100), primary_key=True, default=lambda: f"rsm_{uuid4().hex}" + ) + user_id: Mapped[str] = mapped_column(Text, nullable=False) + namespace: Mapped[str] = mapped_column(String(255), nullable=False) + 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) + payload_zlib: Mapped[bytes] = mapped_column(LargeBinary, nullable=False) + checksum: Mapped[str] = mapped_column(String(64), 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_retrieval_serving_revision_manifests_revision", + ), + Index( + "idx_retrieval_serving_revision_manifests_scope", + "user_id", + "namespace", + "document_id", + "job_result_id", + ), + ) + + +class RetrievalServingRevisionStat(Base): + """Compressed scoring contribution for one document revision.""" + + __tablename__ = "retrieval_serving_revision_stats" + + id: Mapped[str] = mapped_column( + String(100), primary_key=True, default=lambda: f"rss_{uuid4().hex}" + ) + user_id: Mapped[str] = mapped_column(Text, nullable=False) + namespace: Mapped[str] = mapped_column(String(255), nullable=False) + 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) + payload_zlib: Mapped[bytes] = mapped_column(LargeBinary, nullable=False) + checksum: Mapped[str] = mapped_column(String(64), 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_retrieval_serving_revision_stats_revision", + ), + Index( + "idx_retrieval_serving_revision_stats_scope", + "user_id", + "namespace", + "document_id", + "job_result_id", + ), + ) + + +class RetrievalNamespaceStat(Base): + """Compressed aggregate scoring statistics for one namespace generation.""" + + __tablename__ = "retrieval_namespace_stats" + + id: Mapped[str] = mapped_column( + String(100), primary_key=True, default=lambda: f"rns_{uuid4().hex}" + ) + user_id: Mapped[str] = mapped_column(Text, nullable=False) + namespace: Mapped[str] = mapped_column(String(255), nullable=False) + generation: Mapped[int] = mapped_column(BigInteger, nullable=False) + payload_zlib: Mapped[bytes] = mapped_column(LargeBinary, nullable=False) + checksum: Mapped[str] = mapped_column(String(64), nullable=False) + updated_at: Mapped[datetime] = mapped_column( + DateTime, default=utc_now_naive, onupdate=utc_now_naive, nullable=False + ) + + __table_args__ = ( + UniqueConstraint( + "user_id", + "namespace", + name="uq_retrieval_namespace_stats_scope", + ), + ) + + +class RetrievalNamespaceTokenStat(Base): + """Document frequency for one token/channel in a namespace generation.""" + + __tablename__ = "retrieval_namespace_token_stats" + + id: Mapped[str] = mapped_column( + String(100), primary_key=True, default=lambda: f"rnt_{uuid4().hex}" + ) + user_id: Mapped[str] = mapped_column(Text, nullable=False) + namespace: Mapped[str] = mapped_column(String(255), nullable=False) + generation: Mapped[int] = mapped_column(BigInteger, nullable=False) + channel: Mapped[str] = mapped_column(String(32), nullable=False) + token_hash: Mapped[str] = mapped_column(String(64), nullable=False) + document_frequency: Mapped[int] = mapped_column(Integer, nullable=False) + + __table_args__ = ( + UniqueConstraint( + "user_id", + "namespace", + "channel", + "token_hash", + name="uq_retrieval_namespace_token_stats_key", + ), + Index( + "idx_retrieval_namespace_token_stats_lookup", + "user_id", + "namespace", + "generation", + "channel", + "token_hash", + ), + ) + + class GraphNode(Base): """Persisted derived graph node used for routing and expansion.""" diff --git a/packages/shared-python/shared/services/retrieval/execution/plan.py b/packages/shared-python/shared/services/retrieval/execution/plan.py index 8d05753a..adb2a6da 100644 --- a/packages/shared-python/shared/services/retrieval/execution/plan.py +++ b/packages/shared-python/shared/services/retrieval/execution/plan.py @@ -1,6 +1,7 @@ from __future__ import annotations import time +from dataclasses import replace from typing import Any from loguru import logger @@ -16,6 +17,10 @@ set_cached_retrieval_query_result, ) from shared.services.retrieval.execution.routes import run_retrieval_route +from shared.services.retrieval.execution.revision_pins import ( + capture_revision_pins, + is_revision_generation_stable, +) from shared.services.retrieval.stats.recorder import ( schedule_retrieval_hit_stats_update, ) @@ -142,7 +147,26 @@ async def _execute_with_overrides(self, request: RetrievalQuery) -> dict[str, An logger.debug(f" 📦 Cache miss (version={cache_version}), running full pipeline") - outcome = await run_retrieval_route(request.build_route_context()) + route_context = request.build_route_context() + revision_pins = await capture_revision_pins( + request.db, + user_id=request.user_id, + namespace=request.namespace, + ) + if not await is_revision_generation_stable( + request.db, + user_id=request.user_id, + namespace=request.namespace, + pins=revision_pins, + ): + revision_pins = await capture_revision_pins( + request.db, + user_id=request.user_id, + namespace=request.namespace, + ) + outcome = await run_retrieval_route( + replace(route_context, revision_pins=revision_pins) + ) if cache_version is not None: await _write_cached_response( diff --git a/packages/shared-python/shared/services/retrieval/execution/reference_resolver.py b/packages/shared-python/shared/services/retrieval/execution/reference_resolver.py index 802fa278..459ad273 100644 --- a/packages/shared-python/shared/services/retrieval/execution/reference_resolver.py +++ b/packages/shared-python/shared/services/retrieval/execution/reference_resolver.py @@ -1,6 +1,7 @@ from __future__ import annotations from dataclasses import dataclass +from collections.abc import Mapping from typing import Any from sqlalchemy.ext.asyncio import AsyncSession @@ -28,6 +29,7 @@ async def resolve_workflow_references( namespace: str, refs: list[dict[str, Any]], score_by_chunk_id: dict[str, float] | None = None, + revision_pins: Mapping[str, str] | None = None, ) -> ResolvedWorkflowReferences: hydrated_rows = await hydrate_referenced_chunk_rows( db=db, @@ -35,6 +37,7 @@ async def resolve_workflow_references( namespace=namespace, refs=refs, score_by_chunk_id=score_by_chunk_id, + revision_pins=revision_pins, ) resolved = _select_matching_references(refs, hydrated_rows) enriched_rows = await enrich_referenced_chunks_with_asset_url(resolved.rows) diff --git a/packages/shared-python/shared/services/retrieval/execution/revision_pins.py b/packages/shared-python/shared/services/retrieval/execution/revision_pins.py new file mode 100644 index 00000000..e7e9efcc --- /dev/null +++ b/packages/shared-python/shared/services/retrieval/execution/revision_pins.py @@ -0,0 +1,92 @@ +"""Capture and carry one immutable revision set through a retrieval request.""" + +from __future__ import annotations + +from collections.abc import Iterator, Mapping +from dataclasses import dataclass +from types import MappingProxyType + +from sqlalchemy import select +from sqlalchemy.exc import SQLAlchemyError +from sqlalchemy.ext.asyncio import AsyncSession + +from shared.models.database.document import Document +from shared.models.database.document import RetrievalNamespaceGeneration + + +@dataclass(frozen=True) +class RetrievalRevisionPins(Mapping[str, str]): + """The active document revisions admitted to one retrieval request.""" + + revisions: Mapping[str, str] + generation: int | None = None + + def __post_init__(self) -> None: + object.__setattr__(self, "revisions", MappingProxyType(dict(self.revisions))) + + def __getitem__(self, document_id: str) -> str: + return self.revisions[document_id] + + def __iter__(self) -> Iterator[str]: + return iter(self.revisions) + + def __len__(self) -> int: + return len(self.revisions) + + +async def capture_revision_pins( + db: AsyncSession, + *, + user_id: str, + namespace: str, +) -> RetrievalRevisionPins: + """Capture active document revisions in one database read transaction.""" + statement = ( + select(Document.document_id, Document.current_job_result_id) + .where(Document.user_id == user_id) + .where(Document.namespace == namespace) + .where(Document.status == "active") + .where(Document.current_job_result_id.is_not(None)) + .order_by(Document.document_id) + ) + try: + generation_result = await db.execute( + select(RetrievalNamespaceGeneration.generation) + .where(RetrievalNamespaceGeneration.user_id == user_id) + .where(RetrievalNamespaceGeneration.namespace == namespace) + ) + generation_row = generation_result.scalar_one_or_none() + except SQLAlchemyError: + await db.rollback() + generation_row = None + rows = (await db.execute(statement)).all() + revisions = { + str(document_id): str(job_result_id) + for document_id, job_result_id in rows + if document_id and job_result_id + } + return RetrievalRevisionPins( + revisions=revisions, + generation=int(generation_row) if generation_row is not None else 0, + ) + + +async def is_revision_generation_stable( + db: AsyncSession, + *, + user_id: str, + namespace: str, + pins: RetrievalRevisionPins, +) -> bool: + """Return whether the namespace generation is unchanged since capture.""" + try: + result = await db.execute( + select(RetrievalNamespaceGeneration.generation) + .where(RetrievalNamespaceGeneration.user_id == user_id) + .where(RetrievalNamespaceGeneration.namespace == namespace) + ) + current_generation = result.scalar_one_or_none() + except SQLAlchemyError: + await db.rollback() + return True + return int(current_generation or 0) == int(pins.generation or 0) diff --git a/packages/shared-python/shared/services/retrieval/execution/route_types.py b/packages/shared-python/shared/services/retrieval/execution/route_types.py index 59c173e4..13c21c46 100644 --- a/packages/shared-python/shared/services/retrieval/execution/route_types.py +++ b/packages/shared-python/shared/services/retrieval/execution/route_types.py @@ -5,6 +5,8 @@ from sqlalchemy.ext.asyncio import AsyncSession +from shared.services.retrieval.execution.revision_pins import RetrievalRevisionPins + @dataclass(frozen=True) class RetrievalRouteContext: @@ -26,6 +28,7 @@ class RetrievalRouteContext: internal_recall_k: int | None effective_recall_k: int use_agentic: bool | None + revision_pins: RetrievalRevisionPins | None = None @dataclass(frozen=True) diff --git a/packages/shared-python/shared/services/retrieval/execution/routes.py b/packages/shared-python/shared/services/retrieval/execution/routes.py index 97961074..d2b6d211 100644 --- a/packages/shared-python/shared/services/retrieval/execution/routes.py +++ b/packages/shared-python/shared/services/retrieval/execution/routes.py @@ -8,18 +8,29 @@ from sqlalchemy.ext.asyncio import AsyncSession from shared.services.retrieval.search.discovery import bottom_discovery -from shared.services.retrieval.execution.reference_resolver import resolve_workflow_references -from shared.services.retrieval.hydration.result_assembly import assemble_retrieval_results -from shared.services.retrieval.hydration.legacy_evidence import render_legacy_evidence_text +from shared.services.retrieval.execution.reference_resolver import ( + resolve_workflow_references, +) +from shared.services.retrieval.hydration.result_assembly import ( + assemble_retrieval_results, +) +from shared.services.retrieval.hydration.legacy_evidence import ( + render_legacy_evidence_text, +) from shared.services.retrieval.execution.route_types import ( RetrievalRouteContext, RetrievalRouteOutcome, ) from shared.services.retrieval.search.ranking import rank_retrieval_candidates from shared.services.retrieval.search.scoped_corpus import ( + count_manifest_chunks, count_scoped_chunks, load_all_scoped_chunks, ) +from shared.services.retrieval.execution.revision_pins import ( + capture_revision_pins, + is_revision_generation_stable, +) def open_fresh_database_context() -> AbstractAsyncContextManager[AsyncSession]: @@ -46,13 +57,28 @@ async def run_retrieval_route( async def _try_run_small_corpus_route( context: RetrievalRouteContext, ) -> RetrievalRouteOutcome | None: - total_chunk_count = await count_scoped_chunks( - context.db, - user_id=context.user_id, - namespace=context.namespace, - exclude_document_ids=context.exclude_document_ids, - allowed_chunk_types=context.allowed_chunk_types, - ) + total_chunk_count: int | None = None + if ( + context.use_agentic is not False + and context.revision_pins is not None + and not context.exclude_document_ids + and not context.exclude_sections + and context.allowed_chunk_types is None + and not context.signal_paths + ): + total_chunk_count = await count_manifest_chunks( + context.db, + revision_pins=context.revision_pins, + ) + if total_chunk_count is None: + total_chunk_count = await count_scoped_chunks( + context.db, + user_id=context.user_id, + namespace=context.namespace, + exclude_document_ids=context.exclude_document_ids, + allowed_chunk_types=context.allowed_chunk_types, + revision_pins=context.revision_pins, + ) logger.info(f"\n Total chunks in scope: {total_chunk_count}") if total_chunk_count > context.top_k: @@ -71,6 +97,7 @@ async def _try_run_small_corpus_route( allowed_chunk_types=context.allowed_chunk_types, signal_paths=context.signal_paths or [], filter_mode=context.filter_mode, + revision_pins=context.revision_pins, ) logger.info( f" small_corpus load: loaded={len(all_rows)} rows after signal/exclude filters" @@ -81,6 +108,7 @@ async def _try_run_small_corpus_route( exclude_document_ids=context.exclude_document_ids, exclude_sections=context.exclude_sections, allowed_chunk_types=context.allowed_chunk_types, + revision_pins=context.revision_pins, ) results = assembled_rows response = { @@ -117,6 +145,7 @@ async def _run_classic_topk_route( channels=context.channels, channel_weights=context.channel_weights, internal_recall_k=context.internal_recall_k, + revision_pins=context.revision_pins, ) fused_rows = ( @@ -132,6 +161,7 @@ async def _run_classic_topk_route( discovery_rows=fused_rows, routed_rows=[], top_k=context.top_k, + revision_pins=context.revision_pins, ) assembled_rows = await assemble_retrieval_results( @@ -140,6 +170,7 @@ async def _run_classic_topk_route( exclude_document_ids=context.exclude_document_ids, exclude_sections=context.exclude_sections, allowed_chunk_types=context.allowed_chunk_types, + revision_pins=context.revision_pins, ) results = assembled_rows response = { @@ -181,7 +212,9 @@ async def _run_mapnav_route( episode_token_count, episode_workflow_plan, ) + snapshot_started = time.perf_counter() + snapshot_pins = context.revision_pins snapshot = await load_nav_snapshot( context.db, user_id=context.user_id, @@ -189,7 +222,29 @@ async def _run_mapnav_route( exclude_document_ids=context.exclude_document_ids, exclude_sections=context.exclude_sections, lazy=True, + revision_pins=snapshot_pins, ) + if snapshot_pins is not None and not await is_revision_generation_stable( + context.db, + user_id=context.user_id, + namespace=context.namespace, + pins=snapshot_pins, + ): + snapshot.close() + snapshot_pins = await capture_revision_pins( + context.db, + user_id=context.user_id, + namespace=context.namespace, + ) + snapshot = await load_nav_snapshot( + context.db, + user_id=context.user_id, + namespace=context.namespace, + exclude_document_ids=context.exclude_document_ids, + exclude_sections=context.exclude_sections, + lazy=True, + revision_pins=snapshot_pins, + ) snapshot_seconds = time.perf_counter() - snapshot_started logger.info( "retrieval mapnav stage=snapshot_load seconds={:.3f} documents={} refs={}".format( @@ -239,6 +294,7 @@ async def _run_mapnav_route( namespace=context.namespace, refs=refs, score_by_chunk_id=score_by_chunk_id or None, + revision_pins=snapshot.document_revisions, ) assembled_rows = await assemble_retrieval_results( db=final_db, @@ -246,6 +302,7 @@ async def _run_mapnav_route( exclude_document_ids=context.exclude_document_ids, exclude_sections=context.exclude_sections, allowed_chunk_types=context.allowed_chunk_types, + revision_pins=snapshot.document_revisions, ) decision_steps = build_decision_trace( @@ -299,9 +356,7 @@ async def _run_mapnav_route( "decision_trace": decision_trace, } - completion_detail = ( - f"chunks | evidence={len(evidence_text)} chars | router=mapnav" - ) + completion_detail = f"chunks | evidence={len(evidence_text)} chars | router=mapnav" return RetrievalRouteOutcome( response=response, hit_stats_results=resolved.refs, diff --git a/packages/shared-python/shared/services/retrieval/hydration/connected.py b/packages/shared-python/shared/services/retrieval/hydration/connected.py index fcbc647b..2a1f3127 100644 --- a/packages/shared-python/shared/services/retrieval/hydration/connected.py +++ b/packages/shared-python/shared/services/retrieval/hydration/connected.py @@ -1,5 +1,6 @@ from __future__ import annotations +from collections.abc import Mapping from typing import Any from sqlalchemy import and_, or_, select @@ -20,6 +21,7 @@ async def hydrate_connected_target_rows( rows: list[dict[str, Any]], exclude_document_ids: list[str], exclude_sections: list[dict[str, str]], + revision_pins: Mapping[str, str] | None = None, ) -> list[dict[str, Any]]: if db is None: return [] @@ -61,7 +63,25 @@ async def hydrate_connected_target_rows( stmt = ( select(Document, DocumentChunk, DocumentSection, JobResult) - .join(DocumentChunk, DocumentChunk.document_id == Document.document_id) + .join( + DocumentChunk, + ( + (DocumentChunk.document_id == Document.document_id) + if revision_pins is None + else and_( + DocumentChunk.document_id == Document.document_id, + or_( + *[ + and_( + DocumentChunk.document_id == document_id, + DocumentChunk.job_result_id == job_result_id, + ) + for document_id, job_result_id in target_ids_by_revision + ] + ), + ) + ), + ) .outerjoin(DocumentSection, DocumentSection.section_id == DocumentChunk.section_id) .join(JobResult, JobResult.id == DocumentChunk.job_result_id) .where(or_(*revision_filters)) diff --git a/packages/shared-python/shared/services/retrieval/hydration/reference.py b/packages/shared-python/shared/services/retrieval/hydration/reference.py index 46208df5..cd240654 100644 --- a/packages/shared-python/shared/services/retrieval/hydration/reference.py +++ b/packages/shared-python/shared/services/retrieval/hydration/reference.py @@ -1,8 +1,9 @@ from __future__ import annotations +from collections.abc import Mapping from typing import Any -from sqlalchemy import select +from sqlalchemy import and_, or_, select from sqlalchemy.ext.asyncio import AsyncSession from shared.models.database.document import Document, DocumentChunk, DocumentSection @@ -20,6 +21,7 @@ async def hydrate_referenced_chunk_rows( namespace: str, refs: list[dict[str, Any]], score_by_chunk_id: dict[str, float] | None = None, + revision_pins: Mapping[str, str] | None = None, ) -> list[dict[str, Any]]: if db is None or not refs: return [] @@ -39,22 +41,48 @@ async def hydrate_referenced_chunk_rows( document_ids = sorted({document_id for document_id, _, _, _ in ref_keys}) chunk_ids = sorted({chunk_id for _, chunk_id, _, _ in ref_keys}) - stmt = ( - select(Document, DocumentChunk, DocumentSection, JobResult) - .join( - DocumentChunk, + pinned_document_ids = [ + document_id for document_id in document_ids if revision_pins and document_id in revision_pins + ] + if revision_pins is not None and not pinned_document_ids: + return [] + + if revision_pins is None: + chunk_join = ( (DocumentChunk.document_id == Document.document_id) - & (DocumentChunk.job_result_id == Document.current_job_result_id), + & (DocumentChunk.job_result_id == Document.current_job_result_id) + ) + else: + chunk_join = and_( + DocumentChunk.document_id == Document.document_id, + or_( + *[ + and_( + DocumentChunk.document_id == document_id, + DocumentChunk.job_result_id == str(revision_pins[document_id]), + ) + for document_id in pinned_document_ids + ] + ), ) + + stmt = ( + select(Document, DocumentChunk, DocumentSection, JobResult) + .join(DocumentChunk, chunk_join) .outerjoin(DocumentSection, DocumentSection.section_id == DocumentChunk.section_id) .join(JobResult, JobResult.id == DocumentChunk.job_result_id) .where(Document.user_id == user_id) .where(Document.namespace == namespace) - .where(Document.status == 'active') - .where(Document.document_id.in_(document_ids)) + .where( + Document.document_id.in_( + document_ids if revision_pins is None else pinned_document_ids + ) + ) .where(DocumentChunk.chunk_id.in_(chunk_ids)) .order_by(DocumentChunk.sort_order) ) + if revision_pins is None: + stmt = stmt.where(Document.status == 'active') result = await db.execute(stmt) rows_by_key: dict[ReferenceLookupKey, dict[str, Any]] = {} diff --git a/packages/shared-python/shared/services/retrieval/hydration/result_assembly.py b/packages/shared-python/shared/services/retrieval/hydration/result_assembly.py index ed448704..863bfadb 100644 --- a/packages/shared-python/shared/services/retrieval/hydration/result_assembly.py +++ b/packages/shared-python/shared/services/retrieval/hydration/result_assembly.py @@ -1,5 +1,6 @@ from __future__ import annotations +from collections.abc import Mapping from typing import Any from sqlalchemy.ext.asyncio import AsyncSession @@ -21,6 +22,7 @@ async def assemble_retrieval_results( exclude_document_ids: list[str], exclude_sections: list[dict[str, str]], allowed_chunk_types: set[str] | None = None, + revision_pins: Mapping[str, str] | None = None, ) -> list[dict[str, Any]]: filtered_rows = filter_excluded_rows( rows, @@ -37,6 +39,7 @@ async def assemble_retrieval_results( rows=filtered_rows, exclude_document_ids=exclude_document_ids, exclude_sections=exclude_sections, + revision_pins=revision_pins, ) rows_by_chunk_id = { str(row.get('chunk_id') or ''): row diff --git a/packages/shared-python/shared/services/retrieval/manifest_cache.py b/packages/shared-python/shared/services/retrieval/manifest_cache.py new file mode 100644 index 00000000..3861903e --- /dev/null +++ b/packages/shared-python/shared/services/retrieval/manifest_cache.py @@ -0,0 +1,47 @@ +"""Request-scoped decoded serving-manifest reuse.""" + +from __future__ import annotations + +from collections.abc import Mapping +from typing import Any + + +_CACHE_INFO_KEY = "retrieval_serving_manifest_payloads" + + +def manifest_revision_key(revisions: Mapping[str, str]) -> tuple[tuple[str, str], ...]: + """Return a stable key for one immutable revision pin set.""" + return tuple(sorted((str(document_id), str(job_result_id)) for document_id, job_result_id in revisions.items())) + + +def get_cached_manifest_payloads( + db: object, + *, + revisions: Mapping[str, str], +) -> dict[tuple[str, str], dict[str, Any]] | None: + """Return decoded manifests cached on this request's SQLAlchemy session.""" + info = getattr(db, "info", None) + if not isinstance(info, dict): + return None + batches = info.get(_CACHE_INFO_KEY) + if not isinstance(batches, dict): + return None + payloads = batches.get(manifest_revision_key(revisions)) + if not isinstance(payloads, dict): + return None + return payloads + + +def cache_manifest_payloads( + db: object, + *, + revisions: Mapping[str, str], + payloads: dict[tuple[str, str], dict[str, Any]], +) -> None: + """Store one complete decoded manifest batch for this request only.""" + info = getattr(db, "info", None) + if not isinstance(info, dict): + return + batches = info.setdefault(_CACHE_INFO_KEY, {}) + if isinstance(batches, dict): + batches[manifest_revision_key(revisions)] = payloads 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 2aad0b29..e25234ea 100644 --- a/packages/shared-python/shared/services/retrieval/nav/nav_knowhere.py +++ b/packages/shared-python/shared/services/retrieval/nav/nav_knowhere.py @@ -47,6 +47,7 @@ PersistedScoreUnit, tokenize_query_for_ranker, ) +from shared.services.retrieval.serving_manifest import decode_serving_manifest _ASSET_TYPES = ("table", "image") # Knowhere sentinel path for the virtual document container (not a collectable leaf). @@ -54,6 +55,11 @@ _DEFAULT_DSN = "postgresql://root:root123@127.0.0.1:5433/Knowhere" _MAP_UNIT_INDEX_FORMAT_VERSION = 1 _MAP_SCORE_CHANNELS: Tuple[str, str] = ("path", "content") +# PostgreSQL's trigram planner evaluates one LIKE branch per pattern. Once a +# query has more than a handful of patterns, fetching the bounded pinned unit +# set once and evaluating literal substring hits in Python is faster and keeps +# SQL planning time from growing with the planner's subgoal length. +_TERM_SCORE_FULL_SCAN_PATTERN_THRESHOLD = 8 _logger = logging.getLogger(__name__) @@ -220,6 +226,19 @@ def __init__( self._revisions = dict(revisions) self._excluded_sections = set(excluded_sections or ()) self._conn: Optional[_SyncConnection] = None + self._score_manifest_cache: Optional[ + tuple[tuple[tuple[str, str], ...], list[Sequence[object]]] + ] = None + self._score_unit_rows_cache: dict[ + tuple[tuple[str, str], ...], list[Sequence[object]] + ] = {} + self._score_frequency_cache: dict[ + tuple[tuple[str, str], ...], + dict[tuple[str, str], dict[str, int]], + ] = {} + self._score_term_cache: dict[ + tuple[tuple[str, str], ...], dict[str, Tuple[float, ...]] + ] = {} def _connection(self) -> "_SyncConnection": if self._conn is None: @@ -317,26 +336,71 @@ def load_persisted_score_corpus( ] cur = self._connection().cursor() try: - stage_started = time.perf_counter() - 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()) + revision_key = tuple(revisions) + cached_manifests = self._score_manifest_cache + if cached_manifests is not None and cached_manifests[0] == revision_key: + manifests = cached_manifests[1] + _logger.info( + "retrieval map-index load stage=manifests cache_hit rows=%d", + len(manifests), + ) + else: + stage_started = time.perf_counter() + try: + cur.execute( + "SELECT indexes.document_id, indexes.job_result_id, " + "indexes.format_version, indexes.unit_count, indexes.token_count, " + "manifests.payload_zlib, manifests.checksum, manifests.format_version, " + "statistics.payload_zlib, statistics.checksum, statistics.format_version " + "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 " + "JOIN retrieval_serving_revision_manifests AS manifests " + "ON manifests.document_id = indexes.document_id " + "AND manifests.job_result_id = indexes.job_result_id " + "JOIN retrieval_serving_revision_stats AS statistics " + "ON statistics.document_id = indexes.document_id " + "AND statistics.job_result_id = indexes.job_result_id", + revision_params, + ) + except Exception as exc: + if getattr(exc, "pgcode", None) == "42P01": + return None + raise + manifests = list(cur.fetchall()) + self._score_manifest_cache = (revision_key, manifests) _logger.info( "retrieval map-index load stage=manifests seconds=%.3f rows=%d", - time.perf_counter() - stage_started, + time.perf_counter() - stage_started if cached_manifests is None else 0.0, len(manifests), ) if len(manifests) != len(revisions) or any( - int(row[2]) != _MAP_UNIT_INDEX_FORMAT_VERSION for row in manifests + len(row) < 11 + or int(row[2]) != _MAP_UNIT_INDEX_FORMAT_VERSION + or not row[5] + or not row[6] + or not row[8] + or not row[9] + for row in manifests ): return None + decoded_manifests: dict[tuple[str, str], dict[str, Any]] = {} + decoded_statistics: dict[tuple[str, str], dict[str, Any]] = {} + try: + for row in manifests: + decoded_manifests[(str(row[0]), str(row[1]))] = decode_serving_manifest( + bytes(row[5]), + checksum=str(row[6]), + format_version=int(row[7]), + ) + decoded_statistics[(str(row[0]), str(row[1]))] = decode_serving_manifest( + bytes(row[8]), + checksum=str(row[9]), + format_version=int(row[10]), + ) + except ValueError: + return None # The marker is written last in the same transaction that inserts # all units and token rows. A committed marker therefore denotes @@ -355,30 +419,63 @@ def load_persisted_score_corpus( for document_id, section_ids in allowed_by_document.items() for section_id in section_ids ] + unit_cache_key = revision_key + all_unit_rows = self._score_unit_rows_cache.get(unit_cache_key, []) + units_cache_hit = unit_cache_key in self._score_unit_rows_cache 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] - stage_started = time.perf_counter() - 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()) + if not units_cache_hit: + stage_started = time.perf_counter() + manifest_rows: list[Sequence[object]] = [] + for document_id, job_result_id in revisions: + payload = decoded_manifests.get((document_id, job_result_id), {}) + raw_units = payload.get("map_units") + if not isinstance(raw_units, list): + manifest_rows = [] + break + for raw_unit in raw_units: + if not isinstance(raw_unit, dict): + manifest_rows = [] + break + row_id = str(raw_unit.get("row_id") or "").strip() + unit_id = str(raw_unit.get("unit_id") or "").strip() + if not row_id or not unit_id: + manifest_rows = [] + break + manifest_rows.append( + ( + row_id, + document_id, + unit_id, + str(raw_unit.get("section_id") or ""), + int(raw_unit.get("path_token_count") or 0), + int(raw_unit.get("content_token_count") or 0), + ) + ) + if not manifest_rows and raw_units: + break + expected_unit_count = sum(int(row[3]) for row in manifests) + all_unit_rows = ( + manifest_rows + if len(manifest_rows) == expected_unit_count + else [] + ) + if expected_unit_count and not all_unit_rows: + return None + self._score_unit_rows_cache[unit_cache_key] = all_unit_rows + else: + stage_started = time.perf_counter() + allowed_pairs_set = set(allowed_pairs) + unit_rows = [ + row + for row in all_unit_rows + if (str(row[1]), str(row[3])) in allowed_pairs_set + ] _logger.info( - "retrieval map-index load stage=units seconds=%.3f rows=%d", - time.perf_counter() - stage_started, + "retrieval map-index load stage=units seconds=%.3f rows=%d cache_hit=%s", + time.perf_counter() - stage_started if not units_cache_hit else 0.0, len(unit_rows), + units_cache_hit, ) map_unit_ids = [str(row[0]) for row in unit_rows] unique_queries = list(dict.fromkeys(str(query) for query in queries)) @@ -394,6 +491,48 @@ def load_persisted_score_corpus( ) frequencies: Dict[Tuple[str, str], Dict[str, int]] = {} if map_unit_ids and query_tokens: + full_frequency_map = self._score_frequency_cache.get(revision_key) + if full_frequency_map is None: + full_frequency_map = {} + for document_id, job_result_id in revisions: + payload = decoded_statistics.get((document_id, job_result_id), {}) + unit_frequencies = payload.get("unit_frequencies", {}) + if not isinstance(unit_frequencies, dict): + continue + for map_unit_id, by_channel in unit_frequencies.items(): + if not isinstance(by_channel, dict): + continue + for channel in _MAP_SCORE_CHANNELS: + values = by_channel.get(channel, {}) + if isinstance(values, dict): + full_frequency_map[(str(map_unit_id), channel)] = { + str(token): int(value) + for token, value in values.items() + } + self._score_frequency_cache[revision_key] = full_frequency_map + frequencies = { + key: { + token: value + for token, value in values.items() + if token in query_tokens + } + for key, values in full_frequency_map.items() + if key[0] in map_unit_ids + } + expected_units_by_revision = { + (str(row[0]), str(row[1])): int(row[3]) for row in manifests + } + statistics_complete = all( + int( + decoded_statistics.get((document_id, job_result_id), {}).get( + "unit_count", -1 + ) + ) + == expected_units_by_revision.get((document_id, job_result_id), -1) + for document_id, job_result_id in revisions + ) + manifest_frequency_complete = bool(map_unit_ids) and statistics_complete + if map_unit_ids and query_tokens and not manifest_frequency_complete: query_token_hashes = [ sha256(token.encode("utf-8")).hexdigest() for token in query_tokens ] @@ -420,12 +559,16 @@ def load_persisted_score_corpus( len(map_unit_ids), ) - term_scores = self._load_term_scores( - cur, - map_unit_ids=map_unit_ids, - queries=unique_queries, - query_tokens_by_query=query_tokens_by_query, - ) + term_cache_key = (revision_key, tuple(unique_queries)) + term_scores = self._score_term_cache.get(term_cache_key) + if term_scores is None: + term_scores = self._load_term_scores( + cur, + map_unit_ids=map_unit_ids, + queries=unique_queries, + query_tokens_by_query=query_tokens_by_query, + ) + self._score_term_cache[term_cache_key] = term_scores _logger.info( "retrieval map-index load stage=complete units=%d queries=%d", len(unit_rows), @@ -481,42 +624,81 @@ def _load_term_scores( ) -> Dict[str, Tuple[float, ...]]: if not map_unit_ids or not queries: return {} - expressions: List[str] = [] - params: List[object] = [] + # Keep candidate selection in PostgreSQL so the trigram index can + # discard non-matching units, but compute the exact score once per + # returned row in Python. The previous query generated one POSITION + # expression and one OR predicate for every query token. Long planner + # subgoals therefore produced very large SQL statements and repeated + # substring evaluation for the same row. LIKE ANY keeps the SQL shape + # constant while preserving the existing literal substring semantics in + # the final Python scoring pass (LIKE may over-select wildcard matches, + # which are rejected by the literal checks below). + candidate_patterns: list[str] = [] 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" + candidate_patterns.append(f"%{query_lower}%") + candidate_patterns.extend( + f"%{token}%" for token in query_tokens_by_query[query] if token ) - params.append(query_lower) - params.extend(query_tokens_by_query[query]) - params.append(list(map_unit_ids)) + candidate_patterns = list(dict.fromkeys(candidate_patterns)) + if not candidate_patterns: + return {} stage_started = time.perf_counter() - cur.execute( - "SELECT id, " + ", ".join(expressions) + " " - "FROM document_map_units WHERE id = ANY(%s)", - params, - ) + if len(candidate_patterns) > _TERM_SCORE_FULL_SCAN_PATTERN_THRESHOLD: + # Long subgoals make LIKE ANY increasingly expensive even with the + # trigram index. The map-unit id list is already bounded by the + # pinned serving projection, so one text fetch plus literal Python + # checks avoids a query whose shape grows with token count. + cur.execute( + "SELECT id, term_search_text_lower " + "FROM document_map_units WHERE id = ANY(%s)", + (list(map_unit_ids),), + ) + candidate_mode = "full_scan" + else: + cur.execute( + "SELECT id, term_search_text_lower " + "FROM document_map_units " + "WHERE id = ANY(%s) AND term_search_text_lower LIKE ANY(%s)", + (list(map_unit_ids), candidate_patterns), + ) + candidate_mode = "trigram" rows = cur.fetchall() _logger.info( - "retrieval map-index load stage=term_scores units=%d queries=%d seconds=%.3f", + "retrieval map-index load stage=term_scores units=%d queries=%d mode=%s seconds=%.3f", len(map_unit_ids), len(queries), + candidate_mode, time.perf_counter() - stage_started, ) - return { - str(row[0]): tuple(float(value) for value in row[1:]) - for row in rows - } + scores_by_unit: Dict[str, Tuple[float, ...]] = {} + for row in rows: + if len(row) < 2: + continue + unit_id = str(row[0]) + haystack = str(row[1] or "").lower() + scores: list[float] = [] + for query in queries: + query_lower = query.lower().strip() + if not query_lower: + scores.append(0.0) + elif query_lower in haystack: + scores.append(100.0) + else: + scores.append( + float( + sum( + 1 + for token in query_tokens_by_query[query] + if token in haystack + ) + ) + ) + if any(score > 0.0 for score in scores): + scores_by_unit[unit_id] = tuple(scores) + return scores_by_unit def _load_persisted_bm25_stats( self, 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 2e82be08..52cf8867 100644 --- a/packages/shared-python/shared/services/retrieval/nav/nav_orchestrate.py +++ b/packages/shared-python/shared/services/retrieval/nav/nav_orchestrate.py @@ -203,6 +203,8 @@ def _relit_map( """ relit = prepared q = (query or "").strip() + if relit is None and q: + relit = state.relit_map_cache.get(q) if relit is None and q: try: from .nav_map_scores import relight_map_for_query @@ -215,6 +217,7 @@ def _relit_map( ) if scores: relit = (scores, units, highlights) + state.relit_map_cache[q] = relit except Exception: relit = None if relit is None: diff --git a/packages/shared-python/shared/services/retrieval/nav/nav_types.py b/packages/shared-python/shared/services/retrieval/nav/nav_types.py index 8f3366c0..3c6b5241 100644 --- a/packages/shared-python/shared/services/retrieval/nav/nav_types.py +++ b/packages/shared-python/shared/services/retrieval/nav/nav_types.py @@ -305,6 +305,12 @@ class NavState: # PLAN's rewritten retrieval_query per subgoal after widen. Overrides the # planned query for the next harvest, and re-scores the shared map with it. subgoal_refined_queries: Dict[str, str] = field(default_factory=dict) + # Episode-local map scores keyed by retrieval query. Checklist waves may + # revisit the same subgoal query; reuse the exact score snapshot instead of + # rebuilding the persisted index and rescoring the corpus. + relit_map_cache: Dict[ + str, Tuple[Dict[str, float], Dict[str, float], List[str]] + ] = field(default_factory=dict) # Per-subgoal "seen but not selected" section ids — hidden from later map # views for that subgoal so widen surfaces siblings instead of dead ends. subgoal_dismissed_section_ids: Dict[str, set[str]] = field(default_factory=dict) diff --git a/packages/shared-python/shared/services/retrieval/nav_snapshot.py b/packages/shared-python/shared/services/retrieval/nav_snapshot.py index 4eec0404..c68d6d8c 100644 --- a/packages/shared-python/shared/services/retrieval/nav_snapshot.py +++ b/packages/shared-python/shared/services/retrieval/nav_snapshot.py @@ -14,10 +14,25 @@ from collections.abc import Callable, Iterator, Mapping from typing import Any, Protocol -from sqlalchemy import Executable, literal, select, tuple_ +from sqlalchemy import ( + ARRAY, + Executable, + String, + bindparam, + cast, + func, + literal, + select, + tuple_, +) from sqlalchemy.engine import Result -from shared.models.database.document import Document, DocumentChunk, DocumentSection +from shared.models.database.document import ( + Document, + DocumentChunk, + DocumentSection, + RetrievalServingRevisionManifest, +) from shared.models.database.job_result import JobResult from shared.services.retrieval.nav.nav_knowhere import ( LazyKnowhereProvider, @@ -29,12 +44,13 @@ knowhere_database_url, ) from shared.services.retrieval.search.section_filters import is_excluded_section +from shared.services.retrieval.serving_manifest import decode_serving_manifest +from shared.services.retrieval.manifest_cache import get_cached_manifest_payloads # Keep each payload query bounded under the API's 30-second statement timeout. -# Ten-thousand-row keyset pages avoid OFFSET scans while keeping each payload -# statement bounded. The contract benchmark verifies this page size against -# the full 2 KiB content and metadata payload. +# Ten-thousand-row keyset pages avoid OFFSET scans while keeping the reference +# payload bounded under the API's 30-second asyncpg command timeout. _CHUNK_BATCH_SIZE = 10_000 # Keep revision predicates bounded while reducing round trips for large # namespaces. Keyset paging still caps each payload query at 10,000 rows. @@ -48,6 +64,9 @@ class SnapshotSession(Protocol): async def execute(self, statement: Executable) -> Result[tuple[object, ...]]: raise NotImplementedError + async def rollback(self) -> None: + raise NotImplementedError + @dataclass(frozen=True) class NavSnapshot: @@ -57,6 +76,7 @@ class NavSnapshot: chunk_ref_index: Mapping[str, dict[str, Any]] document_ids: list[str] document_titles: dict[str, str] + document_revisions: Mapping[str, str] | None def close(self) -> None: close = getattr(self.provider, "close", None) @@ -70,9 +90,12 @@ def build_nav_snapshot( sections_by_doc: dict[str, list[SectionRow]], units_by_doc: dict[str, list[UnitRow]], chunk_ref_index: dict[str, dict[str, Any]], + document_revisions: Mapping[str, str] | None = None, ) -> NavSnapshot: """Assemble provider + index from already-fetched rows (also used by tests).""" - doc_ids = [did for did in document_titles if did in sections_by_doc or did in units_by_doc] + doc_ids = [ + did for did in document_titles if did in sections_by_doc or did in units_by_doc + ] if not doc_ids: raise ValueError("nav snapshot requires at least one active document") @@ -90,7 +113,18 @@ def build_nav_snapshot( provider=provider, chunk_ref_index=dict(chunk_ref_index), document_ids=list(provider.document_ids()), - document_titles={did: document_titles.get(did, did) for did in provider.document_ids()}, + document_titles={ + did: document_titles.get(did, did) for did in provider.document_ids() + }, + document_revisions=( + { + did: str(document_revisions.get(did, "")) + for did in provider.document_ids() + if document_revisions.get(did) + } + if document_revisions is not None + else None + ), ) @@ -102,19 +136,38 @@ async def load_nav_snapshot( exclude_document_ids: list[str] | None = None, exclude_sections: list[dict[str, str]] | None = None, lazy: bool = False, + revision_pins: Mapping[str, str] | None = None, ) -> NavSnapshot: """Preload namespace current revision into a sync map-nav snapshot.""" - excluded_docs = [str(x).strip() for x in (exclude_document_ids or ()) if str(x).strip()] + excluded_docs = [ + str(x).strip() for x in (exclude_document_ids or ()) if str(x).strip() + ] excluded_secs = list(exclude_sections or ()) - doc_stmt = ( - select(Document.document_id, Document.source_file_name, Document.current_job_result_id) - .where(Document.user_id == user_id) - .where(Document.namespace == namespace) - .where(Document.status == "active") - .where(Document.current_job_result_id.is_not(None)) - .order_by(Document.document_id) - ) + if revision_pins is None: + doc_stmt = ( + select( + Document.document_id, + Document.source_file_name, + Document.current_job_result_id, + ) + .where(Document.user_id == user_id) + .where(Document.namespace == namespace) + .where(Document.status == "active") + .where(Document.current_job_result_id.is_not(None)) + .order_by(Document.document_id) + ) + else: + pinned_document_ids = [str(document_id) for document_id in revision_pins] + if not pinned_document_ids: + raise ValueError("revision pins must include at least one document") + doc_stmt = ( + select(Document.document_id, Document.source_file_name) + .where(Document.user_id == user_id) + .where(Document.namespace == namespace) + .where(Document.document_id.in_(pinned_document_ids)) + .order_by(Document.document_id) + ) if excluded_docs: doc_stmt = doc_stmt.where(Document.document_id.notin_(excluded_docs)) doc_rows = list((await db.execute(doc_stmt)).all()) @@ -127,12 +180,18 @@ async def load_nav_snapshot( document_titles: dict[str, str] = {} current_job_result_ids: set[str] = set() document_revisions: list[tuple[str, str]] = [] - for document_id, source_file_name, current_job_result_id in doc_rows: + for doc_row in doc_rows: + document_id = doc_row[0] + source_file_name = doc_row[1] did = str(document_id) title = str(source_file_name or "").strip() or did document_titles[did] = title - if current_job_result_id: - job_result_id = str(current_job_result_id) + job_result_id = ( + str(revision_pins.get(did, "")) + if revision_pins is not None + else str(doc_row[2] or "") + ) + if job_result_id: current_job_result_ids.add(job_result_id) document_revisions.append((did, job_result_id)) @@ -147,28 +206,47 @@ async def load_nav_snapshot( if job_result_id and job_id } - sections_by_doc, section_path_by_id = await _load_sections( + manifest_sections = await _load_manifest_sections( db, document_revisions=document_revisions, exclude_sections=excluded_secs, + job_id_by_result_id=job_id_by_result_id, ) - if lazy: - chunk_ids_by_doc, chunk_ref_index, remounted_assets = await _load_chunk_index( + if manifest_sections is None: + sections_by_doc, section_path_by_id = await _load_sections( db, document_revisions=document_revisions, exclude_sections=excluded_secs, - section_path_by_id=section_path_by_id, - job_id_by_result_id=job_id_by_result_id, ) + manifest_chunk_index = None + else: + sections_by_doc, section_path_by_id, manifest_chunk_index = manifest_sections + if lazy: + if manifest_chunk_index is None: + chunk_ids_by_doc, chunk_ref_index, remounted_assets = await _load_chunk_index( + db, + document_revisions=document_revisions, + exclude_sections=excluded_secs, + section_path_by_id=section_path_by_id, + job_id_by_result_id=job_id_by_result_id, + ) + else: + chunk_ids_by_doc, chunk_ref_index, remounted_assets = manifest_chunk_index kept_titles = { - did: title for did, title in document_titles.items() if sections_by_doc.get(did) + did: title + for did, title in document_titles.items() + if sections_by_doc.get(did) } if not kept_titles: raise ValueError( f"nav snapshot empty after excludes for " f"user_id={user_id!r} namespace={namespace!r}" ) - revisions = {did: result_id for did, result_id in document_revisions if did in kept_titles} + revisions = { + did: result_id + for did, result_id in document_revisions + if did in kept_titles + } store = ReadOnlyChunkStore( dsn=knowhere_database_url(), revisions=revisions, @@ -192,7 +270,9 @@ async def load_nav_snapshot( chunk_store=store, known_chunk_ids=chunk_ids_by_doc.get(did, ()), root_asset_ids=remounted_assets.get(did, {}).get("root", ()), - remounted_assets_by_section=remounted_assets.get(did, {}).get("owners", {}), + remounted_assets_by_section=remounted_assets.get(did, {}).get( + "owners", {} + ), ) for did in kept_titles ] @@ -215,7 +295,10 @@ async def load_nav_snapshot( resolver=store.load_chunk_reference_metadata, ), document_ids=list(provider.document_ids()), - document_titles={did: kept_titles.get(did, did) for did in provider.document_ids()}, + document_titles={ + did: kept_titles.get(did, did) for did in provider.document_ids() + }, + document_revisions=dict(revisions), ) units_by_doc, chunk_ref_index = await _load_chunks( @@ -228,9 +311,7 @@ async def load_nav_snapshot( # Keep only documents that still have sections after exclude filters. kept_titles = { - did: title - for did, title in document_titles.items() - if sections_by_doc.get(did) + did: title for did, title in document_titles.items() if sections_by_doc.get(did) } if not kept_titles: raise ValueError( @@ -243,9 +324,174 @@ async def load_nav_snapshot( sections_by_doc={did: sections_by_doc.get(did, []) for did in kept_titles}, units_by_doc={did: units_by_doc.get(did, []) for did in kept_titles}, chunk_ref_index=chunk_ref_index, + document_revisions={ + document_id: job_result_id + for document_id, job_result_id in document_revisions + if document_id in kept_titles + }, ) +async def _load_manifest_sections( + db: SnapshotSession, + *, + document_revisions: list[tuple[str, str]], + exclude_sections: list[dict[str, str]], + job_id_by_result_id: dict[str, str], +) -> tuple[ + dict[str, list[SectionRow]], + dict[str, str], + tuple[dict[str, list[str]], dict[str, dict[str, Any]], dict[str, dict[str, Any]]], +] | None: + """Load section metadata from complete serving manifests when available.""" + cached_payloads = get_cached_manifest_payloads( + db, + revisions=dict(document_revisions), + ) + if cached_payloads is not None and len(cached_payloads) == len(document_revisions): + manifest_entries: list[tuple[object, ...]] = [ + (document_id, job_result_id, payload, None, None) + for document_id, job_result_id in document_revisions + for payload in [cached_payloads.get((document_id, job_result_id))] + if payload is not None + ] + if len(manifest_entries) != len(document_revisions): + return None + else: + statement = select( + RetrievalServingRevisionManifest.document_id, + RetrievalServingRevisionManifest.job_result_id, + RetrievalServingRevisionManifest.payload_zlib, + RetrievalServingRevisionManifest.checksum, + RetrievalServingRevisionManifest.format_version, + ).where( + tuple_( + RetrievalServingRevisionManifest.document_id, + RetrievalServingRevisionManifest.job_result_id, + ).in_(document_revisions) + ) + try: + rows = (await db.execute(statement)).all() + except Exception: + await db.rollback() + return None + if len(rows) != len(document_revisions): + return None + manifest_entries = [tuple(row) for row in rows] + by_doc: dict[str, list[SectionRow]] = {} + path_by_id: dict[str, str] = {} + ids_by_doc: dict[str, list[str]] = {} + ref_index: dict[str, dict[str, Any]] = {} + root_assets_by_doc: dict[str, set[str]] = {} + text_connections_by_doc: dict[str, list[tuple[str, str]]] = {} + try: + for document_id, _job_result_id, payload_zlib, checksum, format_version in manifest_entries: + if isinstance(payload_zlib, dict): + payload = payload_zlib + else: + if not isinstance(payload_zlib, (bytes, bytearray, memoryview)): + return None + if checksum is None or format_version is None: + return None + payload = decode_serving_manifest( + bytes(payload_zlib), + checksum=str(checksum), + format_version=int(str(format_version)), + ) + raw_sections = payload.get("sections") + if not isinstance(raw_sections, list): + return None + for raw_section in raw_sections: + if not isinstance(raw_section, dict): + return None + section_path = str(raw_section.get("section_path") or "") + if is_excluded_section( + document_id=str(document_id), + section_path=section_path, + exclude_sections=exclude_sections, + ): + continue + section_id = str(raw_section.get("section_id") or "") + if not section_id or not section_path: + return None + section = SectionRow( + section_id=section_id, + parent_section_id=( + str(raw_section["parent_section_id"]) + if raw_section.get("parent_section_id") + else None + ), + section_path=section_path, + section_title=str(raw_section.get("section_title") or "").strip(), + section_level=int(raw_section.get("section_level") or 0), + summary=str(raw_section.get("summary") or "").strip(), + sort_order=int(raw_section.get("sort_order") or 0), + ) + by_doc.setdefault(str(document_id), []).append(section) + path_by_id[section_id] = section_path + raw_chunks = payload.get("chunks") + if not isinstance(raw_chunks, list): + return None + for raw_chunk in raw_chunks: + if not isinstance(raw_chunk, dict): + return None + chunk_id = str(raw_chunk.get("chunk_id") or "").strip() + if not chunk_id: + return None + section_id = ( + str(raw_chunk["section_id"]) + if raw_chunk.get("section_id") + else None + ) + section_path = path_by_id.get(section_id) if section_id else None + if is_excluded_section( + document_id=str(document_id), + section_path=section_path, + exclude_sections=exclude_sections, + ) or (section_id and section_id not in path_by_id): + continue + chunk_type = str(raw_chunk.get("chunk_type") or "text") + meta = { + "document_id": str(document_id), + "section_path": section_path, + "chunk_type": chunk_type, + "file_path": None, + "job_id": job_id_by_result_id.get(str(_job_result_id)), + } + document_key = str(document_id) + ids_by_doc.setdefault(document_key, []).append(chunk_id) + ref_index[chunk_id] = meta + ref_index[f"{document_key}:{chunk_id}"] = meta + if ( + chunk_type in {"image", "table"} + and section_id + and section_path == "Root" + ): + root_assets_by_doc.setdefault(document_key, set()).add(chunk_id) + connections = raw_chunk.get("connect_to") + if chunk_type == "text" and isinstance(connections, list): + for connection in connections: + target = ( + str(connection.get("target") or "").strip() + if isinstance(connection, dict) + else str(connection or "").strip() + ) + if target: + text_connections_by_doc.setdefault(document_key, []).append( + (section_id or "", target) + ) + except (TypeError, ValueError, KeyError): + return None + remounted: dict[str, dict[str, Any]] = {} + for document_id, asset_ids in root_assets_by_doc.items(): + owners: dict[str, list[str]] = {} + for section_id, target in text_connections_by_doc.get(document_id, ()): + if target in asset_ids: + owners.setdefault(section_id, []).append(target) + remounted[document_id] = {"root": sorted(asset_ids), "owners": owners} + return by_doc, path_by_id, (ids_by_doc, ref_index, remounted) + + async def _load_chunk_index( db: SnapshotSession, *, @@ -431,62 +677,109 @@ async def _load_sections( exclude_sections: list[dict[str, str]], ) -> tuple[dict[str, list[SectionRow]], dict[str, str]]: # Captured pairs replace DocumentSection.job_result_id == Document.current_job_result_id. - stmt = ( - select( - DocumentSection.document_id, - DocumentSection.section_id, - DocumentSection.parent_section_id, - DocumentSection.section_path, - DocumentSection.section_title, - DocumentSection.section_level, - DocumentSection.summary, - DocumentSection.sort_order, + by_doc: dict[str, list[SectionRow]] = {} + path_by_id: dict[str, str] = {} + document_ids = [document_id for document_id, _ in document_revisions] + job_result_ids = [job_result_id for _, job_result_id in document_revisions] + revision_rows = ( + func.unnest( + cast( + bindparam("section_document_ids", value=document_ids), ARRAY(String()) + ), + cast( + bindparam("section_job_result_ids", value=job_result_ids), + ARRAY(String()), + ), ) - .where( - tuple_( + .table_valued("document_id", "job_result_id") + .render_derived(name="revisions") + ) + last_key: tuple[str, str, int, str] | None = None + query_seconds = 0.0 + assembly_seconds = 0.0 + query_count = 0 + row_count = 0 + while True: + stmt = ( + select( DocumentSection.document_id, + DocumentSection.section_id, + DocumentSection.parent_section_id, + DocumentSection.section_path, + DocumentSection.section_title, + DocumentSection.section_level, + DocumentSection.summary, + DocumentSection.sort_order, DocumentSection.job_result_id, - ).in_(document_revisions) - ) - .order_by( - DocumentSection.document_id, - DocumentSection.sort_order, - DocumentSection.section_id, + ) + .join( + revision_rows, + (DocumentSection.document_id == revision_rows.c.document_id) + & (DocumentSection.job_result_id == revision_rows.c.job_result_id), + ) + .order_by( + DocumentSection.document_id, + DocumentSection.job_result_id, + DocumentSection.sort_order, + DocumentSection.section_id, + ) + .limit(_CHUNK_BATCH_SIZE) ) - ) - - by_doc: dict[str, list[SectionRow]] = {} - path_by_id: dict[str, str] = {} - query_started = time.perf_counter() - rows = (await db.execute(stmt)).all() - query_seconds = time.perf_counter() - query_started - assembly_started = time.perf_counter() - for row in rows: - document_id = str(row[0]) - section_path = str(row[3] or "") - if is_excluded_section( - document_id=document_id, - section_path=section_path, - exclude_sections=exclude_sections, - ): - continue - section_id = str(row[1]) - section = SectionRow( - section_id=section_id, - parent_section_id=str(row[2]) if row[2] else None, - section_path=section_path, - section_title=str(row[4] or "").strip(), - section_level=int(row[5] or 0), - summary=str(row[6] or "").strip(), - sort_order=int(row[7] or 0), + if last_key is not None: + stmt = stmt.where( + tuple_( + DocumentSection.document_id, + DocumentSection.job_result_id, + DocumentSection.sort_order, + DocumentSection.section_id, + ) + > tuple_(*[literal(value) for value in last_key]) + ) + query_started = time.perf_counter() + rows = (await db.execute(stmt)).all() + query_seconds += time.perf_counter() - query_started + query_count += 1 + if not rows: + break + row_count += len(rows) + assembly_started = time.perf_counter() + for row in rows: + document_id = str(row[0]) + section_path = str(row[3] or "") + if is_excluded_section( + document_id=document_id, + section_path=section_path, + exclude_sections=exclude_sections, + ): + continue + section_id = str(row[1]) + section = SectionRow( + section_id=section_id, + parent_section_id=str(row[2]) if row[2] else None, + section_path=section_path, + section_title=str(row[4] or "").strip(), + section_level=int(row[5] or 0), + summary=str(row[6] or "").strip(), + sort_order=int(row[7] or 0), + ) + by_doc.setdefault(document_id, []).append(section) + path_by_id[section_id] = section_path + assembly_seconds += time.perf_counter() - assembly_started + last = rows[-1] + last_key = ( + str(last[0]), + str(last[8]), + int(last[7] or 0), + str(last[1]), ) - by_doc.setdefault(document_id, []).append(section) - path_by_id[section_id] = section_path + if len(rows) < _CHUNK_BATCH_SIZE: + break _logger.info( - "retrieval snapshot phase=sections rows=%d query_seconds=%.3f assembly_seconds=%.3f", - len(rows), + "retrieval snapshot phase=sections rows=%d queries=%d query_seconds=%.3f assembly_seconds=%.3f", + row_count, + query_count, query_seconds, - time.perf_counter() - assembly_started, + assembly_seconds, ) return by_doc, path_by_id diff --git a/packages/shared-python/shared/services/retrieval/publication_content.py b/packages/shared-python/shared/services/retrieval/publication_content.py index 19a7f207..34092436 100644 --- a/packages/shared-python/shared/services/retrieval/publication_content.py +++ b/packages/shared-python/shared/services/retrieval/publication_content.py @@ -14,6 +14,7 @@ ) from shared.services.retrieval.map_unit_index import replace_document_map_units from shared.services.retrieval.publication_models import DocumentPublicationScope +from shared.services.retrieval.serving_manifest import persist_revision_serving_state from shared.services.retrieval.search.lexical_text import ( build_content_lexical_text, build_content_search_text, @@ -91,6 +92,8 @@ def replace_document_revision_content( ) db.flush() replace_document_map_units(db, scope=scope) + db.flush() + persist_revision_serving_state(db, scope=scope) class DocumentSectionPublisher: diff --git a/packages/shared-python/shared/services/retrieval/publication_service.py b/packages/shared-python/shared/services/retrieval/publication_service.py index e201916e..d7c879de 100644 --- a/packages/shared-python/shared/services/retrieval/publication_service.py +++ b/packages/shared-python/shared/services/retrieval/publication_service.py @@ -31,6 +31,13 @@ ExistingDocumentScope, PublishedDocumentState, ) +from shared.services.retrieval.serving_generation import ( + advance_namespace_generation, + lock_namespace_generation, +) +from shared.services.retrieval.serving_manifest import ( + rebuild_namespace_serving_statistics, +) def utc_now_naive() -> datetime: @@ -122,6 +129,23 @@ def _publish_document_state_for_job( skipped_all_duplicate=True, ) + existing_namespace = None + if document_id: + existing_namespace = db.execute( + select(Document.namespace).where( + Document.document_id == str(document_id) + ) + ).scalar_one_or_none() + namespaces_to_lock = {namespace} + if existing_namespace: + namespaces_to_lock.add(str(existing_namespace)) + for namespace_to_lock in sorted(namespaces_to_lock): + lock_namespace_generation( + db, + user_id=str(job.user_id), + namespace=namespace_to_lock, + ) + document = self._upsert_document_revision( db, job=job, @@ -156,6 +180,22 @@ def _publish_document_state_for_job( ) db.flush() + rebuild_namespace_serving_statistics( + db, + user_id=scope.user_id, + namespace=scope.namespace, + ) + if existing_namespace and str(existing_namespace) != scope.namespace: + rebuild_namespace_serving_statistics( + db, + user_id=scope.user_id, + namespace=str(existing_namespace), + ) + advance_namespace_generation( + db, + user_id=scope.user_id, + namespace=scope.namespace, + ) return PublishedDocumentState( user_id=str(job.user_id), namespace=namespace, @@ -209,6 +249,7 @@ def _upsert_document_revision( ) return None document.status = "active" + document.namespace = namespace document.archived_at = None document.current_job_result_id = job_result_id document.source_file_name = source_file_name or document.source_file_name diff --git a/packages/shared-python/shared/services/retrieval/search/channels.py b/packages/shared-python/shared/services/retrieval/search/channels.py index 1ddea1be..cdf67326 100644 --- a/packages/shared-python/shared/services/retrieval/search/channels.py +++ b/packages/shared-python/shared/services/retrieval/search/channels.py @@ -8,6 +8,7 @@ from __future__ import annotations import time +from collections.abc import Mapping from typing import Any from loguru import logger @@ -61,7 +62,7 @@ FROM document_chunks dc JOIN documents d ON d.document_id = dc.document_id - AND d.current_job_result_id = dc.job_result_id + {revision_join} LEFT JOIN document_sections ds ON ds.section_id = dc.section_id JOIN job_results jr @@ -69,12 +70,42 @@ WHERE d.user_id = :user_id AND d.namespace = :namespace AND d.status = 'active' + {revision_clause} {exclude_clause} {extra_filters} ) """ +def _build_revision_scope( + revision_pins: Mapping[str, str] | None, +) -> tuple[str, str, dict[str, Any]]: + if revision_pins is None: + return ( + "AND d.current_job_result_id = dc.job_result_id", + "", + {}, + ) + + pairs = [ + (str(document_id).strip(), str(job_result_id).strip()) + for document_id, job_result_id in revision_pins.items() + if str(document_id).strip() and str(job_result_id).strip() + ] + if not pairs: + return "", "AND FALSE", {} + + params: dict[str, Any] = {} + placeholders: list[str] = [] + for index, (document_id, job_result_id) in enumerate(pairs): + document_key = f"_pin_document_{index}" + revision_key = f"_pin_revision_{index}" + placeholders.append(f"(:{document_key}, :{revision_key})") + params[document_key] = document_id + params[revision_key] = job_result_id + return "", f"AND (dc.document_id, dc.job_result_id) IN ({', '.join(placeholders)})", params + + def _build_exclude_clause(exclude_document_ids: list[str]) -> str: if not exclude_document_ids: return "" @@ -228,6 +259,7 @@ async def path_channel( allowed_chunk_types: set[str] | None = None, signal_paths: list[str] | None = None, filter_mode: str = "delete", + revision_pins: Mapping[str, str] | None = None, ) -> list[dict[str, Any]]: """Path channel: BM25 over pre-tokenized path search text. @@ -246,6 +278,7 @@ async def path_channel( signal_paths=signal_paths, filter_mode=filter_mode, search_field="path_search_text", + revision_pins=revision_pins, ) @@ -261,6 +294,7 @@ async def content_channel( allowed_chunk_types: set[str] | None = None, signal_paths: list[str] | None = None, filter_mode: str = "delete", + revision_pins: Mapping[str, str] | None = None, ) -> list[dict[str, Any]]: """Content channel: BM25 over pre-tokenized content search text.""" return await _bm25_channel( @@ -275,6 +309,7 @@ async def content_channel( signal_paths=signal_paths, filter_mode=filter_mode, search_field="content_search_text", + revision_pins=revision_pins, ) @@ -291,6 +326,7 @@ async def _bm25_channel( signal_paths: list[str] | None, filter_mode: str, search_field: str, + revision_pins: Mapping[str, str] | None, ) -> list[dict[str, Any]]: if search_field not in {"content_search_text", "path_search_text"}: raise ValueError(f"Unsupported search_field: {search_field}") @@ -317,7 +353,14 @@ async def _bm25_channel( params.update(extra_params) params.update(section_params) + revision_join, revision_clause, revision_params = _build_revision_scope( + revision_pins + ) + params.update(revision_params) + corpus_cte = _SCOPED_CORPUS_CTE.format( + revision_join=revision_join, + revision_clause=revision_clause, exclude_clause=exclude_clause, extra_filters=extra_sql, ) @@ -409,6 +452,7 @@ async def term_channel( allowed_chunk_types: set[str] | None = None, signal_paths: list[str] | None = None, filter_mode: str = "delete", + revision_pins: Mapping[str, str] | None = None, ) -> list[dict[str, Any]]: """Term/grep channel: substring matching on term_search_text. @@ -435,6 +479,11 @@ async def term_channel( ) params.update(extra_params) + revision_join, revision_clause, revision_params = _build_revision_scope( + revision_pins + ) + params.update(revision_params) + ilike_conditions = [] for i, unit in enumerate(query_tokens): param_key = f"unit_{i}" @@ -448,6 +497,8 @@ async def term_channel( where_clause = " OR ".join(ilike_conditions) sql = ( _SCOPED_CORPUS_CTE.format( + revision_join=revision_join, + revision_clause=revision_clause, exclude_clause=exclude_clause, extra_filters=extra_sql ) + f""" diff --git a/packages/shared-python/shared/services/retrieval/search/discovery.py b/packages/shared-python/shared/services/retrieval/search/discovery.py index 2e3fdd3c..d71ccdd9 100644 --- a/packages/shared-python/shared/services/retrieval/search/discovery.py +++ b/packages/shared-python/shared/services/retrieval/search/discovery.py @@ -18,7 +18,9 @@ from __future__ import annotations +import asyncio import time +from collections.abc import Awaitable, Callable, Mapping from dataclasses import dataclass, field from typing import Any @@ -68,11 +70,13 @@ async def bottom_discovery( channels: list[str] | None = None, channel_weights: dict[str, float] | None = None, internal_recall_k: int | None = None, + revision_pins: Mapping[str, str] | None = None, **_kwargs: Any, ) -> DiscoveryResult: """Run 3-channel BM25 discovery plus RRF fusion.""" t0 = time.monotonic() try: + del db allowed_chunk_types = chunk_types effective_recall_k = ( internal_recall_k @@ -81,13 +85,10 @@ async def bottom_discovery( ) active_channels = set(channels) if channels else {"path", "content", "term"} - path_rows: list[dict[str, Any]] = [] - content_rows: list[dict[str, Any]] = [] - term_rows: list[dict[str, Any]] = [] - - if "path" in active_channels: - path_rows = await path_channel( - db, + path_rows, content_rows, term_rows = await asyncio.gather( + _run_channel( + path_channel, + enabled="path" in active_channels, user_id=user_id, namespace=namespace, query=query, @@ -97,11 +98,11 @@ async def bottom_discovery( allowed_chunk_types=allowed_chunk_types, signal_paths=signal_paths, filter_mode=filter_mode, - ) - - if "content" in active_channels: - content_rows = await content_channel( - db, + revision_pins=revision_pins, + ), + _run_channel( + content_channel, + enabled="content" in active_channels, user_id=user_id, namespace=namespace, query=query, @@ -111,11 +112,11 @@ async def bottom_discovery( allowed_chunk_types=allowed_chunk_types, signal_paths=signal_paths, filter_mode=filter_mode, - ) - - if "term" in active_channels: - term_rows = await term_channel( - db, + revision_pins=revision_pins, + ), + _run_channel( + term_channel, + enabled="term" in active_channels, user_id=user_id, namespace=namespace, query=query, @@ -125,7 +126,9 @@ async def bottom_discovery( allowed_chunk_types=allowed_chunk_types, signal_paths=signal_paths, filter_mode=filter_mode, - ) + revision_pins=revision_pins, + ), + ) default_weights = { "path": CHANNEL_WEIGHT_PATH, @@ -194,3 +197,20 @@ async def bottom_discovery( latency = int((time.monotonic() - t0) * 1000) logger.error(f" search.bottom_discovery failed: {exc}") return DiscoveryResult(status="error", error=str(exc), latency_ms=latency) + + +async def _run_channel( + channel: Callable[..., Awaitable[list[dict[str, Any]]]], + *, + enabled: bool, + **kwargs: Any, +) -> list[dict[str, Any]]: + if not enabled: + return [] + + # Import lazily so discovery remains usable by lightweight contract tests + # without creating a database context until a channel is actually enabled. + from shared.core.database import get_db_context + + async with get_db_context() as channel_db: + return await channel(channel_db, **kwargs) diff --git a/packages/shared-python/shared/services/retrieval/search/ranking.py b/packages/shared-python/shared/services/retrieval/search/ranking.py index 392aa1f1..c4906e46 100644 --- a/packages/shared-python/shared/services/retrieval/search/ranking.py +++ b/packages/shared-python/shared/services/retrieval/search/ranking.py @@ -1,6 +1,7 @@ from __future__ import annotations import math +from collections.abc import Mapping from typing import Any from loguru import logger @@ -15,9 +16,9 @@ def get_candidate_key(row: dict[str, Any]) -> str: path = get_row_path(row) if path: - return f'path:{path}' - chunk_id = str(row.get('chunk_id') or '').strip() - return f'chunk:{chunk_id}' if chunk_id else '' + return f"path:{path}" + chunk_id = str(row.get("chunk_id") or "").strip() + return f"chunk:{chunk_id}" if chunk_id else "" async def load_chunk_importance_scores( @@ -26,12 +27,11 @@ async def load_chunk_importance_scores( user_id: str, namespace: str, rows: list[dict[str, Any]], + revision_pins: Mapping[str, str] | None = None, ) -> dict[str, float]: - chunk_ids = sorted({ - str(row.get('chunk_id') or '').strip() - for row in rows - if row.get('chunk_id') - }) + chunk_ids = sorted( + {str(row.get("chunk_id") or "").strip() for row in rows if row.get("chunk_id")} + ) if not chunk_ids: return {} stmt = ( @@ -43,22 +43,31 @@ async def load_chunk_importance_scores( ) .where(RetrievalHitStat.user_id == user_id) .where(RetrievalHitStat.namespace == namespace) - .where(RetrievalHitStat.hit_kind == 'chunk') + .where(RetrievalHitStat.hit_kind == "chunk") .where(RetrievalHitStat.chunk_id.in_(chunk_ids)) ) + if revision_pins is not None: + pinned_document_ids = {str(document_id) for document_id in revision_pins} + stmt = stmt.where(RetrievalHitStat.document_id.in_(pinned_document_ids)) result = await db.execute(stmt) importance_scores: dict[str, float] = {} for chunk_id, hit_count, last_hit_at, created_at in result.all(): if not chunk_id: continue - importance_scores[str(chunk_id)] = compute_importance_score(hit_count, last_hit_at, created_at) + score = compute_importance_score(hit_count, last_hit_at, created_at) + importance_scores[str(chunk_id)] = score + for row in rows: + if str(row.get("chunk_id") or "") == str(chunk_id): + document_id = str(row.get("document_id") or "") + if document_id: + importance_scores[f"{document_id}:{chunk_id}"] = score return importance_scores def apply_importance_multiplier( rows: list[dict[str, Any]], *, - raw_field: str = 'importance_raw_score', + raw_field: str = "importance_raw_score", low: float = 0.1, high: float = 2.0, ) -> None: @@ -67,7 +76,11 @@ def apply_importance_multiplier( values = sorted(float(row.get(raw_field, 0.0) or 0.0) for row in rows) item_count = len(values) - median = values[item_count // 2] if item_count % 2 else (values[item_count // 2 - 1] + values[item_count // 2]) / 2 + median = ( + values[item_count // 2] + if item_count % 2 + else (values[item_count // 2 - 1] + values[item_count // 2]) / 2 + ) q1 = values[item_count // 4] if item_count >= 4 else values[0] q3 = values[3 * item_count // 4] if item_count >= 4 else values[-1] iqr = q3 - q1 @@ -80,13 +93,13 @@ def apply_importance_multiplier( z_score = (raw_score - median) / iqr sigmoid_score = 1.0 / (1.0 + math.exp(-z_score)) multiplier = low + (high - low) * sigmoid_score - row['importance_multiplier'] = round(multiplier, 4) - row['agent_score'] = round( - float(row.get('agent_score', 0.0) or 0.0) * multiplier, + row["importance_multiplier"] = round(multiplier, 4) + row["agent_score"] = round( + float(row.get("agent_score", 0.0) or 0.0) * multiplier, 6, ) - row['discovery_score'] = round( - float(row.get('discovery_score', 0.0) or 0.0) * multiplier, + row["discovery_score"] = round( + float(row.get("discovery_score", 0.0) or 0.0) * multiplier, 6, ) @@ -107,9 +120,9 @@ def rank_candidates_by_path( if not key: continue candidate = dict(row) - candidate['discovery_score'] = float(row.get('discovery_score', 0.0) or 0.0) - candidate['agent_score'] = 0.0 - candidate.setdefault('hydrate_mode', 'chunks') + candidate["discovery_score"] = float(row.get("discovery_score", 0.0) or 0.0) + candidate["agent_score"] = 0.0 + candidate.setdefault("hydrate_mode", "chunks") merged[key] = candidate insertion_order[key] = counter counter += 1 @@ -118,25 +131,33 @@ def rank_candidates_by_path( key = get_candidate_key(row) if not key: continue - routed_agent_score = float(row.get('agent_score', 0.0) or 0.0) + routed_agent_score = float(row.get("agent_score", 0.0) or 0.0) if key not in merged: candidate = dict(row) - candidate['discovery_score'] = float(row.get('discovery_score', 0.0) or 0.0) - candidate['agent_score'] = routed_agent_score + candidate["discovery_score"] = float(row.get("discovery_score", 0.0) or 0.0) + candidate["agent_score"] = routed_agent_score merged[key] = candidate insertion_order[key] = counter counter += 1 continue candidate = merged[key] - candidate['agent_score'] = max(float(candidate.get('agent_score', 0.0) or 0.0), routed_agent_score) - if not candidate.get('source_chunk_path') and row.get('source_chunk_path'): - candidate['source_chunk_path'] = row.get('source_chunk_path') - if not candidate.get('section_path') and row.get('section_path'): - candidate['section_path'] = row.get('section_path') + candidate["agent_score"] = max( + float(candidate.get("agent_score", 0.0) or 0.0), routed_agent_score + ) + if not candidate.get("source_chunk_path") and row.get("source_chunk_path"): + candidate["source_chunk_path"] = row.get("source_chunk_path") + if not candidate.get("section_path") and row.get("section_path"): + candidate["section_path"] = row.get("section_path") for row in merged.values(): - row['importance_raw_score'] = float( - (importance_scores or {}).get(str(row.get('chunk_id') or ''), 0.0) or 0.0 + document_id = str(row.get("document_id") or "") + chunk_id = str(row.get("chunk_id") or "") + row["importance_raw_score"] = float( + (importance_scores or {}).get( + f"{document_id}:{chunk_id}", + (importance_scores or {}).get(chunk_id, 0.0), + ) + or 0.0 ) apply_importance_multiplier(list(merged.values())) @@ -145,11 +166,13 @@ def rank_candidates_by_path( fallback_rows: list[dict[str, Any]] = [] for key, row in merged.items(): - agent_score = float(row.get('agent_score', 0.0) or 0.0) - discovery_score = float(row.get('discovery_score', 0.0) or 0.0) - row['evidence_score'] = round(agent_score if has_agent_results else max(discovery_score, agent_score), 6) - row['score'] = row['evidence_score'] - row['_candidate_order'] = insertion_order[key] + agent_score = float(row.get("agent_score", 0.0) or 0.0) + discovery_score = float(row.get("discovery_score", 0.0) or 0.0) + row["evidence_score"] = round( + agent_score if has_agent_results else max(discovery_score, agent_score), 6 + ) + row["score"] = row["evidence_score"] + row["_candidate_order"] = insertion_order[key] if has_agent_results and agent_score <= 0.0: fallback_rows.append(row) @@ -158,9 +181,9 @@ def rank_candidates_by_path( def get_sort_key(row: dict[str, Any]) -> tuple[float, float, int]: return ( - float(row.get('agent_score', 0.0) or 0.0), - float(row.get('discovery_score', 0.0) or 0.0), - -int(row.get('_candidate_order', 0) or 0), + float(row.get("agent_score", 0.0) or 0.0), + float(row.get("discovery_score", 0.0) or 0.0), + -int(row.get("_candidate_order", 0) or 0), ) primary_rows.sort(key=get_sort_key, reverse=True) @@ -168,10 +191,10 @@ def get_sort_key(row: dict[str, Any]) -> tuple[float, float, int]: if len(ranked_rows) < top_k and fallback_rows: fallback_rows.sort(key=get_sort_key, reverse=True) - ranked_rows.extend(fallback_rows[:top_k - len(ranked_rows)]) + ranked_rows.extend(fallback_rows[: top_k - len(ranked_rows)]) for row in ranked_rows: - row.pop('_candidate_order', None) + row.pop("_candidate_order", None) return ranked_rows @@ -183,6 +206,7 @@ async def rank_retrieval_candidates( discovery_rows: list[dict[str, Any]], routed_rows: list[dict[str, Any]], top_k: int, + revision_pins: Mapping[str, str] | None = None, ) -> list[dict[str, Any]]: try: importance_scores = await load_chunk_importance_scores( @@ -190,9 +214,12 @@ async def rank_retrieval_candidates( user_id=user_id, namespace=namespace, rows=[*discovery_rows, *routed_rows], + revision_pins=revision_pins, ) except Exception as exc: - logger.warning(f'Failed to load chunk importance scores, continuing without importance: {exc}') + logger.warning( + f"Failed to load chunk importance scores, continuing without importance: {exc}" + ) importance_scores = {} return rank_candidates_by_path( discovery_rows, diff --git a/packages/shared-python/shared/services/retrieval/search/scoped_corpus.py b/packages/shared-python/shared/services/retrieval/search/scoped_corpus.py index 4a60e65b..ed6f2719 100644 --- a/packages/shared-python/shared/services/retrieval/search/scoped_corpus.py +++ b/packages/shared-python/shared/services/retrieval/search/scoped_corpus.py @@ -1,13 +1,72 @@ from __future__ import annotations +from collections.abc import Mapping from typing import Any -from sqlalchemy import func, select +from sqlalchemy import and_, func, select, tuple_ from sqlalchemy.ext.asyncio import AsyncSession +from sqlalchemy.exc import SQLAlchemyError -from shared.models.database.document import Document, DocumentChunk, DocumentSection +from shared.models.database.document import ( + Document, + DocumentChunk, + DocumentSection, + RetrievalServingRevisionManifest, +) from shared.models.database.job_result import JobResult from shared.services.retrieval.search.section_filters import is_excluded_section +from shared.services.retrieval.serving_manifest import decode_serving_manifest +from shared.services.retrieval.manifest_cache import cache_manifest_payloads + + +async def count_manifest_chunks( + db: AsyncSession, + *, + revision_pins: Mapping[str, str], +) -> int | None: + """Count chunks exactly from complete pinned manifests when available.""" + if not revision_pins: + return None + statement = select( + RetrievalServingRevisionManifest.document_id, + RetrievalServingRevisionManifest.job_result_id, + RetrievalServingRevisionManifest.payload_zlib, + RetrievalServingRevisionManifest.checksum, + RetrievalServingRevisionManifest.format_version, + ).where( + tuple_( + RetrievalServingRevisionManifest.document_id, + RetrievalServingRevisionManifest.job_result_id, + ).in_(list(revision_pins.items())) + ) + try: + rows = (await db.execute(statement)).all() + if len(rows) != len(revision_pins): + return None + total = 0 + decoded_payloads: dict[tuple[str, str], dict[str, Any]] = {} + for document_id, job_result_id, payload_zlib, checksum, format_version in rows: + payload = decode_serving_manifest( + bytes(payload_zlib), + checksum=str(checksum), + format_version=int(format_version), + ) + chunks = payload.get("chunks") + if not isinstance(chunks, list): + return None + total += len(chunks) + decoded_payloads[(str(document_id), str(job_result_id))] = payload + cache_manifest_payloads( + db, + revisions=revision_pins, + payloads=decoded_payloads, + ) + return total + except SQLAlchemyError: + await db.rollback() + return None + except (TypeError, ValueError, KeyError): + return None async def count_scoped_chunks( @@ -17,18 +76,32 @@ async def count_scoped_chunks( namespace: str, exclude_document_ids: list[str], allowed_chunk_types: set[str] | None, + revision_pins: Mapping[str, str] | None = None, ) -> int: - stmt = ( - select(func.count(DocumentChunk.id)) - .join( - Document, - (Document.document_id == DocumentChunk.document_id) - & (Document.current_job_result_id == DocumentChunk.job_result_id), + if revision_pins is None: + stmt = ( + select(func.count(DocumentChunk.id)) + .join( + Document, + (Document.document_id == DocumentChunk.document_id) + & (Document.current_job_result_id == DocumentChunk.job_result_id), + ) + .where(Document.user_id == user_id) + .where(Document.namespace == namespace) + .where(Document.status == 'active') + ) + else: + stmt = ( + select(func.count(DocumentChunk.id)) + .join(Document, Document.document_id == DocumentChunk.document_id) + .where(Document.user_id == user_id) + .where(Document.namespace == namespace) + .where( + tuple_(DocumentChunk.document_id, DocumentChunk.job_result_id).in_( + list(revision_pins.items()) + ) + ) ) - .where(Document.user_id == user_id) - .where(Document.namespace == namespace) - .where(Document.status == 'active') - ) if exclude_document_ids: stmt = stmt.where(Document.document_id.notin_(list(exclude_document_ids))) if allowed_chunk_types is not None: @@ -47,21 +120,31 @@ async def load_all_scoped_chunks( allowed_chunk_types: set[str] | None, signal_paths: list[str], filter_mode: str, + revision_pins: Mapping[str, str] | None = None, ) -> list[dict[str, Any]]: - stmt = ( - select(Document, DocumentChunk, DocumentSection, JobResult) - .join( - DocumentChunk, + if revision_pins is None: + chunk_join = ( (DocumentChunk.document_id == Document.document_id) - & (DocumentChunk.job_result_id == Document.current_job_result_id), + & (DocumentChunk.job_result_id == Document.current_job_result_id) + ) + else: + chunk_join = and_( + DocumentChunk.document_id == Document.document_id, + tuple_(DocumentChunk.document_id, DocumentChunk.job_result_id).in_( + list(revision_pins.items()) + ), ) + stmt = ( + select(Document, DocumentChunk, DocumentSection, JobResult) + .join(DocumentChunk, chunk_join) .outerjoin(DocumentSection, DocumentSection.section_id == DocumentChunk.section_id) .join(JobResult, JobResult.id == DocumentChunk.job_result_id) .where(Document.user_id == user_id) .where(Document.namespace == namespace) - .where(Document.status == 'active') .order_by(DocumentChunk.sort_order) ) + if revision_pins is None: + stmt = stmt.where(Document.status == 'active') if exclude_document_ids: stmt = stmt.where(Document.document_id.notin_(list(exclude_document_ids))) if allowed_chunk_types is not None: diff --git a/packages/shared-python/shared/services/retrieval/search/scoring.py b/packages/shared-python/shared/services/retrieval/search/scoring.py index 848a3ada..2746de34 100644 --- a/packages/shared-python/shared/services/retrieval/search/scoring.py +++ b/packages/shared-python/shared/services/retrieval/search/scoring.py @@ -51,14 +51,18 @@ def merge_channels_rrf( for channel_idx, channel_rows in enumerate(channels): weight = weights[channel_idx] if channel_idx < len(weights) else 1.0 - for rank, row in enumerate(channel_rows): + seen_chunk_ids: set[str] = set() + unique_rank = 0 + for row in channel_rows: chunk_id = str(row.get('chunk_id') or '') - if not chunk_id: + if not chunk_id or chunk_id in seen_chunk_ids: continue - rrf_score = weight / (k + rank + 1) + seen_chunk_ids.add(chunk_id) + rrf_score = weight / (k + unique_rank + 1) score_dict[chunk_id] = score_dict.get(chunk_id, 0.0) + rrf_score if chunk_id not in row_by_chunk_id: row_by_chunk_id[chunk_id] = row + unique_rank += 1 ranked = sorted(score_dict.items(), key=lambda x: x[1], reverse=True) results: list[dict[str, Any]] = [] diff --git a/packages/shared-python/shared/services/retrieval/serving_generation.py b/packages/shared-python/shared/services/retrieval/serving_generation.py new file mode 100644 index 00000000..5dc7baf5 --- /dev/null +++ b/packages/shared-python/shared/services/retrieval/serving_generation.py @@ -0,0 +1,60 @@ +"""Namespace generation locking for serving-state lifecycle updates.""" + +from __future__ import annotations + +from hashlib import sha256 + +from sqlalchemy import select +from sqlalchemy.dialects.postgresql import insert +from sqlalchemy.orm import Session + +from shared.models.database.document import RetrievalNamespaceGeneration + + +def lock_namespace_generation( + db: Session, + *, + user_id: str, + namespace: str, +) -> RetrievalNamespaceGeneration: + """Create if needed, then lock and return one namespace generation row.""" + generation_id = f"rng_{sha256(f'{user_id}:{namespace}'.encode()).hexdigest()}" + db.execute( + insert(RetrievalNamespaceGeneration) + .values( + id=generation_id, + user_id=user_id, + namespace=namespace, + generation=0, + ) + .on_conflict_do_nothing( + index_elements=[ + RetrievalNamespaceGeneration.user_id, + RetrievalNamespaceGeneration.namespace, + ] + ) + ) + generation = db.execute( + select(RetrievalNamespaceGeneration) + .where(RetrievalNamespaceGeneration.user_id == user_id) + .where(RetrievalNamespaceGeneration.namespace == namespace) + .with_for_update() + ).scalar_one() + return generation + + +def advance_namespace_generation( + db: Session, + *, + user_id: str, + namespace: str, +) -> int: + """Increment a locked namespace generation and return its new value.""" + generation = lock_namespace_generation( + db, + user_id=user_id, + namespace=namespace, + ) + generation.generation += 1 + db.flush() + return generation.generation diff --git a/packages/shared-python/shared/services/retrieval/serving_manifest.py b/packages/shared-python/shared/services/retrieval/serving_manifest.py new file mode 100644 index 00000000..b2dd40c5 --- /dev/null +++ b/packages/shared-python/shared/services/retrieval/serving_manifest.py @@ -0,0 +1,405 @@ +"""Versioned compression and integrity checks for serving manifests.""" + +from __future__ import annotations + +import hashlib +import json +import zlib +from typing import Any + +from sqlalchemy import delete, select +from sqlalchemy.orm import Session + +from shared.models.database.document import ( + Document, + DocumentChunk, + DocumentMapUnit, + DocumentMapUnitToken, + DocumentSection, + RetrievalNamespaceGeneration, + RetrievalNamespaceStat, + RetrievalNamespaceTokenStat, + RetrievalServingRevisionManifest, + RetrievalServingRevisionStat, +) +from shared.models.database.job_result import JobResult +from shared.services.retrieval.publication_models import DocumentPublicationScope + +SERVING_MANIFEST_FORMAT_VERSION = 1 + + +def build_revision_serving_payload( + db: Session, + *, + scope: DocumentPublicationScope, +) -> dict[str, Any]: + """Build ordered metadata for one published document revision.""" + document = db.execute( + select(Document).where(Document.document_id == scope.document_id) + ).scalar_one() + job_result = db.execute( + select(JobResult).where(JobResult.id == scope.job_result_id) + ).scalar_one() + sections = 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) + ) + ) + chunks = 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 + ) + ) + ) + map_units = list( + db.scalars( + select(DocumentMapUnit) + .where(DocumentMapUnit.document_id == scope.document_id) + .where(DocumentMapUnit.job_result_id == scope.job_result_id) + .order_by(DocumentMapUnit.sort_order, DocumentMapUnit.unit_id) + ) + ) + section_path_by_id = { + section.section_id: section.section_path for section in sections + } + root_asset_ids = { + chunk.chunk_id + for chunk in chunks + if chunk.chunk_type in {"image", "table"} + and chunk.section_id is not None + and section_path_by_id.get(chunk.section_id) == "Root" + } + remounted_assets: dict[str, list[str]] = {} + for chunk in chunks: + if chunk.chunk_type != "text" or not isinstance(chunk.chunk_metadata, dict): + continue + connections = chunk.chunk_metadata.get("connect_to") + if not isinstance(connections, list): + continue + targets = [ + str(connection.get("target") or "").strip() + for connection in connections + if isinstance(connection, dict) + and str(connection.get("target") or "").strip() in root_asset_ids + ] + if targets: + remounted_assets[chunk.section_id or ""] = targets + + return { + "document_id": scope.document_id, + "job_result_id": scope.job_result_id, + "job_id": str(job_result.job_id), + "source_file_name": str( + document.source_file_name or scope.source_file_name or "" + ), + "sections": [ + { + "section_id": section.section_id, + "parent_section_id": section.parent_section_id, + "section_path": section.section_path, + "section_title": section.section_title, + "section_level": section.section_level, + "summary": section.summary, + "sort_order": section.sort_order, + } + for section in sections + ], + "chunks": [ + { + "chunk_id": chunk.chunk_id, + "section_id": chunk.section_id, + "chunk_type": chunk.chunk_type, + "sort_order": chunk.sort_order, + "connect_to": _connection_target_ids(chunk.chunk_metadata), + } + for chunk in chunks + ], + "map_units": [ + { + "row_id": unit.id, + "unit_id": unit.unit_id, + "section_id": unit.section_id, + "unit_kind": unit.unit_kind, + "path_token_count": unit.path_token_count, + "content_token_count": unit.content_token_count, + "sort_order": unit.sort_order, + } + for unit in map_units + ], + "root_asset_ids": sorted(root_asset_ids), + "remounted_assets_by_section": remounted_assets, + } + + +def build_revision_statistics_payload( + db: Session, + *, + scope: DocumentPublicationScope, +) -> dict[str, Any]: + """Build compressed scoring contributions for one revision.""" + units = list( + db.scalars( + select(DocumentMapUnit) + .where(DocumentMapUnit.document_id == scope.document_id) + .where(DocumentMapUnit.job_result_id == scope.job_result_id) + ) + ) + unit_ids = [unit.id for unit in units] + frequencies: dict[str, dict[str, int]] = {"path": {}, "content": {}} + unit_frequencies: dict[str, dict[str, dict[str, int]]] = {} + if unit_ids: + for map_unit_id, channel, token, frequency in db.execute( + select( + DocumentMapUnitToken.map_unit_id, + DocumentMapUnitToken.channel, + DocumentMapUnitToken.token, + DocumentMapUnitToken.frequency, + ).where(DocumentMapUnitToken.map_unit_id.in_(unit_ids)) + ).all(): + channel_key = str(channel) + if channel_key in frequencies: + token_key = str(token) + frequency_value = int(frequency) + frequencies[channel_key][token_key] = ( + frequencies[channel_key].get(token_key, 0) + frequency_value + ) + unit_frequencies.setdefault(str(map_unit_id), {}).setdefault( + channel_key, {} + )[token_key] = frequency_value + return { + "document_id": scope.document_id, + "job_result_id": scope.job_result_id, + "unit_count": len(units), + "path_token_count": sum(int(unit.path_token_count or 0) for unit in units), + "content_token_count": sum( + int(unit.content_token_count or 0) for unit in units + ), + "token_frequencies": frequencies, + "unit_frequencies": unit_frequencies, + } + + +def persist_revision_serving_state( + db: Session, + *, + scope: DocumentPublicationScope, +) -> None: + """Replace manifest and statistics rows for one revision atomically.""" + manifest_payload = build_revision_serving_payload(db, scope=scope) + statistics_payload = build_revision_statistics_payload(db, scope=scope) + manifest_bytes, manifest_checksum, manifest_version = encode_serving_manifest( + manifest_payload + ) + statistics_bytes, statistics_checksum, statistics_version = encode_serving_manifest( + statistics_payload + ) + db.execute( + delete(RetrievalServingRevisionManifest) + .where(RetrievalServingRevisionManifest.document_id == scope.document_id) + .where(RetrievalServingRevisionManifest.job_result_id == scope.job_result_id) + ) + db.execute( + delete(RetrievalServingRevisionStat) + .where(RetrievalServingRevisionStat.document_id == scope.document_id) + .where(RetrievalServingRevisionStat.job_result_id == scope.job_result_id) + ) + db.add( + RetrievalServingRevisionManifest( + user_id=scope.user_id, + namespace=scope.namespace, + document_id=scope.document_id, + job_result_id=scope.job_result_id, + format_version=manifest_version, + payload_zlib=manifest_bytes, + checksum=manifest_checksum, + ) + ) + db.add( + RetrievalServingRevisionStat( + user_id=scope.user_id, + namespace=scope.namespace, + document_id=scope.document_id, + job_result_id=scope.job_result_id, + format_version=statistics_version, + payload_zlib=statistics_bytes, + checksum=statistics_checksum, + ) + ) + + +def rebuild_namespace_serving_statistics( + db: Session, + *, + user_id: str, + namespace: str, +) -> int: + """Recompute namespace aggregates from active current revisions. + + Callers hold the namespace generation lock. The aggregate is prepared for + the generation that the caller will publish next. + """ + generation = db.execute( + select(RetrievalNamespaceGeneration) + .where(RetrievalNamespaceGeneration.user_id == user_id) + .where(RetrievalNamespaceGeneration.namespace == namespace) + .with_for_update() + ).scalar_one() + target_generation = int(generation.generation) + 1 + revisions = { + (str(document_id), str(job_result_id)) + for document_id, job_result_id in db.execute( + select(Document.document_id, Document.current_job_result_id) + .where(Document.user_id == user_id) + .where(Document.namespace == namespace) + .where(Document.status == "active") + .where(Document.current_job_result_id.is_not(None)) + ).all() + if document_id and job_result_id + } + aggregate: dict[str, Any] = { + "document_count": 0, + "unit_count": 0, + "path_token_count": 0, + "content_token_count": 0, + "token_frequencies": {"path": {}, "content": {}}, + } + document_frequencies: dict[tuple[str, str], int] = {} + for row in db.scalars( + select(RetrievalServingRevisionStat) + .where(RetrievalServingRevisionStat.user_id == user_id) + .where(RetrievalServingRevisionStat.namespace == namespace) + ): + if (row.document_id, row.job_result_id) not in revisions: + continue + payload = decode_serving_manifest( + row.payload_zlib, + checksum=row.checksum, + format_version=row.format_version, + ) + aggregate["document_count"] += 1 + aggregate["unit_count"] += int(payload.get("unit_count", 0)) + aggregate["path_token_count"] += int(payload.get("path_token_count", 0)) + aggregate["content_token_count"] += int(payload.get("content_token_count", 0)) + token_frequencies = payload.get("token_frequencies", {}) + if not isinstance(token_frequencies, dict): + continue + for channel, values in token_frequencies.items(): + if channel not in aggregate["token_frequencies"] or not isinstance( + values, dict + ): + continue + for token, value in values.items(): + token_key = str(token) + aggregate["token_frequencies"][channel][token_key] = aggregate[ + "token_frequencies" + ][channel].get(token_key, 0) + int(value) + if int(value) > 0: + key = (str(channel), token_key) + document_frequencies[key] = document_frequencies.get(key, 0) + 1 + + encoded, checksum, _version = encode_serving_manifest(aggregate) + namespace_stat = db.execute( + select(RetrievalNamespaceStat) + .where(RetrievalNamespaceStat.user_id == user_id) + .where(RetrievalNamespaceStat.namespace == namespace) + ).scalar_one_or_none() + if namespace_stat is None: + db.add( + RetrievalNamespaceStat( + user_id=user_id, + namespace=namespace, + generation=target_generation, + payload_zlib=encoded, + checksum=checksum, + ) + ) + else: + namespace_stat.generation = target_generation + namespace_stat.payload_zlib = encoded + namespace_stat.checksum = checksum + db.execute( + delete(RetrievalNamespaceTokenStat) + .where(RetrievalNamespaceTokenStat.user_id == user_id) + .where(RetrievalNamespaceTokenStat.namespace == namespace) + ) + db.add_all( + [ + RetrievalNamespaceTokenStat( + user_id=user_id, + namespace=namespace, + generation=target_generation, + channel=channel, + token_hash=hashlib.sha256(token.encode("utf-8")).hexdigest(), + document_frequency=frequency, + ) + for (channel, token), frequency in document_frequencies.items() + ] + ) + db.flush() + return target_generation + + +def _connection_target_ids(metadata: Any) -> list[str]: + if not isinstance(metadata, dict): + return [] + connections = metadata.get("connect_to") + if not isinstance(connections, list): + return [] + return [ + target + for connection in connections + if isinstance(connection, dict) + for target in [str(connection.get("target") or "").strip()] + if target + ] + + +def encode_serving_manifest(payload: dict[str, Any]) -> tuple[bytes, str, int]: + """Return compressed canonical JSON, checksum, and format version.""" + canonical_payload = json.dumps( + payload, + ensure_ascii=False, + separators=(",", ":"), + sort_keys=True, + ).encode("utf-8") + checksum = hashlib.sha256(canonical_payload).hexdigest() + return ( + zlib.compress(canonical_payload), + checksum, + SERVING_MANIFEST_FORMAT_VERSION, + ) + + +def decode_serving_manifest( + payload_zlib: bytes, + *, + checksum: str, + format_version: int, +) -> dict[str, Any]: + """Validate and decode one persisted serving manifest.""" + if format_version != SERVING_MANIFEST_FORMAT_VERSION: + raise ValueError(f"unsupported serving manifest version: {format_version}") + + try: + canonical_payload = zlib.decompress(payload_zlib) + except zlib.error as exc: + raise ValueError("invalid serving manifest compression") from exc + + actual_checksum = hashlib.sha256(canonical_payload).hexdigest() + if actual_checksum != checksum: + raise ValueError("serving manifest checksum mismatch") + + try: + decoded = json.loads(canonical_payload.decode("utf-8")) + except (UnicodeDecodeError, json.JSONDecodeError) as exc: + raise ValueError("invalid serving manifest JSON") from exc + if not isinstance(decoded, dict): + raise ValueError("serving manifest payload must be an object") + return decoded From dabb25d10da4c47be6d9e06713c2d1cb1bafe93e Mon Sep 17 00:00:00 2001 From: suguanYang Date: Sun, 30 Aug 2026 12:58:01 +0800 Subject: [PATCH 5/5] test: patch active database module in revision race contract --- .../contract/test_retrieval_revision_races_contract.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/apps/api/tests/contract/test_retrieval_revision_races_contract.py b/apps/api/tests/contract/test_retrieval_revision_races_contract.py index dd73d966..3ebb0232 100644 --- a/apps/api/tests/contract/test_retrieval_revision_races_contract.py +++ b/apps/api/tests/contract/test_retrieval_revision_races_contract.py @@ -2,6 +2,7 @@ from __future__ import annotations +import importlib from collections.abc import AsyncGenerator from contextlib import asynccontextmanager from typing import Any, cast @@ -40,7 +41,12 @@ async def fake_channel( observed.append((db, kwargs["revision_pins"])) return [] - monkeypatch.setattr("shared.core.database.get_db_context", fake_context) + # Import the active module explicitly. Some API contract fixtures reload + # shared.core.database between tests, leaving the package attribute pointed + # at an old module object; dotted-string patching can then miss the module + # used by discovery's lazy import. + database_module = importlib.import_module("shared.core.database") + monkeypatch.setattr(database_module, "get_db_context", fake_context) monkeypatch.setattr(discovery, "path_channel", fake_channel) monkeypatch.setattr(discovery, "content_channel", fake_channel) monkeypatch.setattr(discovery, "term_channel", fake_channel)