From 9b54d83c2071625fbe4b8f3223e04eb270236683 Mon Sep 17 00:00:00 2001 From: suguanYang Date: Tue, 1 Sep 2026 05:56:03 +0800 Subject: [PATCH 1/5] fix: harden retrieval serving consistency and fallback --- apps/api/app/api/v1/routes/retrieval.py | 27 +- .../tests/contract/test_retrieval_contract.py | 22 ++ ...etrieval_lazy_snapshot_quality_contract.py | 15 +- .../test_retrieval_map_unit_index_contract.py | 8 +- .../services/retrieval/cache_service.py | 32 +++ .../services/retrieval/nav/nav_knowhere.py | 52 +++- .../services/retrieval/nav/nav_map_scores.py | 84 ++++++ .../retrieval/nav/persisted_score_load.py | 23 ++ .../services/retrieval/publication_service.py | 8 + .../retrieval/search/map_unit_discovery.py | 259 +++++++++++++++++- .../services/retrieval/serving_manifest.py | 15 + 11 files changed, 515 insertions(+), 30 deletions(-) diff --git a/apps/api/app/api/v1/routes/retrieval.py b/apps/api/app/api/v1/routes/retrieval.py index 3c9202dfa..c7d795956 100644 --- a/apps/api/app/api/v1/routes/retrieval.py +++ b/apps/api/app/api/v1/routes/retrieval.py @@ -7,7 +7,7 @@ from app.api.dependencies.current_user import with_current_user from app.services.rate_limit.data_structures import CurrentUser from fastapi import APIRouter, Depends -from pydantic import BaseModel, Field, field_validator +from pydantic import BaseModel, Field, field_validator, model_validator from sqlalchemy.ext.asyncio import AsyncSession from shared.core.database import get_db @@ -60,10 +60,17 @@ class RetrievalQueryRequest(BaseModel): ) channels: list[str] = Field( default_factory=list, - description="Channels to run (empty=all). Options: path, content, term", + description=( + "Deprecated and unsupported by the persisted map-unit route. " + "Leave empty; explicit channel selection is rejected." + ), ) channel_weights: dict[str, float] = Field( - default_factory=dict, description="Per-channel weight overrides" + default_factory=dict, + description=( + "Deprecated and unsupported by the persisted map-unit route. " + "Leave empty; explicit overrides are rejected." + ), ) rerank: bool = Field(False, description="Enable LLM reranking after RRF fusion") threshold: float = Field(0.0, ge=0.0, description="Minimum RRF score threshold") @@ -112,6 +119,20 @@ def validate_chunk_types(cls, v: list[str] | None) -> list[str] | None: def normalize_namespace(cls, namespace: str | None) -> str: return normalize_retrieval_namespace(namespace) + @model_validator(mode="after") + def reject_unsupported_channel_controls(self) -> "RetrievalQueryRequest": + if self.channels: + raise ValueError( + "channels is deprecated and unsupported; omit it and use the " + "persisted path/content map-unit scorer" + ) + if self.channel_weights: + raise ValueError( + "channel_weights is deprecated and unsupported; omit it and use " + "the persisted path/content map-unit scorer" + ) + return self + class RetrievalQueryResponse(BaseModel): namespace: str diff --git a/apps/api/tests/contract/test_retrieval_contract.py b/apps/api/tests/contract/test_retrieval_contract.py index 78dffed98..90dc7e4f6 100644 --- a/apps/api/tests/contract/test_retrieval_contract.py +++ b/apps/api/tests/contract/test_retrieval_contract.py @@ -306,6 +306,28 @@ async def test_should_return_request_validation_failure_for_an_invalid_channel( assert "Invalid channel" in cast(str, violations[0]["description"]) +async def test_should_reject_legacy_channel_controls( + developer_api_client_factory: Callable[ + [], AbstractAsyncContextManager[AsyncClient] + ], +) -> None: + async with developer_api_client_factory() as api_client: + response = await api_client.post( + "/api/v1/retrieval/query", + json={ + "namespace": "default", + "query": "alpha", + "channels": ["content"], + }, + ) + + assert response.status_code == 400 + response_json = cast(dict[str, object], response.json()) + error = cast(dict[str, object], response_json["error"]) + assert error["code"] == "INVALID_ARGUMENT" + assert "deprecated and unsupported" in str(error) + + async def test_should_exclude_matching_document_ids_from_the_response( developer_api_client_factory: Callable[ [], AbstractAsyncContextManager[AsyncClient] 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 b80fb8aee..c796e2a18 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 @@ -165,10 +165,7 @@ def test_lazy_provider_preserves_score_units_and_scores() -> None: lazy_scores = compute_corpus_map_and_unit_scores( lazy, doc_ids=["doc"], query="alpha retrieval" ) - assert eager_scores[1] == {} - assert lazy_scores[1] == {} - assert all(score == 0.0 for score in eager_scores[0].values()) - assert all(score == 0.0 for score in lazy_scores[0].values()) + assert eager_scores == lazy_scores lazy_provider = lazy._provider self_units = getattr(lazy_provider, "self_units") @@ -245,7 +242,7 @@ def build_stats(search_field: str) -> PersistedBm25Stats: assert scored["beta evidence"]["unit-b"] > scored["beta evidence"]["unit-a"] -def test_missing_index_does_not_read_chunk_payloads() -> None: +def test_missing_index_falls_back_to_legacy_payload_scoring() -> None: _eager, lazy, store = _providers() queries: list[str] = ["alpha retrieval", "supporting image"] store.section_loads = 0 @@ -256,8 +253,8 @@ def test_missing_index_does_not_read_chunk_payloads() -> None: ) assert set(actual) == set(queries) - assert all(unit_scores == {} for _map_scores, unit_scores in actual.values()) - assert store.section_loads == 0 + assert any(unit_scores for _map_scores, unit_scores in actual.values()) + assert store.section_loads > 0 def test_missing_index_is_empty_across_documents() -> None: @@ -276,8 +273,8 @@ def test_missing_index_is_empty_across_documents() -> None: ) assert actual == expected - assert all(unit_scores == {} for _map_scores, unit_scores in actual.values()) - assert store.section_loads == 0 + assert any(unit_scores for _map_scores, unit_scores in actual.values()) + assert store.section_loads > 0 def test_native_chunk_store_strips_async_driver_from_database_url( 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 5c653722f..e95efd233 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 @@ -276,8 +276,8 @@ def reject_payload_read( assert any(score > 0.0 for score in actual_scores[1].values()) assert select_map_highlights(actual_scores[1], k=3) - assert fallback_scores[1] == {} - assert all(score == 0.0 for score in fallback_scores[0].values()) + assert any(score > 0.0 for score in fallback_scores[1].values()) + assert any(score > 0.0 for score in fallback_scores[0].values()) async def test_lazy_snapshot_defers_selected_asset_reference_metadata( @@ -476,8 +476,8 @@ def test_incomplete_index_returns_empty_scores() -> None: lazy, doc_ids=["doc-a", "doc-b"], query="alpha beta" ) - assert actual[1] == {} - assert expected[1] == {} + assert set(actual[1]) == {"leaf-a", "leaf-b"} + assert set(expected[1]) == {"leaf-a", "leaf-b"} assert all(score == 0.0 for score in actual[0].values()) assert store.persisted_loads == 1 diff --git a/packages/shared-python/shared/services/retrieval/cache_service.py b/packages/shared-python/shared/services/retrieval/cache_service.py index da76068d1..c327eb468 100644 --- a/packages/shared-python/shared/services/retrieval/cache_service.py +++ b/packages/shared-python/shared/services/retrieval/cache_service.py @@ -10,6 +10,7 @@ _RETRIEVAL_CACHE_TTL_SECONDS = 300 _VERSION_FALLBACK = 0 +_INDEX_READINESS_TTL_SECONDS = 60 def _namespace_version_key(*, user_id: str, namespace: str) -> str: @@ -17,6 +18,37 @@ def _namespace_version_key(*, user_id: str, namespace: str) -> str: return f"retrieval:version:{user_id}:{namespace}" +def _namespace_index_readiness_key(*, user_id: str, namespace: str) -> str: + namespace = normalize_retrieval_namespace(namespace) + return f"retrieval:index-readiness:{user_id}:{namespace}" + + +async def record_retrieval_index_readiness( + *, + user_id: str, + namespace: str, + ready: bool, + expected_revisions: int, + indexed_revisions: int, +) -> None: + """Publish a short-lived index readiness signal for operators and callers. + + Redis is deliberately only a status cache. PostgreSQL generations and + serving rows remain the source of truth, and retrieval must continue to + work if Redis is unavailable. + """ + redis_service = RedisServiceFactory.get_service() + await redis_service.set( + _namespace_index_readiness_key(user_id=user_id, namespace=namespace), + { + "ready": bool(ready), + "expected_revisions": int(expected_revisions), + "indexed_revisions": int(indexed_revisions), + }, + ex=_INDEX_READINESS_TTL_SECONDS, + ) + + def _normalize_exclude_sections(exclude_sections: list[dict[str, str]]) -> list[str]: normalized: list[str] = [] for item in exclude_sections: 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 716ff1ec7..dfbad8717 100644 --- a/packages/shared-python/shared/services/retrieval/nav/nav_knowhere.py +++ b/packages/shared-python/shared/services/retrieval/nav/nav_knowhere.py @@ -300,6 +300,7 @@ def load_persisted_score_corpus( comes from ``document_map_unit_indexes`` (written at index time). """ from shared.services.retrieval.nav.persisted_score_load import ( + average_idf_from_namespace_stats, build_channel_bm25_stats, combine_average_idf, ) @@ -359,12 +360,51 @@ def load_persisted_score_corpus( revision_key ] else: - average_idf_path = combine_average_idf( - [(float(row[4] or 0.0), int(row[3] or 0)) for row in index_rows] - ) - average_idf_content = combine_average_idf( - [(float(row[5] or 0.0), int(row[3] or 0)) for row in index_rows] - ) + total_unit_count = sum(int(row[3] or 0) for row in index_rows) + try: + cur.execute( + "SELECT tokens.channel, tokens.token, " + "COUNT(DISTINCT tokens.map_unit_id) " + "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 " + "WHERE tokens.channel = ANY(%s) " + "GROUP BY tokens.channel, tokens.token", + [*revision_params, ["path", "content"]], + ) + namespace_token_dfs: dict[str, list[int]] = { + "path": [], + "content": [], + } + for channel, _token, document_frequency in cur.fetchall(): + if str(channel) in namespace_token_dfs: + namespace_token_dfs[str(channel)].append( + int(document_frequency) + ) + average_idf_path = average_idf_from_namespace_stats( + unit_count=total_unit_count, + token_document_frequencies=namespace_token_dfs["path"], + ) + average_idf_content = average_idf_from_namespace_stats( + unit_count=total_unit_count, + token_document_frequencies=namespace_token_dfs["content"], + ) + except Exception as exc: + _logger.warning( + "exact namespace IDF load failed; using revision averages: %s", + exc, + ) + average_idf_path = combine_average_idf( + [(float(row[4] or 0.0), int(row[3] or 0)) for row in index_rows] + ) + average_idf_content = combine_average_idf( + [ + (float(row[5] or 0.0), int(row[3] or 0)) + for row in index_rows + ] + ) self._score_average_idf_cache[revision_key] = ( average_idf_path, average_idf_content, 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 c646070ae..2458661b8 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 @@ -8,12 +8,89 @@ build_content_search_text, build_path_search_text, build_term_search_text, + PersistedScoreCorpus, + PersistedScoreUnit, score_persisted_corpus_many, ) +from .persisted_score_load import ( + average_idf_from_unit_dfs, + build_channel_bm25_stats, +) _logger = logging.getLogger(__name__) +def _build_legacy_score_corpus(ts: Any, doc_ids: Sequence[str]) -> PersistedScoreCorpus: + """Build the retired in-memory scorer input when persisted indexes are absent.""" + raw_units: List[dict] = [] + for doc_id in doc_ids: + raw_units.extend(build_score_units(ts, doc_id)) + frequencies: Dict[Tuple[str, str], Dict[str, int]] = {} + unit_rows: List[dict] = [] + path_dfs: Dict[str, int] = {} + content_dfs: Dict[str, int] = {} + for unit in raw_units: + unit_id = str(unit.get("chunk_id") or "").strip() + if not unit_id: + continue + path_tokens = str(unit.get("path_search_text") or "").split() + content_tokens = str(unit.get("content_search_text") or "").split() + path_freq: Dict[str, int] = {} + content_freq: Dict[str, int] = {} + for token in path_tokens: + path_freq[token] = path_freq.get(token, 0) + 1 + for token in content_tokens: + content_freq[token] = content_freq.get(token, 0) + 1 + frequencies[(unit_id, "path")] = path_freq + frequencies[(unit_id, "content")] = content_freq + for token in path_freq: + path_dfs[token] = path_dfs.get(token, 0) + 1 + for token in content_freq: + content_dfs[token] = content_dfs.get(token, 0) + 1 + unit_rows.append( + { + "unit_id": unit_id, + "path_length": len(path_tokens), + "content_length": len(content_tokens), + } + ) + unit_count = len(unit_rows) + return PersistedScoreCorpus( + units=[ + PersistedScoreUnit( + unit_id=str(row["unit_id"]), + path_length=int(row["path_length"]), + content_length=int(row["content_length"]), + path_frequencies=frequencies[(str(row["unit_id"]), "path")], + content_frequencies=frequencies[(str(row["unit_id"]), "content")], + ) + for row in unit_rows + ], + path_stats=build_channel_bm25_stats( + unit_rows=unit_rows, + map_unit_id_field="unit_id", + length_field="path_length", + channel="path", + query_tokens=list(path_dfs), + frequencies=frequencies, + average_idf=average_idf_from_unit_dfs( + unit_count=unit_count, token_document_frequency=path_dfs + ), + ), + content_stats=build_channel_bm25_stats( + unit_rows=unit_rows, + map_unit_id_field="unit_id", + length_field="content_length", + channel="content", + query_tokens=list(content_dfs), + frequencies=frequencies, + average_idf=average_idf_from_unit_dfs( + unit_count=unit_count, token_document_frequency=content_dfs + ), + ), + ) + + def _children_ids(ts: Any, section_id: str, doc_id: str) -> List[str]: children_fn = getattr(ts, "_children_for_section_path", None) if not callable(children_fn): @@ -346,6 +423,13 @@ def compute_corpus_map_and_unit_scores_many( time.perf_counter() - loader_started, persisted_corpus is not None, ) + if persisted_corpus is None: + _logger.warning( + "retrieval map index unavailable; using bounded legacy in-memory scorer " + "documents=%d", + len(valid_doc_ids), + ) + persisted_corpus = _build_legacy_score_corpus(ts, valid_doc_ids) score_started = time.perf_counter() unit_scores_by_query = ( score_persisted_corpus_many(persisted_corpus, unique_queries) diff --git a/packages/shared-python/shared/services/retrieval/nav/persisted_score_load.py b/packages/shared-python/shared/services/retrieval/nav/persisted_score_load.py index 9ce10c2e1..7f1eb2029 100644 --- a/packages/shared-python/shared/services/retrieval/nav/persisted_score_load.py +++ b/packages/shared-python/shared/services/retrieval/nav/persisted_score_load.py @@ -38,6 +38,29 @@ def combine_average_idf(parts: Sequence[tuple[float, int]]) -> float: ) +def average_idf_from_namespace_stats( + *, + unit_count: int, + token_document_frequencies: Sequence[int], +) -> float: + """Compute the exact namespace-level average IDF used by rank_bm25. + + Namespace token statistics already contain one document frequency per + token. Computing the mean from those rows avoids the incorrect + per-revision-average approximation when a namespace contains revisions + with different token distributions. + """ + if unit_count <= 0: + return 0.0 + idfs = [ + math.log(unit_count - int(frequency) + 0.5) + - math.log(int(frequency) + 0.5) + for frequency in token_document_frequencies + if 0 < int(frequency) <= unit_count + ] + return sum(idfs) / len(idfs) if idfs else 0.0 + + def build_channel_bm25_stats( *, unit_rows: Sequence[Mapping[str, Any]], diff --git a/packages/shared-python/shared/services/retrieval/publication_service.py b/packages/shared-python/shared/services/retrieval/publication_service.py index 58fa98dde..84b3b2d27 100644 --- a/packages/shared-python/shared/services/retrieval/publication_service.py +++ b/packages/shared-python/shared/services/retrieval/publication_service.py @@ -200,6 +200,14 @@ def _publish_document_state_for_job( namespace=str(existing_namespace), document_id=document.document_id, ) + # A namespace move mutates both namespace snapshots. Advance the + # old namespace generation as well so request-scoped/process-local + # snapshot caches cannot reuse the pre-move generation. + advance_namespace_generation( + db, + user_id=scope.user_id, + namespace=str(existing_namespace), + ) advance_namespace_generation( db, user_id=scope.user_id, diff --git a/packages/shared-python/shared/services/retrieval/search/map_unit_discovery.py b/packages/shared-python/shared/services/retrieval/search/map_unit_discovery.py index 34d3d74fa..97f5d11e9 100644 --- a/packages/shared-python/shared/services/retrieval/search/map_unit_discovery.py +++ b/packages/shared-python/shared/services/retrieval/search/map_unit_discovery.py @@ -37,9 +37,12 @@ tokenize_query_for_ranker, ) from shared.services.retrieval.nav.persisted_score_load import ( + average_idf_from_namespace_stats, build_channel_bm25_stats, combine_average_idf, ) +from shared.services.retrieval.serving_manifest import decode_serving_manifest +from shared.services.retrieval.cache_service import record_retrieval_index_readiness from shared.services.retrieval.search.scoring import normalize_row_scores from shared.services.retrieval.search.section_filters import is_excluded_section from shared.services.retrieval.settings import ASSET_CHUNK_TYPES @@ -81,6 +84,80 @@ class DiscoveryResult: error: str | None = None +async def _load_exact_namespace_average_idf( + db: AsyncSession, + *, + user_id: str, + namespace: str, +) -> tuple[float, float] | None: + """Load exact namespace IDF floors from the published aggregate tables.""" + generation_row = ( + await db.execute( + text( + "SELECT generation FROM retrieval_namespace_generations " + "WHERE user_id = :user_id AND namespace = :namespace" + ), + {"user_id": user_id, "namespace": namespace}, + ) + ).first() + if generation_row is None: + return None + generation = int(generation_row[0]) + stat_row = ( + await db.execute( + text( + "SELECT payload_zlib, checksum, format_version " + "FROM retrieval_namespace_stats " + "WHERE user_id = :user_id AND namespace = :namespace " + "AND generation = :generation" + ), + { + "user_id": user_id, + "namespace": namespace, + "generation": generation, + }, + ) + ).first() + if stat_row is None: + return None + payload = decode_serving_manifest( + stat_row[0], checksum=str(stat_row[1]), format_version=int(stat_row[2]) + ) + unit_count = int(payload.get("unit_count") or 0) + if unit_count <= 0: + return None + token_rows = ( + await db.execute( + text( + "SELECT channel, document_frequency " + "FROM retrieval_namespace_token_stats " + "WHERE user_id = :user_id AND namespace = :namespace " + "AND generation = :generation AND channel = ANY(:channels)" + ), + { + "user_id": user_id, + "namespace": namespace, + "generation": generation, + "channels": ["path", "content"], + }, + ) + ).all() + frequencies: dict[str, list[int]] = {"path": [], "content": []} + for channel, frequency in token_rows: + if str(channel) in frequencies: + frequencies[str(channel)].append(int(frequency)) + return ( + average_idf_from_namespace_stats( + unit_count=unit_count, + token_document_frequencies=frequencies["path"], + ), + average_idf_from_namespace_stats( + unit_count=unit_count, + token_document_frequencies=frequencies["content"], + ), + ) + + def _build_revision_scope( revision_pins: Mapping[str, str] | None, ) -> tuple[str, str, dict[str, Any]]: @@ -251,15 +328,84 @@ async def map_unit_discovery( (float(path_idf or 0.0), float(content_idf or 0.0), int(unit_count or 0)) for path_idf, content_idf, unit_count in index_result.all() ] - average_idf_path = combine_average_idf( - [(path_idf, unit_count) for path_idf, _content_idf, unit_count in index_parts] - ) - average_idf_content = combine_average_idf( - [ - (content_idf, unit_count) - for _path_idf, content_idf, unit_count in index_parts - ] + expected_revisions = { + (str(row["document_id"]), str(row["job_result_id"])) for row in unit_rows + } + unfiltered_scope = not any( + ( + chunk_types, + signal_paths, + exclude_sections, + exclude_document_ids, + ) ) + index_unit_count_mismatch = unfiltered_scope and sum( + unit_count for _path_idf, _content_idf, unit_count in index_parts + ) != len(unit_rows) + if len(index_parts) != len(expected_revisions) or index_unit_count_mismatch: + try: + await record_retrieval_index_readiness( + user_id=user_id, + namespace=namespace, + ready=False, + expected_revisions=len(expected_revisions), + indexed_revisions=len(index_parts), + ) + except Exception as exc: + logger.warning("retrieval index readiness publish failed: %s", exc) + logger.warning( + "retrieval map index incomplete user_id=%s namespace=%s " + "expected_revisions=%d indexed_revisions=%d fallback=legacy_fts", + user_id, + namespace, + len(expected_revisions), + len(index_parts), + ) + return await _legacy_chunk_discovery( + db, + user_id=user_id, + namespace=namespace, + query=query, + top_k=top_k, + exclude_document_ids=exclude_document_ids, + exclude_sections=exclude_sections, + chunk_types=chunk_types, + signal_paths=signal_paths or [], + filter_mode=filter_mode, + revision_pins=revision_pins, + ) + try: + await record_retrieval_index_readiness( + user_id=user_id, + namespace=namespace, + ready=True, + expected_revisions=len(expected_revisions), + indexed_revisions=len(index_parts), + ) + except Exception as exc: + logger.warning("retrieval index readiness publish failed: %s", exc) + exact_namespace_idf = None + if revision_pins is None: + exact_namespace_idf = await _load_exact_namespace_average_idf( + db, + user_id=user_id, + namespace=namespace, + ) + if exact_namespace_idf is not None: + average_idf_path, average_idf_content = exact_namespace_idf + else: + average_idf_path = combine_average_idf( + [ + (path_idf, unit_count) + for path_idf, _content_idf, unit_count in index_parts + ] + ) + average_idf_content = combine_average_idf( + [ + (content_idf, unit_count) + for _path_idf, content_idf, unit_count in index_parts + ] + ) path_stats = build_channel_bm25_stats( unit_rows=unit_rows, @@ -337,6 +483,103 @@ async def map_unit_discovery( ) +async def _legacy_chunk_discovery( + db: AsyncSession, + *, + user_id: str, + namespace: str, + query: str, + top_k: int, + exclude_document_ids: list[str], + exclude_sections: list[dict[str, str]], + chunk_types: set[str] | None, + signal_paths: list[str], + filter_mode: str, + revision_pins: Mapping[str, str] | None, +) -> DiscoveryResult: + """Bounded lexical fallback used while a serving index is incomplete.""" + clauses = [ + "d.user_id = :user_id", + "d.namespace = :namespace", + "d.status = 'active'", + ] + params: dict[str, Any] = { + "user_id": user_id, + "namespace": namespace, + "query": query, + "limit": max(1, int(top_k)), + } + if revision_pins is None: + clauses.append("d.current_job_result_id = dc.job_result_id") + else: + pairs = [ + (str(document_id), str(job_result_id)) + for document_id, job_result_id in revision_pins.items() + ] + if not pairs: + return DiscoveryResult(status="discovery_done", payload={"fused_rows": []}) + placeholders = [] + for index, (document_id, job_result_id) in enumerate(pairs): + document_key = f"_legacy_doc_{index}" + revision_key = f"_legacy_revision_{index}" + placeholders.append(f"(:{document_key}, :{revision_key})") + params[document_key] = document_id + params[revision_key] = job_result_id + clauses.append(f"(dc.document_id, dc.job_result_id) IN ({', '.join(placeholders)})") + if exclude_document_ids: + clauses.append("d.document_id <> ALL(:excluded_doc_ids)") + params["excluded_doc_ids"] = exclude_document_ids + if chunk_types: + type_keys = [] + for index, chunk_type in enumerate(sorted(chunk_types)): + key = f"_legacy_type_{index}" + type_keys.append(f":{key}") + params[key] = chunk_type + clauses.append(f"LOWER(dc.chunk_type) IN ({', '.join(type_keys)})") + if signal_paths: + signal_parts = [] + for index, signal in enumerate(signal_paths): + key = f"_legacy_signal_{index}" + signal_parts.append("LOWER(COALESCE(ds.section_path, '')) LIKE :" + key) + params[key] = f"%{signal.lower()}%" + combined = " OR ".join(signal_parts) + clauses.append(f"({combined})" if filter_mode == "keep" else f"NOT ({combined})") + for index, item in enumerate(exclude_sections): + document_id = str(item.get("document_id") or "").strip() + section_path = str(item.get("section_path") or "").strip() + if not document_id or not section_path: + continue + doc_key = f"_legacy_exclude_doc_{index}" + path_key = f"_legacy_exclude_path_{index}" + params[doc_key] = document_id + params[path_key] = section_path + clauses.append( + "NOT (dc.document_id = :" + doc_key + " AND (" + "COALESCE(ds.section_path, '') = :" + path_key + " OR " + "POSITION(:" + path_key + " || ' / ' IN COALESCE(ds.section_path, '')) = 1))" + ) + where_sql = " AND ".join(clauses) + statement = text( + "SELECT dc.chunk_id, dc.document_id, dc.section_id, dc.chunk_type, " + "dc.content, dc.source_chunk_path, dc.file_path, dc.chunk_metadata, " + "dc.job_result_id, dc.sort_order, ds.section_path, d.source_file_name, " + "jr.job_id, GREATEST(ts_rank_cd(dc.path_search_tsv, plainto_tsquery('simple', :query)), " + "2 * ts_rank_cd(dc.content_search_tsv, plainto_tsquery('simple', :query))) AS score " + "FROM document_chunks dc JOIN documents d ON d.document_id = dc.document_id " + "LEFT JOIN document_sections ds ON ds.section_id = dc.section_id " + "LEFT JOIN job_results jr ON jr.id = dc.job_result_id " + f"WHERE {where_sql} AND (dc.path_search_tsv @@ plainto_tsquery('simple', :query) " + "OR dc.content_search_tsv @@ plainto_tsquery('simple', :query) " + "OR LOWER(COALESCE(dc.term_search_text, '')) LIKE LOWER(:term_query)) " + "ORDER BY score DESC, dc.sort_order, dc.chunk_id LIMIT :limit" + ) + params["term_query"] = f"%{query}%" + rows = [dict(row._mapping) for row in (await db.execute(statement, params)).all()] + if rows: + normalize_row_scores(rows, source_field="score", target_field="discovery_score", default=0.5) + return DiscoveryResult(status="discovery_done", payload={"fused_rows": rows}) + + def _as_metadata_dict(value: object) -> dict[str, Any]: if isinstance(value, dict): return value diff --git a/packages/shared-python/shared/services/retrieval/serving_manifest.py b/packages/shared-python/shared/services/retrieval/serving_manifest.py index 888113202..467d55c91 100644 --- a/packages/shared-python/shared/services/retrieval/serving_manifest.py +++ b/packages/shared-python/shared/services/retrieval/serving_manifest.py @@ -4,6 +4,8 @@ import hashlib import json +import logging +import time import zlib from typing import Any @@ -26,6 +28,7 @@ from shared.services.retrieval.publication_models import DocumentPublicationScope SERVING_MANIFEST_FORMAT_VERSION = 1 +_logger = logging.getLogger(__name__) def build_revision_serving_payload( @@ -230,6 +233,7 @@ def rebuild_namespace_serving_statistics( Callers hold the namespace generation lock. The aggregate is prepared for the generation that the caller will publish next. """ + started = time.perf_counter() generation = db.execute( select(RetrievalNamespaceGeneration) .where(RetrievalNamespaceGeneration.user_id == user_id) @@ -328,6 +332,17 @@ def rebuild_namespace_serving_statistics( ] ) db.flush() + _logger.info( + "retrieval namespace statistics rebuilt user_id=%s namespace=%s " + "generation=%d documents=%d units=%d token_stats=%d seconds=%.3f", + user_id, + namespace, + target_generation, + aggregate["document_count"], + aggregate["unit_count"], + len(document_frequencies), + time.perf_counter() - started, + ) return target_generation From cfabb8b9eb5ef1275b7b4ea7bbca2d8f88c6487b Mon Sep 17 00:00:00 2001 From: chengke <404835780@qq.com> Date: Tue, 1 Sep 2026 08:34:07 +0800 Subject: [PATCH 2/5] fix: ignore bare markdown hash lines as headings Lines that are only '#' markers have no title text after stripping and must not open an empty section path during chunking. Co-authored-by: Cursor --- .../services/document_parser/structure/heading_candidates.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/apps/worker/app/services/document_parser/structure/heading_candidates.py b/apps/worker/app/services/document_parser/structure/heading_candidates.py index b6f80512b..b77b51e45 100644 --- a/apps/worker/app/services/document_parser/structure/heading_candidates.py +++ b/apps/worker/app/services/document_parser/structure/heading_candidates.py @@ -303,6 +303,10 @@ def _estimate_markdown_heading_level(line: str, meta_ctx: Any | None): if hash_level <= 0: return code_level, code_reason, line_clean + # "#" / "##" with no title text must not become headings via hash_level alone. + if not stripped_line.strip(): + return -1, f"{hash_level}# AND empty-title {code_reason}", line_clean + if isinstance(code_level, int): est_level = max(hash_level, code_level) else: From 1155cd7ba6441ba19f9375e78ab5eed4f479f023 Mon Sep 17 00:00:00 2001 From: chengke <404835780@qq.com> Date: Tue, 1 Sep 2026 10:17:56 +0800 Subject: [PATCH 3/5] refactor: remove RetrievalServingRevisionStat and related statistics handling This commit removes the RetrievalServingRevisionStat model and its associated logic from the document lifecycle and backfill processes. The changes include the removal of calls to rebuild namespace serving statistics and the deletion of related database entries, streamlining the document archiving process. Additionally, the backfill script has been updated to reflect these changes, ensuring it no longer attempts to rebuild statistics for archived documents. --- CONTEXT.md | 2 +- ...d4e_drop_retrieval_namespace_statistics.py | 119 +++++ .../services/documents/lifecycle_service.py | 22 - apps/api/scripts/backfill_map_unit_indexes.py | 10 +- .../tests/contract/test_documents_contract.py | 82 ---- ...mically-publish-retrieval-serving-index.md | 2 +- ...-coherent-retrieval-serving-generations.md | 4 +- docs/design/retrieval-serving-index-plan.md | 425 ------------------ .../shared/models/database/document.py | 99 ---- .../retrieval/namespace_map_snapshot.py | 8 +- .../services/retrieval/publication_service.py | 13 - .../services/retrieval/serving_manifest.py | 203 +-------- 12 files changed, 129 insertions(+), 860 deletions(-) create mode 100644 apps/api/alembic/versions/9f0a1b2c3d4e_drop_retrieval_namespace_statistics.py delete mode 100644 docs/design/retrieval-serving-index-plan.md diff --git a/CONTEXT.md b/CONTEXT.md index b67e16d02..8609579f8 100644 --- a/CONTEXT.md +++ b/CONTEXT.md @@ -221,7 +221,7 @@ 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 +document revisions and their serving-index data. Retrieval captures one generation and retries or falls back if publication changes it during capture. ### Retrieval Semantic Parity diff --git a/apps/api/alembic/versions/9f0a1b2c3d4e_drop_retrieval_namespace_statistics.py b/apps/api/alembic/versions/9f0a1b2c3d4e_drop_retrieval_namespace_statistics.py new file mode 100644 index 000000000..550690148 --- /dev/null +++ b/apps/api/alembic/versions/9f0a1b2c3d4e_drop_retrieval_namespace_statistics.py @@ -0,0 +1,119 @@ +"""Drop unused retrieval namespace statistics tables. + +Query-time BM25 scoring reads only ``document_map_unit_tokens``, +``document_map_units``, and ``document_map_unit_indexes``. The per-revision and +namespace statistics tables were only written by publication/backfill and never +read by the retrieval path, so they are removed here. +""" + +from __future__ import annotations + +from collections.abc import Sequence + +from alembic import op +import sqlalchemy as sa + + +revision: str = "9f0a1b2c3d4e" +down_revision: str | None = "8e9f0a1b2c3d" +branch_labels: Sequence[str] | None = None +depends_on: Sequence[str] | None = None + +__all__ = [ + "revision", + "down_revision", + "branch_labels", + "depends_on", + "upgrade", + "downgrade", +] + + +def upgrade() -> 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) + + +def downgrade() -> None: + bind = op.get_bind() + inspector = sa.inspect(bind) + if not inspector.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", + ), + ) + op.create_index( + "idx_retrieval_serving_revision_stats_scope", + "retrieval_serving_revision_stats", + ["user_id", "namespace", "document_id", "job_result_id"], + ) + if not inspector.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 inspector.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", + ), + ) + op.create_index( + "idx_retrieval_namespace_token_stats_lookup", + "retrieval_namespace_token_stats", + ["user_id", "namespace", "generation", "channel", "token_hash"], + ) diff --git a/apps/api/app/services/documents/lifecycle_service.py b/apps/api/app/services/documents/lifecycle_service.py index 3ea08f05d..fd97a9ba2 100644 --- a/apps/api/app/services/documents/lifecycle_service.py +++ b/apps/api/app/services/documents/lifecycle_service.py @@ -8,13 +8,11 @@ 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, - RetrievalServingRevisionStat, ) from shared.services.retrieval.cache_service import ( invalidate_retrieval_cache_namespaces, @@ -27,9 +25,6 @@ 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 @@ -469,23 +464,6 @@ async def archive_document( ) ) 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: remove_document_from_namespace_map_snapshot( sync_db, diff --git a/apps/api/scripts/backfill_map_unit_indexes.py b/apps/api/scripts/backfill_map_unit_indexes.py index 6eed78eff..ebf1a04d2 100644 --- a/apps/api/scripts/backfill_map_unit_indexes.py +++ b/apps/api/scripts/backfill_map_unit_indexes.py @@ -1,8 +1,8 @@ """Backfill persisted MAP-NAV lexical indexes for existing revisions. Rebuilds, per active revision: the map-unit index, the revision serving -manifest, that document's subtree in the namespace MAP snapshot, namespace -statistics, and the namespace generation. The migrations that create these +manifest, that document's subtree in the namespace MAP snapshot, and the +namespace generation. The migrations that create these derived tables leave them empty intentionally. Run this command after deployment with ``--apply`` so each revision is rebuilt and committed independently; without ``--apply`` it is a read-only inventory. @@ -69,7 +69,6 @@ def _bootstrap_python_path() -> None: from shared.services.retrieval.serving_manifest import ( decode_serving_manifest, persist_revision_serving_state, - rebuild_namespace_serving_statistics, ) @@ -346,11 +345,6 @@ def backfill_map_unit_indexes(*, apply: bool, document_id: str = "") -> int: patch_namespace_map_snapshot( db, scope=scope, manifest_payload=manifest_payload ) - rebuild_namespace_serving_statistics( - db, - user_id=scope.user_id, - namespace=scope.namespace, - ) advance_namespace_generation( db, user_id=scope.user_id, diff --git a/apps/api/tests/contract/test_documents_contract.py b/apps/api/tests/contract/test_documents_contract.py index a71d4029d..3e33f2021 100644 --- a/apps/api/tests/contract/test_documents_contract.py +++ b/apps/api/tests/contract/test_documents_contract.py @@ -13,7 +13,6 @@ 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: @@ -1333,84 +1332,3 @@ 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/docs/adr/0006-atomically-publish-retrieval-serving-index.md b/docs/adr/0006-atomically-publish-retrieval-serving-index.md index 6f0b08008..49e3e4729 100644 --- a/docs/adr/0006-atomically-publish-retrieval-serving-index.md +++ b/docs/adr/0006-atomically-publish-retrieval-serving-index.md @@ -2,5 +2,5 @@ - 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. +- Decision: Build the map-unit index and serving manifest 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 index 89d2ca05b..f875480dc 100644 --- a/docs/adr/0007-use-coherent-retrieval-serving-generations.md +++ b/docs/adr/0007-use-coherent-retrieval-serving-generations.md @@ -1,6 +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. +- Context: A namespace can contain many active document revisions, and publication can replace them while a retrieval request is loading serving metadata. - 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. +- Consequences: Retrieval never combines incompatible revision metadata. 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/design/retrieval-serving-index-plan.md b/docs/design/retrieval-serving-index-plan.md deleted file mode 100644 index 5e66f6d63..000000000 --- a/docs/design/retrieval-serving-index-plan.md +++ /dev/null @@ -1,425 +0,0 @@ -# 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/packages/shared-python/shared/models/database/document.py b/packages/shared-python/shared/models/database/document.py index 739221c48..3d630b917 100644 --- a/packages/shared-python/shared/models/database/document.py +++ b/packages/shared-python/shared/models/database/document.py @@ -443,105 +443,6 @@ class RetrievalServingRevisionManifest(Base): ) -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 RetrievalNamespaceMapSnapshot(Base): """Persisted namespace-level MAP (sections + chunk index + map units). diff --git a/packages/shared-python/shared/services/retrieval/namespace_map_snapshot.py b/packages/shared-python/shared/services/retrieval/namespace_map_snapshot.py index bd8b7767f..41fa6c6b4 100644 --- a/packages/shared-python/shared/services/retrieval/namespace_map_snapshot.py +++ b/packages/shared-python/shared/services/retrieval/namespace_map_snapshot.py @@ -2,9 +2,8 @@ Callers must already hold the namespace generation lock (see ``serving_generation.lock_namespace_generation``) before calling either -function here, exactly as ``rebuild_namespace_serving_statistics`` requires. -Each call only touches one document's subtree; every other document's -subtree in the payload is left byte-for-byte unchanged. +function here. Each call only touches one document's subtree; every other +document's subtree in the payload is left byte-for-byte unchanged. """ from __future__ import annotations @@ -81,8 +80,7 @@ def remove_document_from_namespace_map_snapshot( def _target_generation(db: Session, *, user_id: str, namespace: str) -> int: """Namespace generation this snapshot is prepared for (current + 1). - Mirrors ``rebuild_namespace_serving_statistics``: callers advance the - generation after this write, in the same transaction. + Callers advance the generation after this write, in the same transaction. """ generation = db.execute( select(RetrievalNamespaceGeneration) diff --git a/packages/shared-python/shared/services/retrieval/publication_service.py b/packages/shared-python/shared/services/retrieval/publication_service.py index 84b3b2d27..7ee5dec2b 100644 --- a/packages/shared-python/shared/services/retrieval/publication_service.py +++ b/packages/shared-python/shared/services/retrieval/publication_service.py @@ -38,9 +38,6 @@ advance_namespace_generation, lock_namespace_generation, ) -from shared.services.retrieval.serving_manifest import ( - rebuild_namespace_serving_statistics, -) def utc_now_naive() -> datetime: @@ -183,17 +180,7 @@ 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), - ) remove_document_from_namespace_map_snapshot( db, user_id=scope.user_id, diff --git a/packages/shared-python/shared/services/retrieval/serving_manifest.py b/packages/shared-python/shared/services/retrieval/serving_manifest.py index 467d55c91..c66ae5798 100644 --- a/packages/shared-python/shared/services/retrieval/serving_manifest.py +++ b/packages/shared-python/shared/services/retrieval/serving_manifest.py @@ -4,8 +4,6 @@ import hashlib import json -import logging -import time import zlib from typing import Any @@ -15,20 +13,13 @@ 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 -_logger = logging.getLogger(__name__) def build_revision_serving_payload( @@ -121,82 +112,25 @@ def build_revision_serving_payload( } -def build_revision_statistics_payload( - db: Session, - *, - scope: DocumentPublicationScope, -) -> dict[str, Any]: - """Build compressed scoring contributions for one revision. - - Stores aggregate token frequencies for namespace statistics rebuild. - Per-unit frequencies stay in ``document_map_unit_tokens`` and are loaded - at query time by the query tokens only. - """ - 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": {}} - 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 not in frequencies: - continue - token_key = str(token) - frequencies[channel_key][token_key] = frequencies[channel_key].get( - token_key, 0 - ) + int(frequency) - 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, - } - - def persist_revision_serving_state( db: Session, *, scope: DocumentPublicationScope, ) -> dict[str, Any]: - """Replace manifest and statistics rows for one revision atomically. + """Replace the serving manifest row for one revision atomically. Returns the manifest payload so callers can patch the namespace-level MAP snapshot without rebuilding it. """ 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, @@ -208,144 +142,9 @@ def persist_revision_serving_state( 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, - ) - ) return manifest_payload -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. - """ - started = time.perf_counter() - 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() - _logger.info( - "retrieval namespace statistics rebuilt user_id=%s namespace=%s " - "generation=%d documents=%d units=%d token_stats=%d seconds=%.3f", - user_id, - namespace, - target_generation, - aggregate["document_count"], - aggregate["unit_count"], - len(document_frequencies), - time.perf_counter() - started, - ) - return target_generation - - def _connection_target_ids(metadata: Any) -> list[str]: if not isinstance(metadata, dict): return [] From 8fb4fe861b91b34e66b0c7895c143d3efd11ace9 Mon Sep 17 00:00:00 2001 From: suguanYang Date: Tue, 1 Sep 2026 10:27:37 +0800 Subject: [PATCH 4/5] perf: avoid repeated section scans in map discovery --- .../retrieval/search/map_unit_discovery.py | 41 ++++++++++++++++++- 1 file changed, 39 insertions(+), 2 deletions(-) diff --git a/packages/shared-python/shared/services/retrieval/search/map_unit_discovery.py b/packages/shared-python/shared/services/retrieval/search/map_unit_discovery.py index 97f5d11e9..70a5c1cac 100644 --- a/packages/shared-python/shared/services/retrieval/search/map_unit_discovery.py +++ b/packages/shared-python/shared/services/retrieval/search/map_unit_discovery.py @@ -75,6 +75,25 @@ ) """ +_SCOPED_UNIT_IDS_CTE = """ +WITH scoped_units AS ( + SELECT + dmu.id AS map_unit_id, + dmu.document_id, + dmu.job_result_id + FROM document_map_units dmu + JOIN documents d + ON d.document_id = dmu.document_id + {revision_join} + WHERE d.user_id = :user_id + AND d.namespace = :namespace + AND d.status = 'active' + {revision_clause} + {exclude_clause} + {type_clause} +) +""" + @dataclass class DiscoveryResult: @@ -287,7 +306,16 @@ async def map_unit_discovery( frequency_result = await db.execute( text( - cte + ( + cte + if signal_paths + else _SCOPED_UNIT_IDS_CTE.format( + revision_join=revision_join, + revision_clause=revision_clause, + exclude_clause=exclude_clause, + type_clause=type_clause, + ) + ) + """ SELECT tokens.map_unit_id, tokens.channel, tokens.token, tokens.frequency FROM document_map_unit_tokens AS tokens @@ -310,7 +338,16 @@ async def map_unit_discovery( index_result = await db.execute( text( - cte + ( + cte + if signal_paths + else _SCOPED_UNIT_IDS_CTE.format( + revision_join=revision_join, + revision_clause=revision_clause, + exclude_clause=exclude_clause, + type_clause=type_clause, + ) + ) + """ SELECT indexes.average_idf_path, indexes.average_idf_content, indexes.unit_count From 07ca9d010587493205170de6a0f940754ed758aa Mon Sep 17 00:00:00 2001 From: suguanYang Date: Tue, 1 Sep 2026 10:44:44 +0800 Subject: [PATCH 5/5] fix: remove dropped namespace statistics reads --- .../services/retrieval/nav/nav_knowhere.py | 55 ++------- .../retrieval/nav/persisted_score_load.py | 23 ---- .../retrieval/search/map_unit_discovery.py | 110 ++---------------- 3 files changed, 21 insertions(+), 167 deletions(-) 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 dfbad8717..38cef2249 100644 --- a/packages/shared-python/shared/services/retrieval/nav/nav_knowhere.py +++ b/packages/shared-python/shared/services/retrieval/nav/nav_knowhere.py @@ -300,7 +300,6 @@ def load_persisted_score_corpus( comes from ``document_map_unit_indexes`` (written at index time). """ from shared.services.retrieval.nav.persisted_score_load import ( - average_idf_from_namespace_stats, build_channel_bm25_stats, combine_average_idf, ) @@ -360,51 +359,15 @@ def load_persisted_score_corpus( revision_key ] else: - total_unit_count = sum(int(row[3] or 0) for row in index_rows) - try: - cur.execute( - "SELECT tokens.channel, tokens.token, " - "COUNT(DISTINCT tokens.map_unit_id) " - "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 " - "WHERE tokens.channel = ANY(%s) " - "GROUP BY tokens.channel, tokens.token", - [*revision_params, ["path", "content"]], - ) - namespace_token_dfs: dict[str, list[int]] = { - "path": [], - "content": [], - } - for channel, _token, document_frequency in cur.fetchall(): - if str(channel) in namespace_token_dfs: - namespace_token_dfs[str(channel)].append( - int(document_frequency) - ) - average_idf_path = average_idf_from_namespace_stats( - unit_count=total_unit_count, - token_document_frequencies=namespace_token_dfs["path"], - ) - average_idf_content = average_idf_from_namespace_stats( - unit_count=total_unit_count, - token_document_frequencies=namespace_token_dfs["content"], - ) - except Exception as exc: - _logger.warning( - "exact namespace IDF load failed; using revision averages: %s", - exc, - ) - average_idf_path = combine_average_idf( - [(float(row[4] or 0.0), int(row[3] or 0)) for row in index_rows] - ) - average_idf_content = combine_average_idf( - [ - (float(row[5] or 0.0), int(row[3] or 0)) - for row in index_rows - ] - ) + average_idf_path = combine_average_idf( + [(float(row[4] or 0.0), int(row[3] or 0)) for row in index_rows] + ) + average_idf_content = combine_average_idf( + [ + (float(row[5] or 0.0), int(row[3] or 0)) + for row in index_rows + ] + ) self._score_average_idf_cache[revision_key] = ( average_idf_path, average_idf_content, diff --git a/packages/shared-python/shared/services/retrieval/nav/persisted_score_load.py b/packages/shared-python/shared/services/retrieval/nav/persisted_score_load.py index 7f1eb2029..9ce10c2e1 100644 --- a/packages/shared-python/shared/services/retrieval/nav/persisted_score_load.py +++ b/packages/shared-python/shared/services/retrieval/nav/persisted_score_load.py @@ -38,29 +38,6 @@ def combine_average_idf(parts: Sequence[tuple[float, int]]) -> float: ) -def average_idf_from_namespace_stats( - *, - unit_count: int, - token_document_frequencies: Sequence[int], -) -> float: - """Compute the exact namespace-level average IDF used by rank_bm25. - - Namespace token statistics already contain one document frequency per - token. Computing the mean from those rows avoids the incorrect - per-revision-average approximation when a namespace contains revisions - with different token distributions. - """ - if unit_count <= 0: - return 0.0 - idfs = [ - math.log(unit_count - int(frequency) + 0.5) - - math.log(int(frequency) + 0.5) - for frequency in token_document_frequencies - if 0 < int(frequency) <= unit_count - ] - return sum(idfs) / len(idfs) if idfs else 0.0 - - def build_channel_bm25_stats( *, unit_rows: Sequence[Mapping[str, Any]], diff --git a/packages/shared-python/shared/services/retrieval/search/map_unit_discovery.py b/packages/shared-python/shared/services/retrieval/search/map_unit_discovery.py index 70a5c1cac..ab76977fd 100644 --- a/packages/shared-python/shared/services/retrieval/search/map_unit_discovery.py +++ b/packages/shared-python/shared/services/retrieval/search/map_unit_discovery.py @@ -37,11 +37,9 @@ tokenize_query_for_ranker, ) from shared.services.retrieval.nav.persisted_score_load import ( - average_idf_from_namespace_stats, build_channel_bm25_stats, combine_average_idf, ) -from shared.services.retrieval.serving_manifest import decode_serving_manifest from shared.services.retrieval.cache_service import record_retrieval_index_readiness from shared.services.retrieval.search.scoring import normalize_row_scores from shared.services.retrieval.search.section_filters import is_excluded_section @@ -103,80 +101,6 @@ class DiscoveryResult: error: str | None = None -async def _load_exact_namespace_average_idf( - db: AsyncSession, - *, - user_id: str, - namespace: str, -) -> tuple[float, float] | None: - """Load exact namespace IDF floors from the published aggregate tables.""" - generation_row = ( - await db.execute( - text( - "SELECT generation FROM retrieval_namespace_generations " - "WHERE user_id = :user_id AND namespace = :namespace" - ), - {"user_id": user_id, "namespace": namespace}, - ) - ).first() - if generation_row is None: - return None - generation = int(generation_row[0]) - stat_row = ( - await db.execute( - text( - "SELECT payload_zlib, checksum, format_version " - "FROM retrieval_namespace_stats " - "WHERE user_id = :user_id AND namespace = :namespace " - "AND generation = :generation" - ), - { - "user_id": user_id, - "namespace": namespace, - "generation": generation, - }, - ) - ).first() - if stat_row is None: - return None - payload = decode_serving_manifest( - stat_row[0], checksum=str(stat_row[1]), format_version=int(stat_row[2]) - ) - unit_count = int(payload.get("unit_count") or 0) - if unit_count <= 0: - return None - token_rows = ( - await db.execute( - text( - "SELECT channel, document_frequency " - "FROM retrieval_namespace_token_stats " - "WHERE user_id = :user_id AND namespace = :namespace " - "AND generation = :generation AND channel = ANY(:channels)" - ), - { - "user_id": user_id, - "namespace": namespace, - "generation": generation, - "channels": ["path", "content"], - }, - ) - ).all() - frequencies: dict[str, list[int]] = {"path": [], "content": []} - for channel, frequency in token_rows: - if str(channel) in frequencies: - frequencies[str(channel)].append(int(frequency)) - return ( - average_idf_from_namespace_stats( - unit_count=unit_count, - token_document_frequencies=frequencies["path"], - ), - average_idf_from_namespace_stats( - unit_count=unit_count, - token_document_frequencies=frequencies["content"], - ), - ) - - def _build_revision_scope( revision_pins: Mapping[str, str] | None, ) -> tuple[str, str, dict[str, Any]]: @@ -421,28 +345,18 @@ async def map_unit_discovery( ) except Exception as exc: logger.warning("retrieval index readiness publish failed: %s", exc) - exact_namespace_idf = None - if revision_pins is None: - exact_namespace_idf = await _load_exact_namespace_average_idf( - db, - user_id=user_id, - namespace=namespace, - ) - if exact_namespace_idf is not None: - average_idf_path, average_idf_content = exact_namespace_idf - else: - average_idf_path = combine_average_idf( - [ - (path_idf, unit_count) - for path_idf, _content_idf, unit_count in index_parts - ] - ) - average_idf_content = combine_average_idf( - [ - (content_idf, unit_count) - for _path_idf, content_idf, unit_count in index_parts - ] - ) + average_idf_path = combine_average_idf( + [ + (path_idf, unit_count) + for path_idf, _content_idf, unit_count in index_parts + ] + ) + average_idf_content = combine_average_idf( + [ + (content_idf, unit_count) + for _path_idf, content_idf, unit_count in index_parts + ] + ) path_stats = build_channel_bm25_stats( unit_rows=unit_rows,