From 05e1aaf0b1531a62f704939e91a7d941ffc1fd6e Mon Sep 17 00:00:00 2001 From: suguanYang Date: Tue, 1 Sep 2026 21:44:23 +0800 Subject: [PATCH 1/5] perf: optimize map-nav snapshot and scoring --- apps/api/scripts/backfill_map_unit_indexes.py | 4 +- ...est_retrieval_serving_manifest_contract.py | 55 ++++ .../shared/services/redis/redis_service.py | 43 +++ .../services/retrieval/execution/routes.py | 11 + .../retrieval/namespace_map_snapshot.py | 8 +- .../retrieval/namespace_map_snapshot_cache.py | 43 --- .../retrieval/namespace_map_snapshot_redis.py | 38 +++ .../services/retrieval/nav/nav_map_scores.py | 46 ++-- .../shared/services/retrieval/nav_snapshot.py | 256 ++++++++++++++++-- .../services/retrieval/serving_manifest.py | 121 ++++++++- 10 files changed, 522 insertions(+), 103 deletions(-) delete mode 100644 packages/shared-python/shared/services/retrieval/namespace_map_snapshot_cache.py create mode 100644 packages/shared-python/shared/services/retrieval/namespace_map_snapshot_redis.py diff --git a/apps/api/scripts/backfill_map_unit_indexes.py b/apps/api/scripts/backfill_map_unit_indexes.py index ebf1a04d2..6d3a64e58 100644 --- a/apps/api/scripts/backfill_map_unit_indexes.py +++ b/apps/api/scripts/backfill_map_unit_indexes.py @@ -67,7 +67,7 @@ def _bootstrap_python_path() -> None: lock_namespace_generation, ) from shared.services.retrieval.serving_manifest import ( - decode_serving_manifest, + decode_namespace_map_snapshot, persist_revision_serving_state, ) @@ -160,7 +160,7 @@ def check_fallback_readiness(*, document_id: str = "") -> list[NamespaceFallback snapshot_status = "missing" else: try: - payload = decode_serving_manifest( + payload = decode_namespace_map_snapshot( bytes(snapshot.payload_zlib), checksum=str(snapshot.checksum), format_version=int(snapshot.format_version), diff --git a/apps/api/tests/contract/test_retrieval_serving_manifest_contract.py b/apps/api/tests/contract/test_retrieval_serving_manifest_contract.py index 5bb68c063..d80c688da 100644 --- a/apps/api/tests/contract/test_retrieval_serving_manifest_contract.py +++ b/apps/api/tests/contract/test_retrieval_serving_manifest_contract.py @@ -5,8 +5,11 @@ import pytest from shared.services.retrieval.serving_manifest import ( + NAMESPACE_MAP_SNAPSHOT_FORMAT_VERSION, SERVING_MANIFEST_FORMAT_VERSION, + decode_namespace_map_snapshot, decode_serving_manifest, + encode_namespace_map_snapshot, encode_serving_manifest, ) @@ -49,3 +52,55 @@ def test_serving_manifest_rejects_unknown_version() -> None: checksum=checksum, format_version=SERVING_MANIFEST_FORMAT_VERSION + 1, ) + + +def test_namespace_snapshot_uses_routing_only_v2_and_reads_legacy_v1() -> None: + payload = { + "documents": { + "doc_1": { + "job_result_id": "result_1", + "job_id": "job_1", + "source_file_name": "private.pdf", + "sections": [ + { + "section_id": "sec_1", + "section_path": "Root", + "section_title": "Root", + "section_level": 0, + "summary": "summary", + "sort_order": 0, + "unused": "drop", + } + ], + "chunks": [ + { + "chunk_id": "chunk_1", + "section_id": "sec_1", + "chunk_type": "text", + "sort_order": 0, + "connect_to": [], + "content": "drop", + } + ], + } + } + } + compressed, checksum, version = encode_namespace_map_snapshot(payload) + + assert version == NAMESPACE_MAP_SNAPSHOT_FORMAT_VERSION + decoded = decode_namespace_map_snapshot( + compressed, checksum=checksum, format_version=version + ) + document = decoded["documents"]["doc_1"] + assert "source_file_name" not in document + assert "unused" not in document["sections"][0] + assert "content" not in document["chunks"][0] + + legacy_compressed, legacy_checksum, legacy_version = encode_serving_manifest( + payload + ) + assert decode_namespace_map_snapshot( + legacy_compressed, + checksum=legacy_checksum, + format_version=legacy_version, + ) == payload diff --git a/packages/shared-python/shared/services/redis/redis_service.py b/packages/shared-python/shared/services/redis/redis_service.py index e83e3a7bc..185467ecf 100644 --- a/packages/shared-python/shared/services/redis/redis_service.py +++ b/packages/shared-python/shared/services/redis/redis_service.py @@ -160,6 +160,49 @@ async def _operation(): original_exception=e, ) + async def get_bytes(self, key: str) -> bytes | None: + """Get a value without JSON decoding (for compressed binary blobs).""" + try: + client = await self._get_client() + full_key = self._build_key(key) + + async def _operation() -> bytes | None: + result = await client.get(full_key) + if result is None: + return None + if isinstance(result, bytes): + return result + if isinstance(result, bytearray): + return bytes(result) + return str(result).encode("utf-8") + + return await self._execute_with_retry(_operation) + except Exception as e: + logger.error(f"Redis GET_BYTES operation failed: {e}") + raise RedisOperationError( + internal_message=f"GET_BYTES operation failed: {str(e)}", + operation="GET_BYTES", + original_exception=e, + ) + + async def set_bytes(self, key: str, value: bytes, *, ex: int) -> bool: + """Set a compressed binary value with an explicit TTL.""" + try: + client = await self._get_client() + full_key = self._build_key(key) + + async def _operation() -> bool: + return bool(await client.set(full_key, value, ex=ex)) + + return await self._execute_with_retry(_operation) + except Exception as e: + logger.error(f"Redis SET_BYTES operation failed: {e}") + raise RedisOperationError( + internal_message=f"SET_BYTES operation failed: {str(e)}", + operation="SET_BYTES", + original_exception=e, + ) + async def delete(self, *keys: str) -> int: """Delete keys.""" try: diff --git a/packages/shared-python/shared/services/retrieval/execution/routes.py b/packages/shared-python/shared/services/retrieval/execution/routes.py index 5dbf4f7c5..ca9809a45 100644 --- a/packages/shared-python/shared/services/retrieval/execution/routes.py +++ b/packages/shared-python/shared/services/retrieval/execution/routes.py @@ -1,6 +1,7 @@ from __future__ import annotations import asyncio +import resource import time from contextlib import AbstractAsyncContextManager @@ -173,6 +174,7 @@ async def _run_mapnav_route( context: RetrievalRouteContext, ) -> RetrievalRouteOutcome: """Default agentic path: PLANNER + HARVEST + CONTROL (checklist map-nav).""" + process_started = resource.getrusage(resource.RUSAGE_SELF) from shared.services.retrieval import nav_llm_backend # noqa: F401 from shared.services.retrieval.nav import run_nav_episode from shared.services.retrieval.nav.nav_hierarchy import ProviderToolSpace @@ -202,6 +204,7 @@ async def _run_mapnav_route( exclude_sections=context.exclude_sections, lazy=True, revision_pins=snapshot_pins, + generation=(snapshot_pins.generation if snapshot_pins is not None else None), ) if snapshot_pins is not None and not await is_revision_generation_stable( context.db, @@ -223,6 +226,7 @@ async def _run_mapnav_route( exclude_sections=context.exclude_sections, lazy=True, revision_pins=snapshot_pins, + generation=(snapshot_pins.generation if snapshot_pins is not None else None), ) snapshot_seconds = time.perf_counter() - snapshot_started logger.info( @@ -338,6 +342,13 @@ async def _run_mapnav_route( } completion_detail = f"chunks | evidence={len(evidence_text)} chars | router=mapnav" + process_finished = resource.getrusage(resource.RUSAGE_SELF) + logger.info( + "retrieval mapnav stage=process_resources cpu_seconds=%.3f max_rss_kb=%d", + (process_finished.ru_utime + process_finished.ru_stime) + - (process_started.ru_utime + process_started.ru_stime), + int(process_finished.ru_maxrss), + ) return RetrievalRouteOutcome( response=response, hit_stats_results=resolved.refs, 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 41fa6c6b4..2c31b296f 100644 --- a/packages/shared-python/shared/services/retrieval/namespace_map_snapshot.py +++ b/packages/shared-python/shared/services/retrieval/namespace_map_snapshot.py @@ -19,8 +19,8 @@ ) from shared.services.retrieval.publication_models import DocumentPublicationScope from shared.services.retrieval.serving_manifest import ( - decode_serving_manifest, - encode_serving_manifest, + decode_namespace_map_snapshot, + encode_namespace_map_snapshot, ) @@ -107,7 +107,7 @@ def _decode_documents( if row is None: return {} try: - payload = decode_serving_manifest( + payload = decode_namespace_map_snapshot( row.payload_zlib, checksum=row.checksum, format_version=row.format_version, @@ -127,7 +127,7 @@ def _write_snapshot( documents: dict[str, dict[str, Any]], target_generation: int, ) -> None: - encoded, checksum, format_version = encode_serving_manifest( + encoded, checksum, format_version = encode_namespace_map_snapshot( {"documents": documents} ) if row is None: diff --git a/packages/shared-python/shared/services/retrieval/namespace_map_snapshot_cache.py b/packages/shared-python/shared/services/retrieval/namespace_map_snapshot_cache.py deleted file mode 100644 index 6cc5cab29..000000000 --- a/packages/shared-python/shared/services/retrieval/namespace_map_snapshot_cache.py +++ /dev/null @@ -1,43 +0,0 @@ -"""Process-local cache for decoded namespace MAP snapshot documents. - -Keyed by ``(user_id, namespace, generation)`` so a publish/archive that bumps -the namespace generation invalidates stale entries automatically -- no manual -invalidation call is needed. Bounded LRU keeps memory use predictable across -many namespaces sharing one worker process. -""" - -from __future__ import annotations - -import threading -from collections import OrderedDict -from typing import Any - -_MAX_ENTRIES = 64 -_lock = threading.Lock() -_cache: "OrderedDict[tuple[str, str, int], dict[str, dict[str, Any]]]" = OrderedDict() - - -def get_cached_namespace_documents( - *, user_id: str, namespace: str, generation: int -) -> dict[str, dict[str, Any]] | None: - key = (user_id, namespace, generation) - with _lock: - documents = _cache.get(key) - if documents is not None: - _cache.move_to_end(key) - return documents - - -def cache_namespace_documents( - *, - user_id: str, - namespace: str, - generation: int, - documents: dict[str, dict[str, Any]], -) -> None: - key = (user_id, namespace, generation) - with _lock: - _cache[key] = documents - _cache.move_to_end(key) - while len(_cache) > _MAX_ENTRIES: - _cache.popitem(last=False) diff --git a/packages/shared-python/shared/services/retrieval/namespace_map_snapshot_redis.py b/packages/shared-python/shared/services/retrieval/namespace_map_snapshot_redis.py new file mode 100644 index 000000000..fc8447c6a --- /dev/null +++ b/packages/shared-python/shared/services/retrieval/namespace_map_snapshot_redis.py @@ -0,0 +1,38 @@ +"""Redis cache for compressed namespace MAP routing snapshots.""" + +from __future__ import annotations + +from shared.models.schemas.retrieval_namespace import normalize_retrieval_namespace +from shared.services.redis import RedisServiceFactory + +_CACHE_TTL_SECONDS = 3600 +_KEY_PREFIX = "retrieval:snapshot:v2" + + +def snapshot_cache_key(*, user_id: str, namespace: str, generation: int) -> str: + normalized_namespace = normalize_retrieval_namespace(namespace) + return f"{_KEY_PREFIX}:{user_id}:{normalized_namespace}:g{int(generation)}" + + +async def get_snapshot_blob( + *, user_id: str, namespace: str, generation: int +) -> bytes | None: + service = RedisServiceFactory.get_service() + return await service.get_bytes( + snapshot_cache_key( + user_id=user_id, namespace=namespace, generation=generation + ) + ) + + +async def set_snapshot_blob( + *, user_id: str, namespace: str, generation: int, payload_zlib: bytes +) -> bool: + service = RedisServiceFactory.get_service() + return await service.set_bytes( + snapshot_cache_key( + user_id=user_id, namespace=namespace, generation=generation + ), + payload_zlib, + ex=_CACHE_TTL_SECONDS, + ) 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 28c046001..c3637464c 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 @@ -227,31 +227,31 @@ def _pool_unit_scores_to_tree( unit_scores: Dict[str, float], ) -> Dict[str, float]: """MAX-pool globally comparable unit scores onto one document tree.""" - map_scores = { - leaf_id: float(unit_scores.get(leaf_id, 0.0) or 0.0) for leaf_id in leaves - } - - def score_node(section_id: str) -> float: - if section_id in map_scores: - return map_scores[section_id] + # ``_walk_tree`` inserts every parent before its children. Reversing that + # order is therefore a postorder traversal without allocating descendant + # lists or revisiting nodes. + map_scores: Dict[str, float] = {} + descendant_leaf_max: Dict[str, float] = {} + for section_id in reversed(children_map): kids = children_map.get(section_id) or [] if not kids: - score = float(unit_scores.get(section_id, 0.0) or 0.0) - map_scores[section_id] = score - return score - descendant_leaves = _collect_descendant_leaves(section_id, children_map, leaves) - parts = [ - float(unit_scores.get(leaf_id, 0.0) or 0.0) for leaf_id in descendant_leaves - ] - self_key = f"{section_id}__self" - if self_key in unit_scores: - parts.append(float(unit_scores[self_key])) - score = float(max(parts)) if parts else 0.0 - map_scores[section_id] = score - return score - - for section_id in children_map: - score_node(section_id) + leaf_score = float(unit_scores.get(section_id, 0.0) or 0.0) + descendant_leaf_max[section_id] = leaf_score + score = leaf_score + else: + child_leaf_scores = [ + descendant_leaf_max.get(kid, float(unit_scores.get(kid, 0.0) or 0.0)) + for kid in kids + ] + leaf_score = max(child_leaf_scores, default=0.0) + descendant_leaf_max[section_id] = leaf_score + self_score = float(unit_scores.get(f"{section_id}__self", 0.0) or 0.0) + score = max(leaf_score, self_score) + map_scores[section_id] = float(score) + # Preserve the legacy behavior for leaf ids that are not present in the + # children map (defensive support for sparse providers). + for leaf_id in leaves: + map_scores.setdefault(leaf_id, float(unit_scores.get(leaf_id, 0.0) or 0.0)) return map_scores diff --git a/packages/shared-python/shared/services/retrieval/nav_snapshot.py b/packages/shared-python/shared/services/retrieval/nav_snapshot.py index ea502c54b..50a4b4f89 100644 --- a/packages/shared-python/shared/services/retrieval/nav_snapshot.py +++ b/packages/shared-python/shared/services/retrieval/nav_snapshot.py @@ -9,7 +9,9 @@ import json import logging +import os import time +import zlib from dataclasses import dataclass from collections.abc import Callable, Iterator, Mapping from typing import Any, Protocol @@ -33,6 +35,7 @@ DocumentChunk, DocumentSection, RetrievalNamespaceMapSnapshot, + RetrievalNamespaceGeneration, RetrievalServingRevisionManifest, ) from shared.models.database.job_result import JobResult @@ -46,11 +49,14 @@ 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.serving_manifest import ( + decode_namespace_map_snapshot, + decode_serving_manifest, +) from shared.services.retrieval.manifest_cache import get_cached_manifest_payloads -from shared.services.retrieval.namespace_map_snapshot_cache import ( - cache_namespace_documents, - get_cached_namespace_documents, +from shared.services.retrieval.namespace_map_snapshot_redis import ( + get_snapshot_blob, + set_snapshot_blob, ) @@ -70,6 +76,16 @@ _logger = logging.getLogger(__name__) +def _snapshot_timing_enabled() -> bool: + """Enable detailed snapshot timings for local performance diagnosis.""" + return os.environ.get("RETRIEVAL_SNAPSHOT_TIMING", "0").strip().lower() in { + "1", + "true", + "yes", + "on", + } + + class SnapshotSession(Protocol): """Minimal database interface required by the snapshot loader.""" @@ -149,12 +165,14 @@ async def load_nav_snapshot( exclude_sections: list[dict[str, str]] | None = None, lazy: bool = False, revision_pins: Mapping[str, str] | None = None, + generation: int | 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_secs = list(exclude_sections or ()) + snapshot_started = time.perf_counter() if revision_pins is None: doc_stmt = ( @@ -182,7 +200,9 @@ async def load_nav_snapshot( ) if excluded_docs: doc_stmt = doc_stmt.where(Document.document_id.notin_(excluded_docs)) + doc_query_started = time.perf_counter() doc_rows = list((await db.execute(doc_stmt)).all()) + doc_query_seconds = time.perf_counter() - doc_query_started if not doc_rows: raise ValueError( f"no active documents with current revision for " @@ -207,11 +227,13 @@ async def load_nav_snapshot( current_job_result_ids.add(job_result_id) document_revisions.append((did, job_result_id)) + job_query_started = time.perf_counter() job_result_rows = await db.execute( select(JobResult.id, JobResult.job_id).where( JobResult.id.in_(list(current_job_result_ids)) ) ) + job_query_seconds = time.perf_counter() - job_query_started job_id_by_result_id = { str(job_result_id): str(job_id) for job_result_id, job_id in job_result_rows.all() @@ -223,13 +245,16 @@ async def load_nav_snapshot( user_id=user_id, namespace=namespace, document_revisions=document_revisions, + expected_generation=generation, ) if snapshot_entries is not None: + parse_started = time.perf_counter() manifest_sections = _parse_manifest_entries( snapshot_entries, exclude_sections=excluded_secs, job_id_by_result_id=job_id_by_result_id, ) + parse_seconds = time.perf_counter() - parse_started else: _logger.warning( "retrieval snapshot fallback=manifest_merge user_id=%s namespace=%s documents=%d", @@ -245,6 +270,7 @@ async def load_nav_snapshot( exclude_sections=excluded_secs, job_id_by_result_id=job_id_by_result_id, ) + parse_seconds = 0.0 if manifest_sections is None: _logger.warning( "retrieval snapshot fallback=table_scan user_id=%s namespace=%s documents=%d", @@ -324,21 +350,39 @@ async def load_nav_snapshot( for chunk_id in chunk_ids }, ) + ref_index_started = time.perf_counter() + lazy_ref_index = LazyChunkRefIndex( + chunk_ref_index, + resolver=store.load_chunk_reference_metadata, + ) + ref_index_seconds = time.perf_counter() - ref_index_started except Exception: store.close() raise - return NavSnapshot( + snapshot = NavSnapshot( provider=provider, - chunk_ref_index=LazyChunkRefIndex( - chunk_ref_index, - resolver=store.load_chunk_reference_metadata, - ), + chunk_ref_index=lazy_ref_index, document_ids=list(provider.document_ids()), document_titles={ did: kept_titles.get(did, did) for did in provider.document_ids() }, document_revisions=dict(revisions), ) + if _snapshot_timing_enabled(): + _logger.info( + "retrieval snapshot timing total_seconds=%.3f parse_seconds=%.3f " + "ref_index_copy_seconds=%.3f doc_query_seconds=%.3f " + "job_query_seconds=%.3f documents=%d sections=%d refs=%d mode=lazy", + time.perf_counter() - snapshot_started, + parse_seconds, + ref_index_seconds, + doc_query_seconds, + job_query_seconds, + len(snapshot.document_ids), + sum(len(rows) for rows in sections_by_doc.values()), + len(snapshot.chunk_ref_index), + ) + return snapshot units_by_doc, chunk_ref_index = await _load_chunks( db, @@ -358,7 +402,7 @@ async def load_nav_snapshot( f"user_id={user_id!r} namespace={namespace!r}" ) - return build_nav_snapshot( + snapshot = build_nav_snapshot( document_titles=kept_titles, 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}, @@ -369,6 +413,20 @@ async def load_nav_snapshot( if document_id in kept_titles }, ) + if _snapshot_timing_enabled(): + _logger.info( + "retrieval snapshot timing total_seconds=%.3f parse_seconds=%.3f " + "doc_query_seconds=%.3f job_query_seconds=%.3f documents=%d " + "sections=%d refs=%d mode=eager", + time.perf_counter() - snapshot_started, + parse_seconds, + doc_query_seconds, + job_query_seconds, + len(snapshot.document_ids), + sum(len(rows) for rows in sections_by_doc.values()), + len(snapshot.chunk_ref_index), + ) + return snapshot async def _resolve_namespace_snapshot_entries( @@ -377,21 +435,39 @@ async def _resolve_namespace_snapshot_entries( user_id: str, namespace: str, document_revisions: list[tuple[str, str]], + expected_generation: int | None = None, ) -> list[tuple[object, ...]] | None: """Return manifest-shaped entries from the persisted namespace snapshot. Returns ``None`` (triggering the exact per-revision fallback) when the snapshot row is missing, corrupt, or stale for any requested revision. """ + generation_statement = select(RetrievalNamespaceGeneration.generation).where( + RetrievalNamespaceGeneration.user_id == user_id, + RetrievalNamespaceGeneration.namespace == namespace, + ) + try: + generation_result = await db.execute(generation_statement) + _current_generation = generation_result.scalar_one_or_none() + except SQLAlchemyError as exc: + await db.rollback() + _logger.warning("retrieval snapshot generation lookup failed error=%s", exc) + _current_generation = None + if _snapshot_timing_enabled(): + _logger.info( + "retrieval snapshot generation generation=%s", + _current_generation, + ) + generation_value = int(expected_generation) if expected_generation is not None else None statement = select( RetrievalNamespaceMapSnapshot.generation, - RetrievalNamespaceMapSnapshot.payload_zlib, RetrievalNamespaceMapSnapshot.checksum, RetrievalNamespaceMapSnapshot.format_version, ).where( RetrievalNamespaceMapSnapshot.user_id == user_id, RetrievalNamespaceMapSnapshot.namespace == namespace, ) + lookup_started = time.perf_counter() try: row = (await db.execute(statement)).first() except SQLAlchemyError as exc: @@ -403,30 +479,124 @@ async def _resolve_namespace_snapshot_entries( exc, ) return None + lookup_seconds = time.perf_counter() - lookup_started + if _snapshot_timing_enabled(): + _logger.info( + "retrieval snapshot lookup seconds=%.3f found=%s", + lookup_seconds, + row is not None, + ) if row is None: return None - generation, payload_zlib, checksum, format_version = row - documents = get_cached_namespace_documents( - user_id=user_id, namespace=namespace, generation=int(generation) - ) - if documents is None: - try: - payload = decode_serving_manifest( - bytes(payload_zlib), - checksum=str(checksum), - format_version=int(format_version), + generation, checksum, format_version = row + row_generation = int(generation) + if generation_value is not None and row_generation != generation_value: + _logger.info( + "retrieval snapshot generation mismatch row=%d expected=%d", + row_generation, + generation_value, + ) + return None + cached_blob: bytes | None = None + try: + cached_blob = await get_snapshot_blob( + user_id=user_id, namespace=namespace, generation=row_generation + ) + except Exception as exc: + _logger.warning("retrieval snapshot redis get failed error=%s", exc) + database_blob: bytes | None = None + if cached_blob is None: + payload_result = await db.execute( + select(RetrievalNamespaceMapSnapshot.payload_zlib).where( + RetrievalNamespaceMapSnapshot.user_id == user_id, + RetrievalNamespaceMapSnapshot.namespace == namespace, ) - except (ValueError, TypeError): + ) + payload_row = payload_result.first() + if payload_row is None or payload_row[0] is None: return None - decoded_documents = payload.get("documents") - if not isinstance(decoded_documents, dict): + database_blob = bytes(payload_row[0]) + blob: bytes = database_blob + else: + blob = cached_blob + timing_enabled = _snapshot_timing_enabled() + decode_timings: dict[str, float] | None = {} if timing_enabled else None + decode_started = time.perf_counter() + try: + payload = decode_namespace_map_snapshot( + blob, + checksum=str(checksum), + format_version=int(format_version), + timings=decode_timings, + ) + except (ValueError, TypeError, zlib.error): + if cached_blob is not None: + # A stale/corrupt cache entry must never shadow the PostgreSQL source. + try: + if database_blob is None: + fallback_result = await db.execute( + select(RetrievalNamespaceMapSnapshot.payload_zlib).where( + RetrievalNamespaceMapSnapshot.user_id == user_id, + RetrievalNamespaceMapSnapshot.namespace == namespace, + ) + ) + fallback_row = fallback_result.first() + if fallback_row is None or fallback_row[0] is None: + return None + database_blob = bytes(fallback_row[0]) + assert database_blob is not None + payload = decode_namespace_map_snapshot( + database_blob, + checksum=str(checksum), + format_version=int(format_version), + timings=decode_timings, + ) + blob = database_blob + cached_blob = None + except (ValueError, TypeError, zlib.error): + return None + else: return None - documents = decoded_documents - cache_namespace_documents( - user_id=user_id, - namespace=namespace, - generation=int(generation), - documents=documents, + decoded_documents = payload.get("documents") + if not isinstance(decoded_documents, dict): + return None + documents = decoded_documents + if cached_blob is None: + try: + await set_snapshot_blob( + user_id=user_id, + namespace=namespace, + generation=row_generation, + payload_zlib=blob, + ) + except Exception as exc: + _logger.warning("retrieval snapshot redis set failed error=%s", exc) + if timing_enabled: + _logger.info( + "retrieval snapshot decode cache_hit=false seconds=%.3f " + "compressed_bytes=%d decompressed_bytes=%d " + "decompress_seconds=%.3f checksum_seconds=%.3f " + "json_decode_seconds=%.3f documents=%d", + time.perf_counter() - decode_started, + int((decode_timings or {}).get("compressed_bytes", 0.0)), + int((decode_timings or {}).get("decompressed_bytes", 0.0)), + (decode_timings or {}).get("decompress_seconds", 0.0), + (decode_timings or {}).get("checksum_seconds", 0.0), + (decode_timings or {}).get("json_decode_seconds", 0.0), + len(documents), + ) + elif timing_enabled: + _logger.info( + "retrieval snapshot decode cache_hit=true seconds=%.3f " + "compressed_bytes=%d decompressed_bytes=%d decompress_seconds=%.3f " + "checksum_seconds=%.3f json_decode_seconds=%.3f documents=%d", + time.perf_counter() - decode_started, + int((decode_timings or {}).get("compressed_bytes", 0.0)), + int((decode_timings or {}).get("decompressed_bytes", 0.0)), + (decode_timings or {}).get("decompress_seconds", 0.0), + (decode_timings or {}).get("checksum_seconds", 0.0), + (decode_timings or {}).get("json_decode_seconds", 0.0), + len(documents), ) entries: list[tuple[object, ...]] = [] for document_id, job_result_id in document_revisions: @@ -532,6 +702,11 @@ def _parse_manifest_entries( 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]]] = {} + section_seconds = 0.0 + chunk_seconds = 0.0 + section_count = 0 + chunk_count = 0 + connection_count = 0 try: for document_id, _job_result_id, payload_zlib, checksum, format_version in manifest_entries: if isinstance(payload_zlib, dict): @@ -546,6 +721,7 @@ def _parse_manifest_entries( checksum=str(checksum), format_version=int(str(format_version)), ) + section_started = time.perf_counter() raw_sections = payload.get("sections") if not isinstance(raw_sections, list): return None @@ -577,6 +753,9 @@ def _parse_manifest_entries( ) by_doc.setdefault(str(document_id), []).append(section) path_by_id[section_id] = section_path + section_count += 1 + section_seconds += time.perf_counter() - section_started + chunk_started = time.perf_counter() raw_chunks = payload.get("chunks") if not isinstance(raw_chunks, list): return None @@ -628,15 +807,32 @@ def _parse_manifest_entries( text_connections_by_doc.setdefault(document_key, []).append( (section_id or "", target) ) + connection_count += 1 + chunk_count += 1 + chunk_seconds += time.perf_counter() - chunk_started except (TypeError, ValueError, KeyError): return None remounted: dict[str, dict[str, Any]] = {} + remount_started = time.perf_counter() 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} + if _snapshot_timing_enabled(): + _logger.info( + "retrieval snapshot parse sections=%d chunks=%d connections=%d " + "section_seconds=%.3f chunk_seconds=%.3f remount_seconds=%.3f " + "ref_index_keys=%d", + section_count, + chunk_count, + connection_count, + section_seconds, + chunk_seconds, + time.perf_counter() - remount_started, + len(ref_index), + ) return by_doc, path_by_id, (ids_by_doc, ref_index, remounted) diff --git a/packages/shared-python/shared/services/retrieval/serving_manifest.py b/packages/shared-python/shared/services/retrieval/serving_manifest.py index c66ae5798..9895ca434 100644 --- a/packages/shared-python/shared/services/retrieval/serving_manifest.py +++ b/packages/shared-python/shared/services/retrieval/serving_manifest.py @@ -4,8 +4,9 @@ import hashlib import json +import time import zlib -from typing import Any +from typing import Any, MutableMapping from sqlalchemy import delete, select from sqlalchemy.orm import Session @@ -20,6 +21,7 @@ from shared.services.retrieval.publication_models import DocumentPublicationScope SERVING_MANIFEST_FORMAT_VERSION = 1 +NAMESPACE_MAP_SNAPSHOT_FORMAT_VERSION = 2 def build_revision_serving_payload( @@ -181,24 +183,141 @@ def decode_serving_manifest( *, checksum: str, format_version: int, + timings: MutableMapping[str, float] | None = None, ) -> 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}") + started = time.perf_counter() try: canonical_payload = zlib.decompress(payload_zlib) except zlib.error as exc: raise ValueError("invalid serving manifest compression") from exc + if timings is not None: + timings["decompress_seconds"] = time.perf_counter() - started + checksum_started = time.perf_counter() actual_checksum = hashlib.sha256(canonical_payload).hexdigest() if actual_checksum != checksum: raise ValueError("serving manifest checksum mismatch") + if timings is not None: + timings["checksum_seconds"] = time.perf_counter() - checksum_started + json_started = time.perf_counter() try: decoded = json.loads(canonical_payload.decode("utf-8")) except (UnicodeDecodeError, json.JSONDecodeError) as exc: raise ValueError("invalid serving manifest JSON") from exc + if timings is not None: + timings["json_decode_seconds"] = time.perf_counter() - json_started + timings["compressed_bytes"] = float(len(payload_zlib)) + timings["decompressed_bytes"] = float(len(canonical_payload)) + timings["decode_seconds"] = time.perf_counter() - started if not isinstance(decoded, dict): raise ValueError("serving manifest payload must be an object") return decoded + + +def encode_namespace_map_snapshot( + payload: dict[str, Any], +) -> tuple[bytes, str, int]: + """Encode the routing-only namespace snapshot using its own format version.""" + documents = payload.get("documents") + routing_documents: dict[str, Any] = {} + if isinstance(documents, dict): + for document_id, raw_document in documents.items(): + if not isinstance(raw_document, dict): + continue + sections = [ + { + key: section[key] + for key in ( + "section_id", + "parent_section_id", + "section_path", + "section_title", + "section_level", + "summary", + "sort_order", + ) + if key in section + } + for section in (raw_document.get("sections") or []) + if isinstance(section, dict) + ] + chunks = [ + { + key: chunk[key] + for key in ( + "chunk_id", + "section_id", + "chunk_type", + "sort_order", + "connect_to", + ) + if key in chunk + } + for chunk in (raw_document.get("chunks") or []) + if isinstance(chunk, dict) + ] + routing_documents[str(document_id)] = { + "job_result_id": raw_document.get("job_result_id"), + "job_id": raw_document.get("job_id"), + "sections": sections, + "chunks": chunks, + "root_asset_ids": raw_document.get("root_asset_ids") or [], + "remounted_assets_by_section": raw_document.get( + "remounted_assets_by_section" + ) + or {}, + } + compressed, checksum, _ = encode_serving_manifest({"documents": routing_documents}) + return compressed, checksum, NAMESPACE_MAP_SNAPSHOT_FORMAT_VERSION + + +def decode_namespace_map_snapshot( + payload_zlib: bytes, + *, + checksum: str, + format_version: int, + timings: MutableMapping[str, float] | None = None, +) -> dict[str, Any]: + """Decode namespace snapshots, retaining compatibility with legacy v1 rows.""" + if format_version == SERVING_MANIFEST_FORMAT_VERSION: + return decode_serving_manifest( + payload_zlib, + checksum=checksum, + format_version=SERVING_MANIFEST_FORMAT_VERSION, + timings=timings, + ) + if format_version != NAMESPACE_MAP_SNAPSHOT_FORMAT_VERSION: + raise ValueError(f"unsupported namespace snapshot version: {format_version}") + # The compression/checksum contract is identical to serving manifests; only + # the payload shape/version is different. + started = time.perf_counter() + try: + canonical_payload = zlib.decompress(payload_zlib) + except zlib.error as exc: + raise ValueError("invalid namespace snapshot compression") from exc + if timings is not None: + timings["decompress_seconds"] = time.perf_counter() - started + timings["compressed_bytes"] = float(len(payload_zlib)) + timings["decompressed_bytes"] = float(len(canonical_payload)) + checksum_started = time.perf_counter() + actual_checksum = hashlib.sha256(canonical_payload).hexdigest() + if actual_checksum != checksum: + raise ValueError("namespace snapshot checksum mismatch") + if timings is not None: + timings["checksum_seconds"] = time.perf_counter() - checksum_started + json_started = time.perf_counter() + try: + decoded = json.loads(canonical_payload.decode("utf-8")) + except (UnicodeDecodeError, json.JSONDecodeError) as exc: + raise ValueError("invalid namespace snapshot JSON") from exc + if timings is not None: + timings["json_decode_seconds"] = time.perf_counter() - json_started + timings["decode_seconds"] = time.perf_counter() - started + if not isinstance(decoded, dict): + raise ValueError("namespace snapshot payload must be an object") + return decoded From 25c074c70e247ac398055681b669ffcb5bc3d07e Mon Sep 17 00:00:00 2001 From: suguanYang Date: Tue, 1 Sep 2026 21:55:21 +0800 Subject: [PATCH 2/5] fix: enforce snapshot generation and validation --- .../retrieval/namespace_map_snapshot_redis.py | 6 +- .../shared/services/retrieval/nav_snapshot.py | 37 +++-- .../services/retrieval/serving_manifest.py | 140 +++++++++--------- 3 files changed, 99 insertions(+), 84 deletions(-) diff --git a/packages/shared-python/shared/services/retrieval/namespace_map_snapshot_redis.py b/packages/shared-python/shared/services/retrieval/namespace_map_snapshot_redis.py index fc8447c6a..b977e5f45 100644 --- a/packages/shared-python/shared/services/retrieval/namespace_map_snapshot_redis.py +++ b/packages/shared-python/shared/services/retrieval/namespace_map_snapshot_redis.py @@ -9,7 +9,7 @@ _KEY_PREFIX = "retrieval:snapshot:v2" -def snapshot_cache_key(*, user_id: str, namespace: str, generation: int) -> str: +def build_snapshot_cache_key(*, user_id: str, namespace: str, generation: int) -> str: normalized_namespace = normalize_retrieval_namespace(namespace) return f"{_KEY_PREFIX}:{user_id}:{normalized_namespace}:g{int(generation)}" @@ -19,7 +19,7 @@ async def get_snapshot_blob( ) -> bytes | None: service = RedisServiceFactory.get_service() return await service.get_bytes( - snapshot_cache_key( + build_snapshot_cache_key( user_id=user_id, namespace=namespace, generation=generation ) ) @@ -30,7 +30,7 @@ async def set_snapshot_blob( ) -> bool: service = RedisServiceFactory.get_service() return await service.set_bytes( - snapshot_cache_key( + build_snapshot_cache_key( user_id=user_id, namespace=namespace, generation=generation ), payload_zlib, diff --git a/packages/shared-python/shared/services/retrieval/nav_snapshot.py b/packages/shared-python/shared/services/retrieval/nav_snapshot.py index 50a4b4f89..3210b55f7 100644 --- a/packages/shared-python/shared/services/retrieval/nav_snapshot.py +++ b/packages/shared-python/shared/services/retrieval/nav_snapshot.py @@ -58,6 +58,7 @@ get_snapshot_blob, set_snapshot_blob, ) +from shared.core.exceptions.redis_exceptions import RedisOperationError # Keep each payload query bounded under the API's 30-second statement timeout. @@ -76,7 +77,7 @@ _logger = logging.getLogger(__name__) -def _snapshot_timing_enabled() -> bool: +def _is_snapshot_timing_enabled() -> bool: """Enable detailed snapshot timings for local performance diagnosis.""" return os.environ.get("RETRIEVAL_SNAPSHOT_TIMING", "0").strip().lower() in { "1", @@ -368,7 +369,7 @@ async def load_nav_snapshot( }, document_revisions=dict(revisions), ) - if _snapshot_timing_enabled(): + if _is_snapshot_timing_enabled(): _logger.info( "retrieval snapshot timing total_seconds=%.3f parse_seconds=%.3f " "ref_index_copy_seconds=%.3f doc_query_seconds=%.3f " @@ -413,7 +414,7 @@ async def load_nav_snapshot( if document_id in kept_titles }, ) - if _snapshot_timing_enabled(): + if _is_snapshot_timing_enabled(): _logger.info( "retrieval snapshot timing total_seconds=%.3f parse_seconds=%.3f " "doc_query_seconds=%.3f job_query_seconds=%.3f documents=%d " @@ -453,7 +454,7 @@ async def _resolve_namespace_snapshot_entries( await db.rollback() _logger.warning("retrieval snapshot generation lookup failed error=%s", exc) _current_generation = None - if _snapshot_timing_enabled(): + if _is_snapshot_timing_enabled(): _logger.info( "retrieval snapshot generation generation=%s", _current_generation, @@ -480,7 +481,7 @@ async def _resolve_namespace_snapshot_entries( ) return None lookup_seconds = time.perf_counter() - lookup_started - if _snapshot_timing_enabled(): + if _is_snapshot_timing_enabled(): _logger.info( "retrieval snapshot lookup seconds=%.3f found=%s", lookup_seconds, @@ -497,12 +498,24 @@ async def _resolve_namespace_snapshot_entries( generation_value, ) return None + if ( + generation_value is None + and _current_generation is not None + and int(_current_generation) > 0 + and row_generation != int(_current_generation) + ): + _logger.info( + "retrieval snapshot generation mismatch row=%d current=%d", + row_generation, + int(_current_generation), + ) + return None cached_blob: bytes | None = None try: cached_blob = await get_snapshot_blob( user_id=user_id, namespace=namespace, generation=row_generation ) - except Exception as exc: + except RedisOperationError as exc: _logger.warning("retrieval snapshot redis get failed error=%s", exc) database_blob: bytes | None = None if cached_blob is None: @@ -519,8 +532,8 @@ async def _resolve_namespace_snapshot_entries( blob: bytes = database_blob else: blob = cached_blob - timing_enabled = _snapshot_timing_enabled() - decode_timings: dict[str, float] | None = {} if timing_enabled else None + is_timing_enabled = _is_snapshot_timing_enabled() + decode_timings: dict[str, float] | None = {} if is_timing_enabled else None decode_started = time.perf_counter() try: payload = decode_namespace_map_snapshot( @@ -569,9 +582,9 @@ async def _resolve_namespace_snapshot_entries( generation=row_generation, payload_zlib=blob, ) - except Exception as exc: + except RedisOperationError as exc: _logger.warning("retrieval snapshot redis set failed error=%s", exc) - if timing_enabled: + if is_timing_enabled: _logger.info( "retrieval snapshot decode cache_hit=false seconds=%.3f " "compressed_bytes=%d decompressed_bytes=%d " @@ -585,7 +598,7 @@ async def _resolve_namespace_snapshot_entries( (decode_timings or {}).get("json_decode_seconds", 0.0), len(documents), ) - elif timing_enabled: + elif is_timing_enabled: _logger.info( "retrieval snapshot decode cache_hit=true seconds=%.3f " "compressed_bytes=%d decompressed_bytes=%d decompress_seconds=%.3f " @@ -820,7 +833,7 @@ def _parse_manifest_entries( if target in asset_ids: owners.setdefault(section_id, []).append(target) remounted[document_id] = {"root": sorted(asset_ids), "owners": owners} - if _snapshot_timing_enabled(): + if _is_snapshot_timing_enabled(): _logger.info( "retrieval snapshot parse sections=%d chunks=%d connections=%d " "section_seconds=%.3f chunk_seconds=%.3f remount_seconds=%.3f " diff --git a/packages/shared-python/shared/services/retrieval/serving_manifest.py b/packages/shared-python/shared/services/retrieval/serving_manifest.py index 9895ca434..8480b8f49 100644 --- a/packages/shared-python/shared/services/retrieval/serving_manifest.py +++ b/packages/shared-python/shared/services/retrieval/serving_manifest.py @@ -189,18 +189,41 @@ def decode_serving_manifest( if format_version != SERVING_MANIFEST_FORMAT_VERSION: raise ValueError(f"unsupported serving manifest version: {format_version}") + return _decode_compressed_json( + payload_zlib, + checksum=checksum, + timings=timings, + compression_error="invalid serving manifest compression", + checksum_error="serving manifest checksum mismatch", + json_error="invalid serving manifest JSON", + object_error="serving manifest payload must be an object", + ) + + +def _decode_compressed_json( + payload_zlib: bytes, + *, + checksum: str, + timings: MutableMapping[str, float] | None, + compression_error: str, + checksum_error: str, + json_error: str, + object_error: str, +) -> dict[str, Any]: + """Decompress, validate, and decode a canonical JSON payload.""" + started = time.perf_counter() try: canonical_payload = zlib.decompress(payload_zlib) except zlib.error as exc: - raise ValueError("invalid serving manifest compression") from exc + raise ValueError(compression_error) from exc if timings is not None: timings["decompress_seconds"] = time.perf_counter() - started checksum_started = time.perf_counter() actual_checksum = hashlib.sha256(canonical_payload).hexdigest() if actual_checksum != checksum: - raise ValueError("serving manifest checksum mismatch") + raise ValueError(checksum_error) if timings is not None: timings["checksum_seconds"] = time.perf_counter() - checksum_started @@ -208,14 +231,14 @@ def decode_serving_manifest( try: decoded = json.loads(canonical_payload.decode("utf-8")) except (UnicodeDecodeError, json.JSONDecodeError) as exc: - raise ValueError("invalid serving manifest JSON") from exc + raise ValueError(json_error) from exc if timings is not None: timings["json_decode_seconds"] = time.perf_counter() - json_started timings["compressed_bytes"] = float(len(payload_zlib)) timings["decompressed_bytes"] = float(len(canonical_payload)) timings["decode_seconds"] = time.perf_counter() - started if not isinstance(decoded, dict): - raise ValueError("serving manifest payload must be an object") + raise ValueError(object_error) return decoded @@ -224,54 +247,52 @@ def encode_namespace_map_snapshot( ) -> tuple[bytes, str, int]: """Encode the routing-only namespace snapshot using its own format version.""" documents = payload.get("documents") + if not isinstance(documents, dict): + raise ValueError("namespace snapshot documents must be an object") routing_documents: dict[str, Any] = {} - if isinstance(documents, dict): - for document_id, raw_document in documents.items(): - if not isinstance(raw_document, dict): - continue - sections = [ + for document_id, raw_document in documents.items(): + if not isinstance(raw_document, dict): + raise ValueError(f"namespace snapshot document is not an object: {document_id}") + raw_sections = raw_document.get("sections") + raw_chunks = raw_document.get("chunks") + if not isinstance(raw_sections, list) or not isinstance(raw_chunks, list): + raise ValueError(f"namespace snapshot records are invalid: {document_id}") + sections = [] + for section in raw_sections: + if not isinstance(section, dict) or not str(section.get("section_id") or ""): + raise ValueError(f"namespace snapshot section is invalid: {document_id}") + sections.append( { key: section[key] for key in ( - "section_id", - "parent_section_id", - "section_path", - "section_title", - "section_level", - "summary", - "sort_order", + "section_id", "parent_section_id", "section_path", + "section_title", "section_level", "summary", "sort_order", ) if key in section } - for section in (raw_document.get("sections") or []) - if isinstance(section, dict) - ] - chunks = [ + ) + chunks = [] + for chunk in raw_chunks: + if not isinstance(chunk, dict) or not str(chunk.get("chunk_id") or ""): + raise ValueError(f"namespace snapshot chunk is invalid: {document_id}") + chunks.append( { key: chunk[key] - for key in ( - "chunk_id", - "section_id", - "chunk_type", - "sort_order", - "connect_to", - ) + for key in ("chunk_id", "section_id", "chunk_type", "sort_order", "connect_to") if key in chunk } - for chunk in (raw_document.get("chunks") or []) - if isinstance(chunk, dict) - ] - routing_documents[str(document_id)] = { - "job_result_id": raw_document.get("job_result_id"), - "job_id": raw_document.get("job_id"), - "sections": sections, - "chunks": chunks, - "root_asset_ids": raw_document.get("root_asset_ids") or [], - "remounted_assets_by_section": raw_document.get( - "remounted_assets_by_section" - ) - or {}, - } + ) + routing_documents[str(document_id)] = { + "job_result_id": raw_document.get("job_result_id"), + "job_id": raw_document.get("job_id"), + "sections": sections, + "chunks": chunks, + "root_asset_ids": raw_document.get("root_asset_ids") or [], + "remounted_assets_by_section": raw_document.get( + "remounted_assets_by_section" + ) + or {}, + } compressed, checksum, _ = encode_serving_manifest({"documents": routing_documents}) return compressed, checksum, NAMESPACE_MAP_SNAPSHOT_FORMAT_VERSION @@ -293,31 +314,12 @@ def decode_namespace_map_snapshot( ) if format_version != NAMESPACE_MAP_SNAPSHOT_FORMAT_VERSION: raise ValueError(f"unsupported namespace snapshot version: {format_version}") - # The compression/checksum contract is identical to serving manifests; only - # the payload shape/version is different. - started = time.perf_counter() - try: - canonical_payload = zlib.decompress(payload_zlib) - except zlib.error as exc: - raise ValueError("invalid namespace snapshot compression") from exc - if timings is not None: - timings["decompress_seconds"] = time.perf_counter() - started - timings["compressed_bytes"] = float(len(payload_zlib)) - timings["decompressed_bytes"] = float(len(canonical_payload)) - checksum_started = time.perf_counter() - actual_checksum = hashlib.sha256(canonical_payload).hexdigest() - if actual_checksum != checksum: - raise ValueError("namespace snapshot checksum mismatch") - if timings is not None: - timings["checksum_seconds"] = time.perf_counter() - checksum_started - json_started = time.perf_counter() - try: - decoded = json.loads(canonical_payload.decode("utf-8")) - except (UnicodeDecodeError, json.JSONDecodeError) as exc: - raise ValueError("invalid namespace snapshot JSON") from exc - if timings is not None: - timings["json_decode_seconds"] = time.perf_counter() - json_started - timings["decode_seconds"] = time.perf_counter() - started - if not isinstance(decoded, dict): - raise ValueError("namespace snapshot payload must be an object") - return decoded + return _decode_compressed_json( + payload_zlib, + checksum=checksum, + timings=timings, + compression_error="invalid namespace snapshot compression", + checksum_error="namespace snapshot checksum mismatch", + json_error="invalid namespace snapshot JSON", + object_error="namespace snapshot payload must be an object", + ) From 1223551ffeabf73fe569df5080e9f69d12792b31 Mon Sep 17 00:00:00 2001 From: suguanYang Date: Tue, 1 Sep 2026 22:38:24 +0800 Subject: [PATCH 3/5] fix: harden map-nav snapshot cache --- ...est_retrieval_map_score_parity_contract.py | 88 +++++++++ ...retrieval_snapshot_consistency_contract.py | 28 +++ .../test_retrieval_snapshot_redis_contract.py | 181 ++++++++++++++++++ .../retrieval-serving-snapshot-cache.md | 48 +++++ .../shared/services/redis/redis_service.py | 34 +++- .../services/retrieval/execution/routes.py | 3 +- .../retrieval/namespace_map_snapshot_redis.py | 56 +++--- .../shared/services/retrieval/nav_snapshot.py | 14 +- .../services/retrieval/serving_manifest.py | 6 +- 9 files changed, 417 insertions(+), 41 deletions(-) create mode 100644 apps/api/tests/contract/test_retrieval_map_score_parity_contract.py create mode 100644 apps/api/tests/contract/test_retrieval_snapshot_redis_contract.py create mode 100644 docs/design/retrieval-serving-snapshot-cache.md diff --git a/apps/api/tests/contract/test_retrieval_map_score_parity_contract.py b/apps/api/tests/contract/test_retrieval_map_score_parity_contract.py new file mode 100644 index 000000000..c7d651c07 --- /dev/null +++ b/apps/api/tests/contract/test_retrieval_map_score_parity_contract.py @@ -0,0 +1,88 @@ +"""Contract tests for map score pooling semantics.""" + +from __future__ import annotations + +from typing import Final + +import pytest + +from shared.services.retrieval.nav.nav_map_scores import _pool_unit_scores_to_tree + + +_CASES: Final[tuple[tuple[dict[str, list[str]], set[str], dict[str, float]], ...]] = ( + ( + {"root-a": ["section-a", "section-b"], "section-a": [], "section-b": []}, + {"section-a", "section-b"}, + {"section-a": 0.4, "section-b": 0.8, "root-a__self": 0.2}, + ), + ( + { + "root-a": ["parent-a"], + "parent-a": ["leaf-a", "leaf-b"], + "leaf-a": [], + "leaf-b": [], + "root-b": ["leaf-c"], + "leaf-c": [], + }, + {"leaf-a", "leaf-b", "leaf-c"}, + { + "leaf-a": 0.9, + "leaf-b": 0.3, + "leaf-c": 0.7, + "parent-a__self": 0.95, + }, + ), +) + + +def _legacy_pool( + children_map: dict[str, list[str]], + leaves: set[str], + unit_scores: dict[str, float], +) -> dict[str, float]: + map_scores = { + leaf_id: float(unit_scores.get(leaf_id, 0.0) or 0.0) for leaf_id in leaves + } + + def score_node(section_id: str) -> float: + if section_id in map_scores: + return map_scores[section_id] + children = children_map.get(section_id) or [] + if not children: + score = float(unit_scores.get(section_id, 0.0) or 0.0) + map_scores[section_id] = score + return score + descendants: list[str] = [] + + def collect(section: str) -> None: + nested = children_map.get(section) or [] + if not nested: + if section in leaves: + descendants.append(section) + return + for child in nested: + collect(child) + + collect(section_id) + parts = [float(unit_scores.get(leaf_id, 0.0) or 0.0) for leaf_id in descendants] + self_key = f"{section_id}__self" + if self_key in unit_scores: + parts.append(float(unit_scores[self_key])) + score = float(max(parts)) if parts else 0.0 + map_scores[section_id] = score + return score + + for section_id in children_map: + score_node(section_id) + return map_scores + + +@pytest.mark.parametrize("children_map, leaves, unit_scores", _CASES) +def test_map_score_pooling_preserves_legacy_semantics( + children_map: dict[str, list[str]], + leaves: set[str], + unit_scores: dict[str, float], +) -> None: + assert _pool_unit_scores_to_tree(children_map, leaves, unit_scores) == _legacy_pool( + children_map, leaves, unit_scores + ) 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 bbe03e059..a30af3c31 100644 --- a/apps/api/tests/contract/test_retrieval_snapshot_consistency_contract.py +++ b/apps/api/tests/contract/test_retrieval_snapshot_consistency_contract.py @@ -8,11 +8,13 @@ from httpx import AsyncClient import pytest from sqlalchemy import Executable, Result +from sqlalchemy.exc import SQLAlchemyError from shared.services.retrieval.execution.reference_resolver import ( resolve_workflow_references, ) from shared.services.retrieval.nav_snapshot import SnapshotSession, load_nav_snapshot +from shared.services.retrieval.nav_snapshot import _resolve_namespace_snapshot_entries from tests.support.retrieval_snapshot_support import contract_db_session from tests.support.contract_database import ContractDatabase @@ -20,6 +22,32 @@ _USER_ID = "local-dev-user" +class _GenerationUnavailableSession: + def __init__(self) -> None: + self.rollback_count = 0 + + async def execute(self, _statement: Executable) -> Result[tuple[object, ...]]: + raise SQLAlchemyError("generation table unavailable") + + async def rollback(self) -> None: + self.rollback_count += 1 + + +@pytest.mark.asyncio +async def test_snapshot_loader_falls_back_when_generation_cannot_be_verified() -> None: + session = _GenerationUnavailableSession() + + result = await _resolve_namespace_snapshot_entries( + session, + user_id=_USER_ID, + namespace="default", + document_revisions=[("doc-a", "result-a")], + ) + + assert result is None + assert session.rollback_count == 1 + + class _PublishingSession: def __init__( self, diff --git a/apps/api/tests/contract/test_retrieval_snapshot_redis_contract.py b/apps/api/tests/contract/test_retrieval_snapshot_redis_contract.py new file mode 100644 index 000000000..299341336 --- /dev/null +++ b/apps/api/tests/contract/test_retrieval_snapshot_redis_contract.py @@ -0,0 +1,181 @@ +"""Contract tests for binary namespace snapshot caching.""" + +from __future__ import annotations + +import pytest +from sqlalchemy import Executable + +from shared.core.config.redis import RedisConfig, RedisConfigManager +from shared.services.redis.redis_service_factory import RedisServiceFactory +from shared.services.redis.redis_service import RedisService +from shared.services.retrieval.nav_snapshot import _resolve_namespace_snapshot_entries +from shared.services.retrieval.namespace_map_snapshot_redis import ( + NamespaceMapSnapshotRedisCache, +) +from shared.services.retrieval.serving_manifest import encode_namespace_map_snapshot + + +class _FakeRedisClient: + def __init__(self, *, decode_responses: bool) -> None: + self.decode_responses = decode_responses + self.values: dict[str, bytes] = {} + self.ttls: dict[str, int] = {} + + async def get(self, key: str) -> bytes | None: + return self.values.get(key) + + async def set(self, key: str, value: bytes, *, ex: int) -> bool: + self.values[key] = value + self.ttls[key] = ex + return True + + async def aclose(self) -> None: + return None + + +@pytest.mark.asyncio +async def test_binary_redis_operations_preserve_compressed_snapshot_bytes( + monkeypatch: pytest.MonkeyPatch, +) -> None: + clients: list[_FakeRedisClient] = [] + + def create_client(*args: object, **kwargs: object) -> _FakeRedisClient: + client = _FakeRedisClient( + decode_responses=bool(kwargs.get("decode_responses")) + ) + clients.append(client) + return client + + monkeypatch.setattr( + "shared.services.redis.redis_service.redis.from_url", create_client + ) + service = RedisService(RedisConfigManager(RedisConfig())) + payload = b"\x78\x9c\x00\xffcompressed-snapshot" + + assert await service.set_bytes("contract:snapshot", payload, ex=3600) + assert await service.get_bytes("contract:snapshot") == payload + assert len(clients) == 1 + assert clients[0].decode_responses is False + assert clients[0].ttls["knowhere-api:contract:snapshot"] == 3600 + + await service.close() + + +@pytest.mark.asyncio +async def test_snapshot_cache_scopes_reads_and_writes_by_generation( + monkeypatch: pytest.MonkeyPatch, +) -> None: + class _FakeSnapshotService: + def __init__(self) -> None: + self.get_keys: list[str] = [] + self.set_calls: list[tuple[str, bytes, int]] = [] + + async def get_bytes(self, key: str) -> bytes | None: + self.get_keys.append(key) + return b"snapshot" + + async def set_bytes(self, key: str, value: bytes, *, ex: int) -> bool: + self.set_calls.append((key, value, ex)) + return True + + fake_service = _FakeSnapshotService() + monkeypatch.setattr( + RedisServiceFactory, + "get_service", + classmethod(lambda cls: fake_service), + ) + + assert ( + await NamespaceMapSnapshotRedisCache.get( + user_id="user", + namespace=" ", + generation=7, + ) + == b"snapshot" + ) + assert await NamespaceMapSnapshotRedisCache.set( + user_id="user", + namespace="default", + generation=8, + payload_zlib=b"compressed", + ) + + assert fake_service.get_keys == ["retrieval:snapshot:v2:user:default:g7"] + assert fake_service.set_calls == [ + ( + "retrieval:snapshot:v2:user:default:g8", + b"compressed", + 3600, + ) + ] + + +class _SnapshotResult: + def __init__(self, *, scalar: object = None, row: tuple[object, ...] | None = None): + self._scalar = scalar + self._row = row + + def scalar_one_or_none(self) -> object: + return self._scalar + + def first(self) -> tuple[object, ...] | None: + return self._row + + +class _SnapshotSequenceSession: + def __init__(self, results: list[_SnapshotResult]) -> None: + self._results = iter(results) + + async def execute(self, _statement: Executable) -> _SnapshotResult: + return next(self._results) + + async def rollback(self) -> None: + return None + + +@pytest.mark.asyncio +async def test_corrupt_redis_blob_falls_back_to_postgres_snapshot( + monkeypatch: pytest.MonkeyPatch, +) -> None: + payload_zlib, checksum, format_version = encode_namespace_map_snapshot( + { + "documents": { + "doc-a": { + "job_result_id": "result-a", + "job_id": "job-a", + "sections": [], + "chunks": [], + } + } + } + ) + set_calls: list[bytes] = [] + + async def get_corrupt_blob(**_: object) -> bytes: + return b"corrupt" + + async def record_repaired_blob(**kwargs: object) -> bool: + set_calls.append(bytes(kwargs["payload_zlib"])) + return True + + monkeypatch.setattr(NamespaceMapSnapshotRedisCache, "get", get_corrupt_blob) + monkeypatch.setattr(NamespaceMapSnapshotRedisCache, "set", record_repaired_blob) + session = _SnapshotSequenceSession( + [ + _SnapshotResult(scalar=3), + _SnapshotResult(row=(3, checksum, format_version)), + _SnapshotResult(row=(payload_zlib,)), + ] + ) + + entries = await _resolve_namespace_snapshot_entries( + session, + user_id="user", + namespace="default", + document_revisions=[("doc-a", "result-a")], + expected_generation=3, + ) + + assert entries is not None + assert entries[0][0:2] == ("doc-a", "result-a") + assert set_calls == [payload_zlib] diff --git a/docs/design/retrieval-serving-snapshot-cache.md b/docs/design/retrieval-serving-snapshot-cache.md new file mode 100644 index 000000000..36f23f9df --- /dev/null +++ b/docs/design/retrieval-serving-snapshot-cache.md @@ -0,0 +1,48 @@ +# Retrieval serving snapshot cache + +Map-nav reads a namespace routing snapshot that contains section/chunk +relationships, ordering, chunk types, connection references, and remount +ownership. It intentionally does not contain chunk body text; final hydration +still reads pinned revision content from PostgreSQL. + +## Persisted formats + +- Namespace snapshots written by current publication code use format version 2 + and contain routing metadata only. +- The reader remains compatible with version 1 snapshots so existing rows can + be served during rollout. No unconditional namespace rebuild is required + solely because the reader was upgraded. +- A missing, corrupt, stale, or generation-mismatched snapshot falls back to + the exact manifest/table loading path. + +## Redis cache + +The map-nav reader caches the compressed snapshot bytes, not a decoded Python +dictionary. Keys are scoped by user, normalized namespace, and serving +generation: + +```text +retrieval:snapshot:v2:{user_id}:{namespace}:g{generation} +``` + +The cache TTL is one hour. Redis errors, misses, and invalid blobs are +non-fatal: PostgreSQL remains the source of truth and the reader repopulates +Redis after a successful database read. Binary snapshot operations use a Redis +connection with response decoding disabled; normal JSON Redis operations keep +their existing text-decoding connection. + +## Generation coherence + +Retrieval carries one captured revision set and namespace generation through +snapshot loading, map-nav, reference resolution, and final hydration. A +generation lookup failure is treated as inability to establish coherence and +uses the fallback path. A generation mismatch is never served as a valid +snapshot. + +## Diagnostics + +Set `RETRIEVAL_SNAPSHOT_TIMING=1` for detailed snapshot decode timings during +local or staging diagnosis. The route CPU metric is process-level CPU time for +the map-nav route. `ru_maxrss` is a process high-water mark, not a request-level +memory peak, and must not be interpreted as one. + diff --git a/packages/shared-python/shared/services/redis/redis_service.py b/packages/shared-python/shared/services/redis/redis_service.py index 185467ecf..c25a3e1d6 100644 --- a/packages/shared-python/shared/services/redis/redis_service.py +++ b/packages/shared-python/shared/services/redis/redis_service.py @@ -38,6 +38,7 @@ def __init__(self, config_manager: Optional[RedisConfigManager] = None): config_manager = RedisConfigManager(settings) self.config_manager = config_manager self._client: Optional[redis.Redis] = None + self._binary_client: Optional[redis.Redis] = None self._health_checker: Optional[RedisHealthChecker] = None self._lock = asyncio.Lock() @@ -60,6 +61,26 @@ async def _get_client(self) -> redis.Redis: ) return self._client + async def _get_binary_client(self) -> redis.Redis: + """Get a Redis client that preserves arbitrary binary responses.""" + if self._binary_client is None: + async with self._lock: + if self._binary_client is None: + try: + connection_params = self.config_manager.get_connection_params() + connection_params["decode_responses"] = False + self._binary_client = redis.from_url( + self.config_manager.get_connection_url(), + **connection_params, + ) + logger.debug("Redis binary client initialized") + except Exception as e: + raise RedisConnectionError( + internal_message=f"Redis binary client initialization failed: {str(e)}", + original_exception=e, + ) + return self._binary_client + async def _execute_with_retry( self, operation: Callable[[], Awaitable[ResponseT]] ) -> ResponseT: @@ -163,7 +184,7 @@ async def _operation(): async def get_bytes(self, key: str) -> bytes | None: """Get a value without JSON decoding (for compressed binary blobs).""" try: - client = await self._get_client() + client = await self._get_binary_client() full_key = self._build_key(key) async def _operation() -> bytes | None: @@ -188,7 +209,7 @@ async def _operation() -> bytes | None: async def set_bytes(self, key: str, value: bytes, *, ex: int) -> bool: """Set a compressed binary value with an explicit TTL.""" try: - client = await self._get_client() + client = await self._get_binary_client() full_key = self._build_key(key) async def _operation() -> bool: @@ -651,18 +672,21 @@ async def is_healthy(self) -> bool: async def close(self): """Close the Redis connection.""" - if self._client: + clients = [client for client in (self._client, self._binary_client) if client] + for client in clients: close_client = cast( Callable[[], Awaitable[None]] | None, - getattr(self._client, "aclose", None), + getattr(client, "aclose", None), ) if close_client is not None: await close_client() else: - await self._client.close() + await client.close() + if clients: self._client = None + self._binary_client = None self._health_checker = None logger.info("Redis connection closed") diff --git a/packages/shared-python/shared/services/retrieval/execution/routes.py b/packages/shared-python/shared/services/retrieval/execution/routes.py index ca9809a45..dc16c0f16 100644 --- a/packages/shared-python/shared/services/retrieval/execution/routes.py +++ b/packages/shared-python/shared/services/retrieval/execution/routes.py @@ -344,7 +344,8 @@ async def _run_mapnav_route( completion_detail = f"chunks | evidence={len(evidence_text)} chars | router=mapnav" process_finished = resource.getrusage(resource.RUSAGE_SELF) logger.info( - "retrieval mapnav stage=process_resources cpu_seconds=%.3f max_rss_kb=%d", + "retrieval mapnav stage=process_resources cpu_seconds=%.3f " + "process_max_rss_kb=%d", (process_finished.ru_utime + process_finished.ru_stime) - (process_started.ru_utime + process_started.ru_stime), int(process_finished.ru_maxrss), diff --git a/packages/shared-python/shared/services/retrieval/namespace_map_snapshot_redis.py b/packages/shared-python/shared/services/retrieval/namespace_map_snapshot_redis.py index b977e5f45..162b85c25 100644 --- a/packages/shared-python/shared/services/retrieval/namespace_map_snapshot_redis.py +++ b/packages/shared-python/shared/services/retrieval/namespace_map_snapshot_redis.py @@ -9,30 +9,34 @@ _KEY_PREFIX = "retrieval:snapshot:v2" -def build_snapshot_cache_key(*, user_id: str, namespace: str, generation: int) -> str: - normalized_namespace = normalize_retrieval_namespace(namespace) - return f"{_KEY_PREFIX}:{user_id}:{normalized_namespace}:g{int(generation)}" - - -async def get_snapshot_blob( - *, user_id: str, namespace: str, generation: int -) -> bytes | None: - service = RedisServiceFactory.get_service() - return await service.get_bytes( - build_snapshot_cache_key( - user_id=user_id, namespace=namespace, generation=generation +class NamespaceMapSnapshotRedisCache: + """Access generation-scoped compressed namespace snapshots in Redis.""" + + @staticmethod + def _build_key(*, user_id: str, namespace: str, generation: int) -> str: + normalized_namespace = normalize_retrieval_namespace(namespace) + return f"{_KEY_PREFIX}:{user_id}:{normalized_namespace}:g{int(generation)}" + + @classmethod + async def get( + cls, *, user_id: str, namespace: str, generation: int + ) -> bytes | None: + service = RedisServiceFactory.get_service() + return await service.get_bytes( + cls._build_key( + user_id=user_id, namespace=namespace, generation=generation + ) + ) + + @classmethod + async def set( + cls, *, user_id: str, namespace: str, generation: int, payload_zlib: bytes + ) -> bool: + service = RedisServiceFactory.get_service() + return await service.set_bytes( + cls._build_key( + user_id=user_id, namespace=namespace, generation=generation + ), + payload_zlib, + ex=_CACHE_TTL_SECONDS, ) - ) - - -async def set_snapshot_blob( - *, user_id: str, namespace: str, generation: int, payload_zlib: bytes -) -> bool: - service = RedisServiceFactory.get_service() - return await service.set_bytes( - build_snapshot_cache_key( - user_id=user_id, namespace=namespace, generation=generation - ), - payload_zlib, - ex=_CACHE_TTL_SECONDS, - ) diff --git a/packages/shared-python/shared/services/retrieval/nav_snapshot.py b/packages/shared-python/shared/services/retrieval/nav_snapshot.py index 3210b55f7..b15425e39 100644 --- a/packages/shared-python/shared/services/retrieval/nav_snapshot.py +++ b/packages/shared-python/shared/services/retrieval/nav_snapshot.py @@ -55,8 +55,7 @@ ) from shared.services.retrieval.manifest_cache import get_cached_manifest_payloads from shared.services.retrieval.namespace_map_snapshot_redis import ( - get_snapshot_blob, - set_snapshot_blob, + NamespaceMapSnapshotRedisCache, ) from shared.core.exceptions.redis_exceptions import RedisOperationError @@ -452,8 +451,11 @@ async def _resolve_namespace_snapshot_entries( _current_generation = generation_result.scalar_one_or_none() except SQLAlchemyError as exc: await db.rollback() - _logger.warning("retrieval snapshot generation lookup failed error=%s", exc) - _current_generation = None + _logger.warning( + "retrieval snapshot generation lookup failed; using fallback error=%s", + exc, + ) + return None if _is_snapshot_timing_enabled(): _logger.info( "retrieval snapshot generation generation=%s", @@ -512,7 +514,7 @@ async def _resolve_namespace_snapshot_entries( return None cached_blob: bytes | None = None try: - cached_blob = await get_snapshot_blob( + cached_blob = await NamespaceMapSnapshotRedisCache.get( user_id=user_id, namespace=namespace, generation=row_generation ) except RedisOperationError as exc: @@ -576,7 +578,7 @@ async def _resolve_namespace_snapshot_entries( documents = decoded_documents if cached_blob is None: try: - await set_snapshot_blob( + await NamespaceMapSnapshotRedisCache.set( user_id=user_id, namespace=namespace, generation=row_generation, diff --git a/packages/shared-python/shared/services/retrieval/serving_manifest.py b/packages/shared-python/shared/services/retrieval/serving_manifest.py index 8480b8f49..77a33e1cd 100644 --- a/packages/shared-python/shared/services/retrieval/serving_manifest.py +++ b/packages/shared-python/shared/services/retrieval/serving_manifest.py @@ -249,7 +249,7 @@ def encode_namespace_map_snapshot( documents = payload.get("documents") if not isinstance(documents, dict): raise ValueError("namespace snapshot documents must be an object") - routing_documents: dict[str, Any] = {} + routing_documents: dict[str, dict[str, object]] = {} for document_id, raw_document in documents.items(): if not isinstance(raw_document, dict): raise ValueError(f"namespace snapshot document is not an object: {document_id}") @@ -257,7 +257,7 @@ def encode_namespace_map_snapshot( raw_chunks = raw_document.get("chunks") if not isinstance(raw_sections, list) or not isinstance(raw_chunks, list): raise ValueError(f"namespace snapshot records are invalid: {document_id}") - sections = [] + sections: list[dict[str, object]] = [] for section in raw_sections: if not isinstance(section, dict) or not str(section.get("section_id") or ""): raise ValueError(f"namespace snapshot section is invalid: {document_id}") @@ -271,7 +271,7 @@ def encode_namespace_map_snapshot( if key in section } ) - chunks = [] + chunks: list[dict[str, object]] = [] for chunk in raw_chunks: if not isinstance(chunk, dict) or not str(chunk.get("chunk_id") or ""): raise ValueError(f"namespace snapshot chunk is invalid: {document_id}") From 6ed286d33f4d4e40bfc683ac4ce4214bda4dde3c Mon Sep 17 00:00:00 2001 From: suguanYang Date: Tue, 1 Sep 2026 23:06:12 +0800 Subject: [PATCH 4/5] perf: narrow map index frequency lookup --- .../test_retrieval_map_unit_index_contract.py | 8 +++++-- .../services/retrieval/nav/nav_knowhere.py | 23 +++++++++++++------ 2 files changed, 22 insertions(+), 9 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 b4b047d86..8b0214d9d 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 @@ -105,11 +105,15 @@ def close(self) -> None: ] assert len(frequency_executions) == 1 statement, parameters = frequency_executions[0] - assert "map_unit_id = ANY" in statement + assert "scoped_units AS MATERIALIZED" in statement + assert "JOIN scoped_units" in statement + assert "channel = ANY" in statement assert "token_hash = ANY" in statement + assert "map_unit_id = ANY" not in statement assert isinstance(parameters, list) assert parameters[0] == ["unit-frequency"] - assert parameters[1] == [ + assert parameters[1] == ["path", "content"] + assert parameters[2] == [ "6e51d6a3d90b6a3243d38e6da6b3f31f49867c1360beba83da8ca9630f9672c7" ] 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 bd4e24f76..debfb3c86 100644 --- a/packages/shared-python/shared/services/retrieval/nav/nav_knowhere.py +++ b/packages/shared-python/shared/services/retrieval/nav/nav_knowhere.py @@ -435,18 +435,27 @@ def load_persisted_score_corpus( frequencies: Dict[Tuple[str, str], Dict[str, int]] = {} if unit_rows and query_tokens: stage_started = time.perf_counter() - # Restrict the token scan to this episode's map units instead of - # matching token_hash across the whole table then filtering. + # Keep the episode scope explicit without materializing every + # token row matching a common query term. PostgreSQL can choose + # either the scoped-unit side or the channel/token_hash-leading + # index, while the scope still limits results to the pinned + # revision and allowed sections. allowed_map_unit_ids = [str(row["map_unit_id"]) for row in unit_rows] cur.execute( - "SELECT map_unit_id, channel, token, frequency " - "FROM document_map_unit_tokens " - "WHERE map_unit_id = ANY(%s) " - "AND token_hash = ANY(%s) AND channel = ANY(%s)", + "WITH scoped_units AS MATERIALIZED (" + "SELECT unnest(%s::text[]) AS map_unit_id" + ") " + "SELECT tokens.map_unit_id, tokens.channel, tokens.token, " + "tokens.frequency " + "FROM document_map_unit_tokens AS tokens " + "JOIN scoped_units " + "ON scoped_units.map_unit_id = tokens.map_unit_id " + "WHERE tokens.channel = ANY(%s) " + "AND tokens.token_hash = ANY(%s)", [ allowed_map_unit_ids, - list(query_token_hashes), list(_MAP_SCORE_CHANNELS), + list(query_token_hashes), ], ) for map_unit_id, channel, token, frequency in cur.fetchall(): From a3ad93a5b3a5dd042d774a51639eef8c143e32bd Mon Sep 17 00:00:00 2001 From: suguanYang Date: Wed, 2 Sep 2026 00:11:16 +0800 Subject: [PATCH 5/5] perf: avoid legacy job chunk hydration --- .../test_retrieval_map_unit_index_contract.py | 97 ++++++++++++++++++- .../services/retrieval/hydration/connected.py | 10 +- 2 files changed, 103 insertions(+), 4 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 8b0214d9d..d571ad43b 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 @@ -6,7 +6,7 @@ from uuid import uuid4 from httpx import AsyncClient -from sqlalchemy import delete, select, text +from sqlalchemy import Engine, delete, event, select, text from shared.models.database.document import ( DocumentMapUnit, @@ -21,6 +21,7 @@ compute_corpus_map_and_unit_scores, select_map_highlights, ) +from shared.services.retrieval.hydration.connected import hydrate_connected_target_rows from shared.services.retrieval.nav.nav_knowhere import ( KnowhereProvider, LazyKnowhereProvider, @@ -497,6 +498,100 @@ def record_reference_load( snapshot.close() +async def test_connected_hydration_does_not_load_legacy_job_chunks( + developer_api_client_factory: Callable[ + [], AbstractAsyncContextManager[AsyncClient] + ], +) -> None: + identifier = uuid4().hex[:8] + namespace = f"connected-job-{identifier}" + document_id = f"doc_connected_{identifier}" + job_id = f"job_connected_{identifier}" + job_result_id = f"result_connected_{identifier}" + statements: list[str] = [] + + def capture_job_chunk_query( + _connection: Any, + _cursor: Any, + statement: str, + _parameters: Any, + _context: Any, + _executemany: bool, + ) -> None: + if "job_chunks" in statement.lower(): + statements.append(statement) + + 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="connected.pdf", + ) + chunks = [ + { + "chunk_id": "body-connected", + "type": "text", + "content": "body connected evidence", + "path": "connected.pdf/Root/Section/body", + "order": 1, + "metadata": {"connect_to": [{"target": "asset-connected"}]}, + }, + { + "chunk_id": "asset-connected", + "type": "image", + "content": "asset connected summary", + "path": "images/asset-connected.png", + "order": 2, + "file_path": "images/asset-connected.png", + "metadata": {}, + }, + ] + async with contract_db_session() as db: + await db.run_sync( + lambda sync_db: _publish_revision_with_generation_lock( + sync_db, + scope=scope, + chunks=chunks, + ) + ) + await db.commit() + + event.listen(Engine, "before_cursor_execute", capture_job_chunk_query) + try: + async with contract_db_session() as db: + hydrated = await hydrate_connected_target_rows( + db=db, + rows=[ + { + "document_id": document_id, + "job_result_id": job_result_id, + "chunk_id": "body-connected", + "chunk_type": "text", + "chunk_metadata": { + "connect_to": [{"target": "asset-connected"}] + }, + } + ], + exclude_document_ids=[], + exclude_sections=[], + revision_pins={document_id: job_result_id}, + ) + finally: + event.remove(Engine, "before_cursor_execute", capture_job_chunk_query) + + assert [row["chunk_id"] for row in hydrated] == ["asset-connected"] + assert hydrated[0]["job_id"] == job_id + assert statements == [] + + def test_incomplete_index_returns_empty_scores() -> None: first_sections = [ SectionRow("root-a", None, "Root A", "Root A", 0, "", 0), diff --git a/packages/shared-python/shared/services/retrieval/hydration/connected.py b/packages/shared-python/shared/services/retrieval/hydration/connected.py index 2a1f31276..8e6379573 100644 --- a/packages/shared-python/shared/services/retrieval/hydration/connected.py +++ b/packages/shared-python/shared/services/retrieval/hydration/connected.py @@ -62,7 +62,11 @@ async def hydrate_connected_target_rows( return [] stmt = ( - select(Document, DocumentChunk, DocumentSection, JobResult) + # Select only the job identifier needed for the public projection. + # Selecting the JobResult entity triggers its ``chunks`` selectin + # relationship, loading the entire legacy job-chunk collection for + # every connected revision during final hydration. + select(Document, DocumentChunk, DocumentSection, JobResult.job_id) .join( DocumentChunk, ( @@ -90,7 +94,7 @@ async def hydrate_connected_target_rows( result = await db.execute(stmt) hydrated_rows: list[dict[str, Any]] = [] - for document, chunk, section, job_result in result.all(): + for document, chunk, section, job_id in result.all(): section_path = section.section_path if section else None hydrated_rows.append( { @@ -105,7 +109,7 @@ async def hydrate_connected_target_rows( 'file_path': chunk.file_path, 'chunk_metadata': chunk.chunk_metadata or {}, 'job_result_id': chunk.job_result_id, - 'job_id': job_result.job_id if job_result else None, + 'job_id': job_id, 'sort_order': chunk.sort_order, } )