From 926397b52f6b7f4151d4d25cd017c22089daa0da Mon Sep 17 00:00:00 2001 From: suguanYang Date: Wed, 30 Sep 2026 14:19:10 +0800 Subject: [PATCH 1/2] perf: reuse demo publication artifacts --- apps/api/.env.example | 6 + ...a8b_gate_legacy_token_lookup_retirement.py | 56 ++++ .../api/app/services/demo/canonical_bundle.py | 259 ++++++++++++++++++ .../services/demo/canonical_bundle_result.py | 12 + apps/api/app/services/demo/source_catalog.py | 12 +- .../app/services/demo/source_materializer.py | 98 ++++--- .../services/documents/lifecycle_service.py | 116 +++++++- apps/api/scripts/prepare_demo_bundles.py | 116 ++++++++ .../test_canonical_demo_bundle_contract.py | 234 ++++++++++++++++ .../test_demo_bundle_preparation_contract.py | 48 ++++ .../test_demo_preparation_cache_contract.py | 167 +++++++++++ ...t_demo_publication_trace_owner_contract.py | 25 +- ...legacy_token_lookup_retirement_contract.py | 148 ++++++++++ .../test_page_memory_retrieval_contract.py | 19 ++ docs/design/canonical-demo-bundle-reuse.md | 136 +++++++++ docs/design/canonical-demo-preparation.md | 27 ++ .../demo-publication-preparation-cache.md | 61 +++++ docs/design/legacy-token-lookup-retirement.md | 122 +++++++++ .../publication-scale-cold-cache-findings.md | 173 ++++++++++++ .../shared/core/config/storage.py | 14 + .../retrieval/agent_tools/tools/grep.py | 3 + .../agent_tools/tools/query_table.py | 8 +- .../retrieval/agent_tools/tools/read.py | 14 +- .../services/retrieval/graph/query_service.py | 4 + .../services/retrieval/hydration/assets.py | 27 +- .../services/retrieval/hydration/connected.py | 10 +- .../retrieval/hydration/evidence_compose.py | 6 + .../services/retrieval/hydration/path.py | 4 + .../services/retrieval/hydration/reference.py | 10 +- .../retrieval/hydration/table_grid.py | 6 + .../maintenance/legacy_token_lookup.py | 101 +++++++ .../services/retrieval/map_unit_index.py | 25 +- .../publication_preparation_cache.py | 143 ++++++++++ .../retrieval/scoring/knowhere_hybrid.py | 9 +- .../services/retrieval/search/lexical_text.py | 11 +- .../retrieval/search/map_unit_discovery.py | 3 +- .../retrieval/search/scoped_corpus.py | 4 + .../shared/services/storage/page_pdf_crop.py | 117 +++++++- .../services/storage/raw_prefix_arguments.py | 7 + .../shared/services/storage/result_storage.py | 111 +++++++- ...est_agent_tools_discovery_read_contract.py | 1 + .../test_agent_tools_explore_table_image.py | 1 + .../shared/tests/test_agent_tools_grep.py | 3 + scripts/retire_legacy_token_lookup.py | 74 +++++ 44 files changed, 2440 insertions(+), 111 deletions(-) create mode 100644 apps/api/alembic/versions/3c4d5e6f7a8b_gate_legacy_token_lookup_retirement.py create mode 100644 apps/api/app/services/demo/canonical_bundle.py create mode 100644 apps/api/app/services/demo/canonical_bundle_result.py create mode 100644 apps/api/scripts/prepare_demo_bundles.py create mode 100644 apps/api/tests/contract/test_canonical_demo_bundle_contract.py create mode 100644 apps/api/tests/contract/test_demo_bundle_preparation_contract.py create mode 100644 apps/api/tests/contract/test_demo_preparation_cache_contract.py create mode 100644 apps/api/tests/migrations/test_legacy_token_lookup_retirement_contract.py create mode 100644 docs/design/canonical-demo-bundle-reuse.md create mode 100644 docs/design/canonical-demo-preparation.md create mode 100644 docs/design/demo-publication-preparation-cache.md create mode 100644 docs/design/legacy-token-lookup-retirement.md create mode 100644 docs/design/publication-scale-cold-cache-findings.md create mode 100644 packages/shared-python/shared/services/retrieval/maintenance/legacy_token_lookup.py create mode 100644 packages/shared-python/shared/services/retrieval/publication_preparation_cache.py create mode 100644 packages/shared-python/shared/services/storage/raw_prefix_arguments.py create mode 100644 scripts/retire_legacy_token_lookup.py diff --git a/apps/api/.env.example b/apps/api/.env.example index f5a003ee2..bba577426 100644 --- a/apps/api/.env.example +++ b/apps/api/.env.example @@ -136,6 +136,12 @@ MOESIF_APPLICATION_ID= FREE_PLAN_INITIAL_CREDITS=5 FRONTEND_URL=http://localhost:3000 +# Reuse canonical demo bundles; disable to restore per-job uploads while keeping +# existing canonical artifact URLs readable. Does not change publication strategy. +DEMO_CANONICAL_BUNDLE_ENABLED=false +# Optional per-process tokenization/frequency reuse for versioned demo sources. +DEMO_PUBLICATION_PREPARATION_CACHE_ENABLED=false + # Required for specific features: parsing providers MINERU_API_KEYS= ILOVEAPI_PUBLIC_KEY= diff --git a/apps/api/alembic/versions/3c4d5e6f7a8b_gate_legacy_token_lookup_retirement.py b/apps/api/alembic/versions/3c4d5e6f7a8b_gate_legacy_token_lookup_retirement.py new file mode 100644 index 000000000..5d04d1313 --- /dev/null +++ b/apps/api/alembic/versions/3c4d5e6f7a8b_gate_legacy_token_lookup_retirement.py @@ -0,0 +1,56 @@ +"""Opt in to retiring the legacy token lookup after replacement validation.""" + +from __future__ import annotations + +from collections.abc import Sequence + +from alembic import context, op + +from shared.services.retrieval.maintenance.legacy_token_lookup import ( + LegacyTokenLookupMaintenance, +) + +revision: str = "3c4d5e6f7a8b" +down_revision: str | None = "2b3c4d5e6f7a" +branch_labels: Sequence[str] | None = None +depends_on: Sequence[str] | None = None + + +def _run_maintenance(*, is_restore: bool) -> None: + if context.is_offline_mode(): + raise RuntimeError("Legacy token index maintenance requires live catalog checks") + is_external_transaction: bool = bool( + op.get_context().opts.get("knowhere_external_transaction", False) + ) + + def execute() -> None: + maintenance: LegacyTokenLookupMaintenance = LegacyTokenLookupMaintenance( + op.get_bind(), is_concurrent=not is_external_transaction + ) + if is_restore: + maintenance.restore_index() + else: + maintenance.apply_retirement() + + if is_external_transaction: + execute() + return + with op.get_context().autocommit_block(): + execute() + + +def upgrade() -> None: + """Default deployment records the revision and retains the legacy index.""" + gate: str | None = context.get_x_argument(as_dictionary=True).get( + "retire_legacy_token_lookup" + ) + if gate is None or gate == "false": + return + if gate != "true": + raise ValueError("retire_legacy_token_lookup must be exactly true or false") + _run_maintenance(is_restore=False) + + +def downgrade() -> None: + """Restore the exact legacy definition, including after operator retirement.""" + _run_maintenance(is_restore=True) diff --git a/apps/api/app/services/demo/canonical_bundle.py b/apps/api/app/services/demo/canonical_bundle.py new file mode 100644 index 000000000..432c9c0ea --- /dev/null +++ b/apps/api/app/services/demo/canonical_bundle.py @@ -0,0 +1,259 @@ +"""Publish immutable demo artifacts once per source content version.""" + +from __future__ import annotations + +import asyncio +import json +import re +import shutil +import tempfile +import time +from functools import lru_cache +from hashlib import sha256 +from io import BytesIO +from pathlib import Path +from uuid import uuid4 + +from loguru import logger + +from app.services.demo.canonical_bundle_result import CanonicalDemoBundle +from shared.services.redis.redis_service import RedisService +from shared.services.storage.job_file_storage import JobFileStorage +from shared.services.storage.result_storage import ( + ResultStorage, + UploadedResultBundle, + get_result_storage, +) +from shared.services.storage.storage_adapter import StorageAdapter + +_READY_FILE = ".bundle-ready.json" +_LEASE_SECONDS = 180 +_WAIT_SECONDS = 600 +_RENEW_SCRIPT = """ +if redis.call('get', KEYS[1]) == ARGV[1] then + return redis.call('expire', KEYS[1], ARGV[2]) +end +return 0 +""" +_RELEASE_SCRIPT = """ +if redis.call('get', KEYS[1]) == ARGV[1] then + return redis.call('del', KEYS[1]) +end +return 0 +""" +_SourceSignature = tuple[tuple[str, int, int, int], ...] + + +class CanonicalDemoBundleStore: + """Own versioning, upload coordination and the durable completion marker.""" + + def __init__(self, *, redis_service: RedisService) -> None: + self._redis: RedisService = redis_service + self._storage: ResultStorage = get_result_storage() + self._files: JobFileStorage = JobFileStorage() + + async def ensure_bundle( + self, *, source_id: str, source_directory: Path + ) -> CanonicalDemoBundle: + if re.fullmatch(r"[a-zA-Z0-9][a-zA-Z0-9_-]*", source_id) is None: + raise ValueError("Invalid canonical demo source identifier") + directory: Path = source_directory.resolve(strict=True) + signature: _SourceSignature = await asyncio.to_thread( + _read_source_signature, directory + ) + version: str = await asyncio.to_thread( + _calculate_content_version, directory, signature + ) + storage_id: str = f"demo-canonical/{source_id}/{version}" + bundle: CanonicalDemoBundle | None = await asyncio.to_thread( + self._read_ready_bundle, storage_id, version + ) + if bundle is not None: + return bundle + + lock_key: str = f"lock:demo_bundle:{source_id}:{version}" + owner: str = uuid4().hex + deadline: float = time.monotonic() + _WAIT_SECONDS + while not await self._redis.set_nx(lock_key, owner, ex=_LEASE_SECONDS): + bundle = await asyncio.to_thread( + self._read_ready_bundle, storage_id, version + ) + if bundle is not None: + return bundle + if time.monotonic() >= deadline: + raise TimeoutError("Timed out waiting for canonical demo artifacts") + await asyncio.sleep(0.5) + + stopped: asyncio.Event = asyncio.Event() + renewal: asyncio.Task[None] = asyncio.create_task( + self._renew_lease(lock_key, owner, stopped) + ) + try: + # A previous owner can finish between the first read and SET NX. + bundle = await asyncio.to_thread( + self._read_ready_bundle, storage_id, version + ) + if bundle is not None: + return bundle + upload: asyncio.Task[CanonicalDemoBundle] = asyncio.create_task( + asyncio.to_thread( + self._upload_bundle, storage_id, version, directory, signature + ) + ) + try: + uploaded_bundle: CanonicalDemoBundle = await asyncio.shield(upload) + except asyncio.CancelledError: + # A thread cannot be cancelled; retain the lease until it stops. + await upload + raise + if renewal.done(): + renewal.result() + if not await self._redis.eval( + _RENEW_SCRIPT, keys=[lock_key], args=[owner, _LEASE_SECONDS] + ): + raise RuntimeError("Canonical demo upload lease was lost") + marker_upload: asyncio.Task[None] = asyncio.create_task( + asyncio.to_thread(self._write_ready_marker, uploaded_bundle) + ) + try: + await asyncio.shield(marker_upload) + except asyncio.CancelledError: + await marker_upload + raise + return uploaded_bundle + finally: + stopped.set() + try: + await renewal + except Exception as error: + logger.warning("Canonical demo lease renewal failed: {}", error) + try: + await self._redis.eval(_RELEASE_SCRIPT, keys=[lock_key], args=[owner]) + except Exception as error: + # Ownership expires even if Redis is unavailable during cleanup. + logger.warning("Canonical demo lease release failed: {}", error) + + async def _renew_lease( + self, lock_key: str, owner: str, stopped: asyncio.Event + ) -> None: + while not stopped.is_set(): + try: + await asyncio.wait_for(stopped.wait(), timeout=_LEASE_SECONDS / 3) + except TimeoutError: + if not await self._redis.eval( + _RENEW_SCRIPT, keys=[lock_key], args=[owner, _LEASE_SECONDS] + ): + raise RuntimeError("Canonical demo upload lease was lost") + + def _read_ready_bundle( + self, storage_id: str, version: str + ) -> CanonicalDemoBundle | None: + raw_prefix: str = self._files.build_result_raw_prefix(job_id=storage_id) + marker_key: str = f"{raw_prefix}{_READY_FILE}" + adapter: StorageAdapter = self._files.storage_adapter + bucket: str = self._files.results_bucket + if not adapter.exists(marker_key, bucket): + return None + try: + marker: object = json.loads(adapter.download_fileobj(marker_key, bucket)) + except (ValueError, UnicodeDecodeError): + return None + if not isinstance(marker, dict) or marker.get("content_version") != version: + return None + zip_size: object = marker.get("zip_size") + if type(zip_size) is not int or zip_size <= 0: + return None + zip_key: str = self._files.build_result_zip_key(job_id=storage_id) + if adapter.get_object_size(zip_key, bucket) != zip_size: + return None + return CanonicalDemoBundle( + zip_key=zip_key, + zip_size=zip_size, + raw_prefix=raw_prefix, + content_version=version, + reused=True, + ) + + def _upload_bundle( + self, + storage_id: str, + version: str, + directory: Path, + signature: _SourceSignature, + ) -> CanonicalDemoBundle: + marker_key: str = ( + f"{self._files.build_result_raw_prefix(job_id=storage_id)}{_READY_FILE}" + ) + if self._files.storage_adapter.exists(marker_key, self._files.results_bucket): + # A malformed marker or missing ZIP must stay incomplete during repair. + self._files.delete_object(marker_key, bucket=self._files.results_bucket) + with tempfile.TemporaryDirectory(prefix="knowhere-demo-result-") as temporary: + archive_path: Path = Path( + shutil.make_archive( + str(Path(temporary) / "bundle"), "zip", root_dir=directory + ) + ) + zip_size: int = archive_path.stat().st_size + uploaded: UploadedResultBundle = self._storage.upload( + job_id=storage_id, + result_dir=str(directory), + zip_file_path=str(archive_path), + ) + if _read_source_signature(directory) != signature: + raise RuntimeError("Canonical demo source changed during artifact upload") + return CanonicalDemoBundle( + zip_key=uploaded.zip_key, + zip_size=zip_size, + raw_prefix=uploaded.raw_prefix, + content_version=version, + reused=False, + ) + + def _write_ready_marker(self, bundle: CanonicalDemoBundle) -> None: + # The marker is the commit point; a ZIP alone is never a complete bundle. + payload: bytes = json.dumps( + {"content_version": bundle.content_version, "zip_size": bundle.zip_size} + ).encode("utf-8") + self._files.upload_fileobj( + BytesIO(payload), + f"{bundle.raw_prefix}{_READY_FILE}", + bucket=self._files.results_bucket, + content_type="application/json", + ) + + +def _read_source_signature(directory: Path) -> _SourceSignature: + files: list[tuple[str, int, int, int]] = [] + for path in sorted(directory.rglob("*")): + if path.is_symlink(): + raise ValueError("Canonical demo sources cannot contain symbolic links") + if path.is_file(): + status = path.stat() + files.append( + ( + path.relative_to(directory).as_posix(), + status.st_size, + status.st_mtime_ns, + status.st_ctime_ns, + ) + ) + if not files: + raise ValueError("Canonical demo source directory is empty") + return tuple(files) + + +@lru_cache(maxsize=128) +def _calculate_content_version(directory: Path, signature: _SourceSignature) -> str: + # Stat changes invalidate the local memo; the durable key uses actual bytes. + digest = sha256(b"knowhere-demo-bundle-v1\0") + for relative_path, size, _modified, _changed in signature: + path_bytes: bytes = relative_path.encode("utf-8") + digest.update(len(path_bytes).to_bytes(8, "big")) + digest.update(path_bytes) + digest.update(size.to_bytes(8, "big")) + with (directory / relative_path).open("rb") as source_file: + while block := source_file.read(1024 * 1024): + digest.update(block) + if _read_source_signature(directory) != signature: + raise RuntimeError("Canonical demo source changed while computing its version") + return digest.hexdigest() diff --git a/apps/api/app/services/demo/canonical_bundle_result.py b/apps/api/app/services/demo/canonical_bundle_result.py new file mode 100644 index 000000000..a67fdf142 --- /dev/null +++ b/apps/api/app/services/demo/canonical_bundle_result.py @@ -0,0 +1,12 @@ +"""Immutable object-storage references for one canonical demo version.""" + +from dataclasses import dataclass + + +@dataclass(frozen=True) +class CanonicalDemoBundle: + zip_key: str + zip_size: int + raw_prefix: str + content_version: str + reused: bool diff --git a/apps/api/app/services/demo/source_catalog.py b/apps/api/app/services/demo/source_catalog.py index 59c59abea..774234d07 100644 --- a/apps/api/app/services/demo/source_catalog.py +++ b/apps/api/app/services/demo/source_catalog.py @@ -651,9 +651,19 @@ def _normalize_asset_path(asset_path: str) -> Path | None: return Path(*parts) -@lru_cache(maxsize=8) def _load_source_chunks(source: DemoSourceDefinition) -> tuple[dict[str, Any], ...]: chunks_path = (_DATA_ROOT / source.asset_directory) / "chunks.json" + status = chunks_path.stat() + return _load_source_chunks_version( + chunks_path, status.st_size, status.st_mtime_ns, status.st_ctime_ns + ) + + +@lru_cache(maxsize=8) +def _load_source_chunks_version( + chunks_path: Path, size: int, modified_at: int, changed_at: int +) -> tuple[dict[str, Any], ...]: + """Invalidate parsed source content when the canonical input file changes.""" with chunks_path.open("r", encoding="utf-8") as file: payload = json.load(file) diff --git a/apps/api/app/services/demo/source_materializer.py b/apps/api/app/services/demo/source_materializer.py index d27f86054..022746046 100644 --- a/apps/api/app/services/demo/source_materializer.py +++ b/apps/api/app/services/demo/source_materializer.py @@ -2,6 +2,7 @@ from __future__ import annotations +import asyncio import shutil import tempfile import time @@ -16,6 +17,8 @@ import logfire from app.services.demo.source_catalog import DemoSourceCatalog, DemoSourceDefinition +from app.services.demo.canonical_bundle import CanonicalDemoBundleStore +from app.services.demo.canonical_bundle_result import CanonicalDemoBundle from sqlalchemy.exc import IntegrityError from sqlalchemy import delete, func, select from sqlalchemy.ext.asyncio import AsyncConnection, AsyncSession @@ -33,7 +36,9 @@ from shared.services.retrieval.cache_service import invalidate_retrieval_cache_namespaces from shared.services.retrieval.publication_service import RetrievalPublicationService from shared.services.retrieval.publication_models import DocumentPublicationScope -from shared.services.redis import RedisPublicationSemaphore, RedisServiceFactory +from shared.services.retrieval.publication_preparation_cache import PublicationPreparationCache +from shared.services.retrieval.scoring.knowhere_hybrid import MAP_UNIT_INDEX_FORMAT_VERSION +from shared.services.redis import RedisPublicationSemaphore, RedisService, RedisServiceFactory from shared.core.config import settings from shared.services.storage.result_storage import get_result_storage @@ -157,14 +162,18 @@ async def _materialize_source( bind.engine if isinstance(bind, AsyncConnection) else bind ) stage_started_at = time.perf_counter() - result_bundle = _upload_demo_result_bundle( + result_bundle = await _upload_demo_result_bundle( job_id=job_id, + demo_source_id=source.demo_source_id, source_directory=self._catalog.source_directory(source), + redis_service=self._redis_service, ) logfire.info( "Demo materialization source bundle upload completed", demo_source_id=source.demo_source_id, duration_seconds=time.perf_counter() - stage_started_at, + reused=result_bundle.reused, + content_version=result_bundle.content_version, ) db.add( @@ -197,10 +206,15 @@ async def _materialize_source( document_metadata={ "source_file_name": source.title, "demo_source_id": source.demo_source_id, + **( + {"result_raw_prefix": result_bundle.raw_prefix} + if result_bundle.content_version + else {} + ), }, inline_payload={"source": "canonical_demo"}, - result_s3_key=result_bundle["zip_key"], - result_size=result_bundle["zip_size"], + result_s3_key=result_bundle.zip_key, + result_size=result_bundle.zip_size, created_at=timestamp, updated_at=timestamp, ) @@ -259,16 +273,26 @@ async def _materialize_source( demo_source_id=source.demo_source_id, duration_seconds=time.perf_counter() - base_rows_started_at, ) - published_state = await db.run_sync( - lambda sync_db: self._publication_service.publish_document_state( - sync_db, - job_id=job_id, - job_result_id=job_result_id, - chunks=[dict(chunk) for chunk in chunks], - update_namespace_snapshot=False, - trace=trace, + with ( + PublicationPreparationCache.scope( + source_id=source.demo_source_id, + content_version=result_bundle.content_version, + index_format_version=MAP_UNIT_INDEX_FORMAT_VERSION, + ) + if settings.DEMO_PUBLICATION_PREPARATION_CACHE_ENABLED + and result_bundle.content_version + else nullcontext() + ): + published_state = await db.run_sync( + lambda sync_db: self._publication_service.publish_document_state( + sync_db, + job_id=job_id, + job_result_id=job_result_id, + chunks=[dict(chunk) for chunk in chunks], + update_namespace_snapshot=False, + trace=trace, + ) ) - ) await db.run_sync( lambda sync_db: self._publication_service.publish_document_graph( sync_db, @@ -493,32 +517,38 @@ def _materialized_source_payload( ) -def _upload_demo_result_bundle( +async def _upload_demo_result_bundle( *, job_id: str, + demo_source_id: str, source_directory: Path, -) -> dict[str, int | str]: - storage = get_result_storage() - with tempfile.TemporaryDirectory(prefix="knowhere-demo-result-") as temp_directory: - zip_base_path = Path(temp_directory) / job_id - zip_file_path = Path( - shutil.make_archive( - str(zip_base_path), - "zip", - root_dir=source_directory, - ) - ) - zip_size = zip_file_path.stat().st_size - bundle = storage.upload( - job_id=job_id, - result_dir=str(source_directory), - zip_file_path=str(zip_file_path), + redis_service: RedisService, +) -> CanonicalDemoBundle: + if not settings.DEMO_CANONICAL_BUNDLE_ENABLED: + return await asyncio.to_thread( + _upload_job_result_bundle, job_id=job_id, source_directory=source_directory ) + return await CanonicalDemoBundleStore(redis_service=redis_service).ensure_bundle( + source_id=demo_source_id, source_directory=source_directory + ) + - return { - "zip_key": bundle.zip_key, - "zip_size": zip_size, - } +def _upload_job_result_bundle( + *, job_id: str, source_directory: Path +) -> CanonicalDemoBundle: + with tempfile.TemporaryDirectory(prefix="knowhere-demo-result-") as temporary: + archive_path: Path = Path(shutil.make_archive( + str(Path(temporary) / job_id), "zip", root_dir=source_directory + )) + zip_size: int = archive_path.stat().st_size + bundle = get_result_storage().upload( + job_id=job_id, result_dir=str(source_directory), + zip_file_path=str(archive_path), + ) + return CanonicalDemoBundle( + zip_key=bundle.zip_key, zip_size=zip_size, raw_prefix=bundle.raw_prefix, + content_version="", reused=False, + ) def _utc_now() -> datetime: diff --git a/apps/api/app/services/documents/lifecycle_service.py b/apps/api/app/services/documents/lifecycle_service.py index fd97a9ba2..9e1bd839e 100644 --- a/apps/api/app/services/documents/lifecycle_service.py +++ b/apps/api/app/services/documents/lifecycle_service.py @@ -43,6 +43,7 @@ def _document_chunk_asset_url( *, chunk_type: str, job_id: str | None, + raw_prefix: str | None, file_path: str | None, include_asset_urls: bool, result_storage: ResultStorage | None, @@ -57,8 +58,10 @@ def _document_chunk_asset_url( return None try: - return result_storage.generate_artifact_url( + return _generate_artifact_url( + result_storage, job_id=job_id, + raw_prefix=raw_prefix, artifact_ref=file_path, expires_in=_DOCUMENT_CHUNK_ASSET_URL_EXPIRES_SECONDS, ) @@ -71,6 +74,7 @@ def _document_page_assets( *, metadata: dict[str, Any] | None, job_id: str | None, + raw_prefix: str | None, include_asset_urls: bool, result_storage: ResultStorage | None, ) -> list[dict[str, Any]]: @@ -90,6 +94,7 @@ def _document_page_assets( if include_asset_urls and job_id and result_storage is not None: asset_url = _page_asset_url( job_id=job_id, + raw_prefix=raw_prefix, artifact_ref=asset["artifact_ref"], result_storage=result_storage, ) @@ -125,6 +130,7 @@ def _normalize_page_asset(raw_asset: dict[str, Any]) -> dict[str, Any] | None: def _page_asset_url( *, job_id: str, + raw_prefix: str | None, artifact_ref: str, result_storage: ResultStorage, ) -> str | None: @@ -132,8 +138,10 @@ def _page_asset_url( if not normalized_ref or not normalized_ref.startswith("page_citation_assets/"): return None try: - return result_storage.generate_artifact_url( + return _generate_artifact_url( + result_storage, job_id=job_id, + raw_prefix=raw_prefix, artifact_ref=normalized_ref, expires_in=_DOCUMENT_CHUNK_ASSET_URL_EXPIRES_SECONDS, ) @@ -150,6 +158,76 @@ def _positive_int(value: Any) -> int | None: return number if number > 0 else None +def _result_raw_prefix(metadata: object) -> str | None: + if not isinstance(metadata, dict): + return None + value = metadata.get("result_raw_prefix") + if not isinstance(value, str): + return None + normalized = value.strip() + return normalized or None + + +def _generate_artifact_url( + storage: ResultStorage, + *, + job_id: str, + raw_prefix: str | None, + artifact_ref: str, + expires_in: int, +) -> str | None: + if raw_prefix is None: + return storage.generate_artifact_url( + job_id=job_id, + artifact_ref=artifact_ref, + expires_in=expires_in, + ) + return storage.generate_artifact_url( + job_id=job_id, + raw_prefix=raw_prefix, + artifact_ref=artifact_ref, + expires_in=expires_in, + ) + + +def _verify_raw_exists( + storage: ResultStorage, + *, + job_id: str, + raw_prefix: str | None, + relative_path: str, +) -> bool: + if raw_prefix is None: + return storage.verify_raw_exists(job_id=job_id, relative_path=relative_path) + return storage.verify_raw_exists( + job_id=job_id, + raw_prefix=raw_prefix, + relative_path=relative_path, + ) + + +def _generate_raw_file_url( + storage: ResultStorage, + *, + job_id: str, + raw_prefix: str | None, + relative_path: str, + expires_in: int, +) -> str | None: + if raw_prefix is None: + return storage.generate_raw_file_url( + job_id=job_id, + relative_path=relative_path, + expires_in=expires_in, + ) + return storage.generate_raw_file_url( + job_id=job_id, + raw_prefix=raw_prefix, + relative_path=relative_path, + expires_in=expires_in, + ) + + def document_payload(document) -> dict[str, Any]: return { "document_id": document.document_id, @@ -266,6 +344,7 @@ async def list_document_chunks( chunk=chunk, section=section, job_id=job_result.job_id, + raw_prefix=_result_raw_prefix(job_result.document_metadata), include_asset_urls=include_asset_urls, result_storage=result_storage, ) @@ -324,6 +403,7 @@ async def get_document_chunk( chunk=chunk, section=section, job_id=job_result.job_id, + raw_prefix=_result_raw_prefix(job_result.document_metadata), include_asset_urls=include_asset_urls, result_storage=result_storage, ), @@ -365,17 +445,26 @@ async def get_document_page_citation_source( return None result_storage = self._result_storage or get_result_storage() - if not result_storage.verify_raw_exists( - job_id=job_result.job_id, - relative_path=_PAGE_CITATION_SOURCE_FILE_NAME, - ): - return None + raw_prefix = _result_raw_prefix(job_result.document_metadata) + try: + if not _verify_raw_exists( + result_storage, + job_id=job_result.job_id, + raw_prefix=raw_prefix, + relative_path=_PAGE_CITATION_SOURCE_FILE_NAME, + ): + return None - source_url = result_storage.generate_raw_file_url( - job_id=job_result.job_id, - relative_path=_PAGE_CITATION_SOURCE_FILE_NAME, - expires_in=_PAGE_CITATION_SOURCE_EXPIRES_SECONDS, - ) + source_url = _generate_raw_file_url( + result_storage, + job_id=job_result.job_id, + raw_prefix=raw_prefix, + relative_path=_PAGE_CITATION_SOURCE_FILE_NAME, + expires_in=_PAGE_CITATION_SOURCE_EXPIRES_SECONDS, + ) + except Exception as exc: + logger.warning(f"Failed to generate page citation source URL (ignored): {exc}") + return None if not source_url: return None @@ -400,6 +489,7 @@ def _chunk_payload( chunk: DocumentChunk, section: DocumentSection | None, job_id: str | None, + raw_prefix: str | None, include_asset_urls: bool, result_storage: ResultStorage | None, ) -> dict[str, Any]: @@ -409,6 +499,7 @@ def _chunk_payload( page_assets = _document_page_assets( metadata=raw_metadata, job_id=job_id, + raw_prefix=raw_prefix, include_asset_urls=include_asset_urls, result_storage=result_storage, ) @@ -429,6 +520,7 @@ def _chunk_payload( "asset_url": _document_chunk_asset_url( chunk_type=chunk_type, job_id=job_id, + raw_prefix=raw_prefix, file_path=file_path, include_asset_urls=include_asset_urls, result_storage=result_storage, diff --git a/apps/api/scripts/prepare_demo_bundles.py b/apps/api/scripts/prepare_demo_bundles.py new file mode 100644 index 000000000..ef914caaf --- /dev/null +++ b/apps/api/scripts/prepare_demo_bundles.py @@ -0,0 +1,116 @@ +"""Prepare canonical demo result artifacts without materializing documents.""" + +from __future__ import annotations + +import argparse +import asyncio +import os +import sys +from pathlib import Path + + +def _bootstrap_python_path() -> None: + api_root: Path = Path(__file__).resolve().parents[1] + shared_root: Path = api_root.parents[1] / "packages" / "shared-python" + for root in (api_root, shared_root): + value: str = os.fspath(root) + if value not in sys.path: + sys.path.insert(0, value) + + +_bootstrap_python_path() + +from app.services.demo.canonical_bundle import CanonicalDemoBundleStore # noqa: E402 +from app.services.demo.source_catalog import ( # noqa: E402 + DemoSourceCatalog, + DemoSourceDefinition, +) +from shared.services.redis.redis_service import RedisService # noqa: E402 + + +def _build_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("source_ids", metavar="SOURCE_ID", nargs="*") + parser.add_argument( + "--all", action="store_true", help="Select every catalog demo source." + ) + parser.add_argument( + "--upload", + action="store_true", + help="Upload missing canonical bundles to the configured result bucket.", + ) + return parser + + +def _select_sources( + parser: argparse.ArgumentParser, + catalog: DemoSourceCatalog, + *, + source_ids: list[str], + all_sources: bool, +) -> list[DemoSourceDefinition]: + if all_sources and source_ids: + parser.error("--all cannot be combined with source IDs") + if not all_sources and not source_ids: + parser.error("provide source IDs or --all") + if all_sources: + return list(catalog.list_sources()) + + sources: list[DemoSourceDefinition] = [] + for source_id in dict.fromkeys(source_ids): + source: DemoSourceDefinition | None = catalog.get_source(source_id) + if source is None: + parser.error(f"unknown demo source ID: {source_id}") + sources.append(source) + return sources + + +async def _upload_sources( + *, catalog: DemoSourceCatalog, sources: list[DemoSourceDefinition] +) -> None: + redis_service: RedisService = RedisService() + try: + store: CanonicalDemoBundleStore = CanonicalDemoBundleStore( + redis_service=redis_service + ) + for source in sources: + bundle = await store.ensure_bundle( + source_id=source.demo_source_id, + source_directory=catalog.source_directory(source), + ) + state: str = "reused" if bundle.reused else "created" + print( + f"{state} {source.demo_source_id} " + f"version={bundle.content_version} " + f"zip={bundle.zip_key} raw_prefix={bundle.raw_prefix}" + ) + finally: + await redis_service.close() + + +def main(argv: list[str] | None = None) -> int: + parser: argparse.ArgumentParser = _build_parser() + arguments: argparse.Namespace = parser.parse_args(argv) + catalog: DemoSourceCatalog = DemoSourceCatalog() + sources: list[DemoSourceDefinition] = _select_sources( + parser, + catalog, + source_ids=arguments.source_ids, + all_sources=arguments.all, + ) + if not arguments.upload: + for source in sources: + print(f"dry-run {source.demo_source_id} {source.title}") + print(f"Selected {len(sources)} source(s); pass --upload to prepare bundles.") + return 0 + + try: + asyncio.run(_upload_sources(catalog=catalog, sources=sources)) + except Exception as error: + print(f"Canonical demo preparation failed: {error}", file=sys.stderr) + return 1 + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/apps/api/tests/contract/test_canonical_demo_bundle_contract.py b/apps/api/tests/contract/test_canonical_demo_bundle_contract.py new file mode 100644 index 000000000..9a075609d --- /dev/null +++ b/apps/api/tests/contract/test_canonical_demo_bundle_contract.py @@ -0,0 +1,234 @@ +"""Contract coverage for immutable canonical demo artifact reuse.""" + +from __future__ import annotations + +import asyncio +from collections.abc import Callable +from contextlib import AbstractAsyncContextManager +from pathlib import Path + +import pytest +from httpx import AsyncClient +from pytest import MonkeyPatch + +from tests.support.contract_database import ContractDatabase + + +def _write_source(directory: Path, content: str) -> None: + directory.mkdir(parents=True, exist_ok=True) + (directory / "source.md").write_text(content, encoding="utf-8") + + +@pytest.mark.asyncio +async def test_canonical_bundle_reuses_and_versions_real_filesystem_objects( + api_client_factory: Callable[[], AbstractAsyncContextManager[AsyncClient]], + tmp_path: Path, +) -> None: + async with api_client_factory(): + from app.services.demo.canonical_bundle import CanonicalDemoBundleStore + from shared.services.redis import RedisServiceFactory + + source_directory = tmp_path / "source" + _write_source(source_directory, "version one") + store = CanonicalDemoBundleStore( + redis_service=RedisServiceFactory.get_service() + ) + + first = await store.ensure_bundle( + source_id="contract-source", + source_directory=source_directory, + ) + reused = await store.ensure_bundle( + source_id="contract-source", + source_directory=source_directory, + ) + + assert first.reused is False + assert reused.reused is True + assert reused.content_version == first.content_version + assert reused.zip_key == first.zip_key + assert reused.raw_prefix == first.raw_prefix + + _write_source(source_directory, "version two") + changed = await store.ensure_bundle( + source_id="contract-source", + source_directory=source_directory, + ) + + assert changed.reused is False + assert changed.content_version != first.content_version + assert changed.zip_key != first.zip_key + assert changed.raw_prefix != first.raw_prefix + + +@pytest.mark.asyncio +async def test_canonical_bundle_retries_when_completion_marker_is_missing( + api_client_factory: Callable[[], AbstractAsyncContextManager[AsyncClient]], + tmp_path: Path, +) -> None: + async with api_client_factory(): + from app.services.demo.canonical_bundle import CanonicalDemoBundleStore + from shared.services.redis import RedisServiceFactory + + source_directory = tmp_path / "source" + _write_source(source_directory, "marker recovery") + store = CanonicalDemoBundleStore( + redis_service=RedisServiceFactory.get_service() + ) + initial = await store.ensure_bundle( + source_id="contract-marker-recovery", + source_directory=source_directory, + ) + marker_key = f"{initial.raw_prefix}.bundle-ready.json" + assert store._files.storage_adapter.delete_object( + marker_key, store._files.results_bucket + ) + + original_writer = store._write_ready_marker + + def fail_marker(bundle: object) -> None: + raise RuntimeError("marker write failed") + + store._write_ready_marker = fail_marker # type: ignore[method-assign] + with pytest.raises(RuntimeError, match="marker write failed"): + await store.ensure_bundle( + source_id="contract-marker-recovery", + source_directory=source_directory, + ) + + store._write_ready_marker = original_writer # type: ignore[method-assign] + recovered = await store.ensure_bundle( + source_id="contract-marker-recovery", + source_directory=source_directory, + ) + + assert recovered.reused is False + assert recovered.zip_key == initial.zip_key + assert recovered.content_version == initial.content_version + + +@pytest.mark.asyncio +async def test_canonical_bundle_concurrent_callers_publish_once( + api_client_factory: Callable[[], AbstractAsyncContextManager[AsyncClient]], + tmp_path: Path, +) -> None: + async with api_client_factory(): + from app.services.demo.canonical_bundle import CanonicalDemoBundleStore + from shared.services.redis import RedisServiceFactory + + source_directory = tmp_path / "source" + _write_source(source_directory, "concurrent publication") + redis_service = RedisServiceFactory.get_service() + first_store = CanonicalDemoBundleStore(redis_service=redis_service) + second_store = CanonicalDemoBundleStore(redis_service=redis_service) + + results = await asyncio.gather( + first_store.ensure_bundle( + source_id="contract-concurrent", + source_directory=source_directory, + ), + second_store.ensure_bundle( + source_id="contract-concurrent", + source_directory=source_directory, + ), + ) + + assert len({result.zip_key for result in results}) == 1 + assert len({result.raw_prefix for result in results}) == 1 + assert sum(not result.reused for result in results) == 1 + + +@pytest.mark.asyncio +async def test_http_materialization_reuses_canonical_bundle_across_namespaces( + developer_api_client_factory: Callable[ + [], AbstractAsyncContextManager[AsyncClient] + ], + monkeypatch: MonkeyPatch, +) -> None: + monkeypatch.setenv("DEMO_CANONICAL_BUNDLE_ENABLED", "true") + async with developer_api_client_factory() as api_client: + import app.services.demo.source_materializer as source_materializer_module + + first = await api_client.post( + "/api/v1/demo/materializations", + json={ + "namespace": "canonical-contract-one", + "demo_source_ids": ["demo-tsla-q4-2025"], + }, + ) + second = await api_client.post( + "/api/v1/demo/materializations", + json={ + "namespace": "canonical-contract-two", + "demo_source_ids": ["demo-tsla-q4-2025"], + }, + ) + assert first.status_code == 200 + assert second.status_code == 200 + first_document_id = str(first.json()["sources"][0]["document_id"]) + second_document_id = str(second.json()["sources"][0]["document_id"]) + chunks_response = await api_client.get( + f"/api/v1/documents/{first_document_id}/chunks?page_size=200" + "&include_asset_urls=true" + ) + retrieval_response = await api_client.post( + "/api/v1/retrieval/query", + json={ + "namespace": "canonical-contract-one", + "query": "Tesla investment", + "top_k": 3, + "use_agentic": False, + }, + ) + assert retrieval_response.status_code == 200 + retrieval_asset_urls: list[str] = [ + str(result["asset_url"]) + for result in retrieval_response.json()["results"] + if result.get("asset_url") + ] + assert retrieval_asset_urls + assert all("demo-canonical" in url for url in retrieval_asset_urls) + rows = await ContractDatabase.fetch_all( + """ + SELECT d.document_id, jr.result_s3_key, jr.result_size, + jr.document_metadata ->> 'result_raw_prefix' AS result_raw_prefix + FROM job_results AS jr + JOIN documents AS d ON d.current_job_result_id = jr.id + WHERE d.document_id IN (:first_document_id, :second_document_id) + ORDER BY d.document_id + """, + { + "first_document_id": first_document_id, + "second_document_id": second_document_id, + }, + ) + + monkeypatch.setattr( + source_materializer_module.settings, + "DEMO_CANONICAL_BUNDLE_ENABLED", + False, + ) + existing_document_response = await api_client.get( + f"/api/v1/documents/{first_document_id}/chunks?page_size=200" + "&include_asset_urls=true" + ) + + chunks_body = chunks_response.json() + asset_urls = [ + str(chunk["asset_url"]) + for chunk in chunks_body["chunks"] + if chunk.get("asset_url") + ] + assert asset_urls + assert any("demo-canonical" in asset_url for asset_url in asset_urls) + assert existing_document_response.status_code == 200 + existing_asset_urls = [ + str(chunk["asset_url"]) + for chunk in existing_document_response.json()["chunks"] + if chunk.get("asset_url") + ] + assert existing_asset_urls == asset_urls + assert len(rows) == 2 + assert len({str(row["result_s3_key"]) for row in rows}) == 1 + assert len({str(row["result_raw_prefix"]) for row in rows}) == 1 + assert all(int(row["result_size"]) > 0 for row in rows) diff --git a/apps/api/tests/contract/test_demo_bundle_preparation_contract.py b/apps/api/tests/contract/test_demo_bundle_preparation_contract.py new file mode 100644 index 000000000..d2e8dbcb7 --- /dev/null +++ b/apps/api/tests/contract/test_demo_bundle_preparation_contract.py @@ -0,0 +1,48 @@ +from __future__ import annotations + +import importlib.util +from pathlib import Path + + +def _load_script(): + script_path = ( + Path(__file__).resolve().parents[2] / "scripts" / "prepare_demo_bundles.py" + ) + spec = importlib.util.spec_from_file_location("prepare_demo_bundles", script_path) + if spec is None or spec.loader is None: + raise AssertionError("could not load preparation script") + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +def test_preparation_defaults_to_explicit_dry_run(capsys) -> None: + module = _load_script() + + assert module.main(["demo-spacex-s1"]) == 0 + output = capsys.readouterr().out + assert "dry-run demo-spacex-s1" in output + assert "--upload" in output + + +def test_preparation_requires_source_selection() -> None: + module = _load_script() + + try: + module.main([]) + except SystemExit as error: + assert error.code == 2 + else: + raise AssertionError("missing source selection should fail") + + +def test_preparation_rejects_all_with_explicit_source(capsys) -> None: + module = _load_script() + + try: + module.main(["--all", "demo-spacex-s1"]) + except SystemExit as error: + assert error.code == 2 + else: + raise AssertionError("mixed source selection should fail") + assert "cannot be combined" in capsys.readouterr().err diff --git a/apps/api/tests/contract/test_demo_preparation_cache_contract.py b/apps/api/tests/contract/test_demo_preparation_cache_contract.py new file mode 100644 index 000000000..20005952a --- /dev/null +++ b/apps/api/tests/contract/test_demo_preparation_cache_contract.py @@ -0,0 +1,167 @@ +"""HTTP and database contract for demo preparation-cache parity.""" + +from __future__ import annotations + +from collections.abc import Callable +from contextlib import AbstractAsyncContextManager +from typing import Any, cast + +import pytest +from httpx import AsyncClient +from pytest import MonkeyPatch + +from tests.support.contract_database import ContractDatabase + +_SOURCE_ID = "demo-tsla-q4-2025" + + +async def _publish( + client: AsyncClient, + *, + namespace: str, +) -> str: + response = await client.post( + "/api/v1/demo/materializations", + json={"namespace": namespace, "demo_source_ids": [_SOURCE_ID]}, + ) + assert response.status_code == 200 + return str(response.json()["sources"][0]["document_id"]) + + +async def _snapshot(document_id: str) -> dict[str, list[dict[str, Any]]]: + chunks = await ContractDatabase.fetch_all( + """ + SELECT chunk_id, chunk_type, content, content_search_text, + path_search_text, term_search_text, source_chunk_path, + file_path, sort_order + FROM document_chunks + WHERE document_id = :document_id + ORDER BY sort_order, chunk_id, id + """, + {"document_id": document_id}, + ) + units = await ContractDatabase.fetch_all( + """ + SELECT s.section_path, m.unit_kind, m.path_token_count, + m.content_token_count, m.has_image, m.has_table, m.sort_order + FROM document_map_units AS m + JOIN document_sections AS s ON s.section_id = m.section_id + WHERE m.document_id = :document_id + ORDER BY m.sort_order, s.section_path, m.unit_kind, m.id + """, + {"document_id": document_id}, + ) + tokens = await ContractDatabase.fetch_all( + """ + SELECT s.section_path, m.unit_kind, m.sort_order, + t.channel, t.token, t.token_hash, t.frequency + FROM document_map_unit_tokens AS t + JOIN document_map_units AS m ON m.id = t.map_unit_id + JOIN document_sections AS s ON s.section_id = m.section_id + WHERE m.document_id = :document_id + ORDER BY m.sort_order, s.section_path, m.unit_kind, t.channel, t.token + """, + {"document_id": document_id}, + ) + indexes = await ContractDatabase.fetch_all( + """ + SELECT format_version, unit_count, token_count, + average_idf_path, average_idf_content, + path_document_count, path_total_length, + content_document_count, content_total_length + FROM document_map_unit_indexes + WHERE document_id = :document_id + """, + {"document_id": document_id}, + ) + units = sorted(units, key=lambda row: tuple(str(value) for value in row.values())) + tokens = sorted( + tokens, key=lambda row: tuple(str(value) for value in row.values()) + ) + return {"chunks": chunks, "units": units, "tokens": tokens, "indexes": indexes} + + +@pytest.mark.asyncio +async def test_demo_preparation_cache_preserves_persisted_semantics_across_modes( + developer_api_client_factory: Callable[ + [], AbstractAsyncContextManager[AsyncClient] + ], + monkeypatch: MonkeyPatch, +) -> None: + monkeypatch.setenv("DEMO_CANONICAL_BUNDLE_ENABLED", "true") + monkeypatch.setenv("DEMO_PUBLICATION_PREPARATION_CACHE_ENABLED", "false") + async with developer_api_client_factory() as client: + from shared.services.retrieval.publication_preparation_cache import ( + PublicationPreparationCache, + ) + + PublicationPreparationCache.clear() + cold_off_id = await _publish(client, namespace="prep-cache-off") + off_snapshot = await _snapshot(cold_off_id) + + import shared.core.config as config_module + + monkeypatch.setattr( + config_module.settings, + "DEMO_PUBLICATION_PREPARATION_CACHE_ENABLED", + True, + ) + PublicationPreparationCache.clear() + cache_counts = {"hits": 0, "misses": 0} + original_get = PublicationPreparationCache._get + + def observe_get(key: object) -> object: + value = original_get(key) # type: ignore[arg-type] + cache_counts["hits" if value is not None else "misses"] += 1 + return value + + monkeypatch.setattr( + PublicationPreparationCache, + "_get", + classmethod(lambda _cls, key: observe_get(key)), + ) + cold_on_id = await _publish(client, namespace="prep-cache-cold") + cold_cache_hits = cache_counts["hits"] + cold_cache_misses = cache_counts["misses"] + warm_on_id = await _publish(client, namespace="prep-cache-warm") + assert cold_cache_misses > 0 + assert cache_counts["hits"] > cold_cache_hits + assert cache_counts["misses"] == cold_cache_misses + assert len(PublicationPreparationCache._entries) > 0 + cold_snapshot = await _snapshot(cold_on_id) + warm_snapshot = await _snapshot(warm_on_id) + + monkeypatch.setattr( + config_module.settings, + "DEMO_PUBLICATION_PREPARATION_CACHE_ENABLED", + False, + ) + rollback_id = await _publish(client, namespace="prep-cache-rollback") + rollback_snapshot = await _snapshot(rollback_id) + + for namespace, expected_document_id in ( + ("prep-cache-off", cold_off_id), + ("prep-cache-cold", cold_on_id), + ("prep-cache-warm", warm_on_id), + ("prep-cache-rollback", rollback_id), + ): + retrieval = await client.post( + "/api/v1/retrieval/query", + json={ + "namespace": namespace, + "query": "Tesla investment", + "top_k": 3, + "use_agentic": False, + }, + ) + assert retrieval.status_code == 200 + result_documents = { + str(result["source"]["document_id"]) + for result in cast(list[dict[str, Any]], retrieval.json()["results"]) + } + assert result_documents + assert result_documents == {expected_document_id} + + assert cold_snapshot == off_snapshot + assert warm_snapshot == off_snapshot + assert rollback_snapshot == off_snapshot diff --git a/apps/api/tests/contract/test_demo_publication_trace_owner_contract.py b/apps/api/tests/contract/test_demo_publication_trace_owner_contract.py index c457e96db..e6d21e389 100644 --- a/apps/api/tests/contract/test_demo_publication_trace_owner_contract.py +++ b/apps/api/tests/contract/test_demo_publication_trace_owner_contract.py @@ -21,13 +21,14 @@ ensure_import_paths() from app.services.demo.source_catalog import DemoSourceCatalog, DemoSourceDefinition # noqa: E402 +from app.services.demo.canonical_bundle_result import CanonicalDemoBundle # noqa: E402 from app.services.demo.source_materializer import DemoSourceMaterializer # noqa: E402 from shared.models.database.demo_materialization import DemoMaterialization # noqa: E402 from shared.models.database.job import Job # noqa: E402 from shared.services.jobs.lifecycle.publication_trace import PublicationTrace # noqa: E402 from shared.services.retrieval.publication_models import PublishedDocumentState # noqa: E402 from shared.services.retrieval.publication_service import RetrievalPublicationService # noqa: E402 -from shared.services.redis import RedisPublicationSemaphore # noqa: E402 +from shared.services.redis import RedisPublicationSemaphore, RedisService # noqa: E402 class FakeCatalog: @@ -138,6 +139,7 @@ def make_materializer( materializer = object.__new__(materializer_type) materializer._catalog = cast(DemoSourceCatalog, FakeCatalog()) materializer._publication_service = cast(RetrievalPublicationService, publication) + materializer._redis_service = cast(RedisService, object()) return materializer @@ -164,11 +166,16 @@ async def test_demo_owner_finishes_trace_at_publication_commit_or_rollback( updated_at=None, ) - def use_bundle(**kwargs: object) -> dict[str, int | str]: - return { - "zip_key": "zip/123", - "zip_size": 123, - } + async def use_bundle(**kwargs: object) -> CanonicalDemoBundle: + assert kwargs["demo_source_id"] == source.demo_source_id + assert kwargs["source_directory"] == Path("/tmp") + return CanonicalDemoBundle( + zip_key="zip/123", + zip_size=123, + raw_prefix="results/demo-canonical/", + content_version="v1", + reused=False, + ) monkeypatch.setattr(materializer_module, "_upload_demo_result_bundle", use_bundle) monkeypatch.setattr( @@ -179,7 +186,11 @@ def use_bundle(**kwargs: object) -> dict[str, int | str]: monkeypatch.setattr( materializer_module, "settings", - SimpleNamespace(KNOWHERE_PUBLICATION_TRACE_ENABLED=True), + SimpleNamespace( + KNOWHERE_PUBLICATION_TRACE_ENABLED=True, + DEMO_CANONICAL_BUNDLE_ENABLED=True, + DEMO_PUBLICATION_PREPARATION_CACHE_ENABLED=False, + ), ) terminal_payloads: list[dict[str, object]] = [] diff --git a/apps/api/tests/migrations/test_legacy_token_lookup_retirement_contract.py b/apps/api/tests/migrations/test_legacy_token_lookup_retirement_contract.py new file mode 100644 index 000000000..f3ff96473 --- /dev/null +++ b/apps/api/tests/migrations/test_legacy_token_lookup_retirement_contract.py @@ -0,0 +1,148 @@ +"""Real PostgreSQL contracts for gated and replayable token index maintenance.""" + +from __future__ import annotations + +import argparse +import json +import os +from pathlib import Path +import subprocess +import sys + +import pytest +from alembic import command +from alembic.config import Config +from sqlalchemy import text +from sqlalchemy.engine import Engine + +from shared.services.retrieval.maintenance.legacy_token_lookup import ( + LegacyTokenLookupMaintenance, +) + +_ROOT: Path = Path(__file__).resolve().parents[4] +_PARENT_REVISION: str = "2b3c4d5e6f7a" +_REVISION: str = "3c4d5e6f7a8b" +_LEGACY_INDEX: str = "idx_document_map_unit_tokens_lookup" +_BINARY_INDEX: str = "idx_document_map_unit_tokens_token_lookup_binary" +_UNIT_LOOKUP_INDEX: str = "idx_document_map_unit_tokens_unit_lookup" + + +def _build_config() -> Config: + config: Config = Config(str(_ROOT / "apps/api/alembic.ini")) + config.set_main_option("script_location", str(_ROOT / "apps/api/alembic")) + return config + + +def _inspect_indexes(engine: Engine) -> dict[str, dict[str, str | bool]]: + with engine.connect() as connection: + return LegacyTokenLookupMaintenance(connection).inspect_indexes() + + +def _run_operator(*arguments: str) -> subprocess.CompletedProcess[str]: + return subprocess.run( + [sys.executable, str(_ROOT / "scripts/retire_legacy_token_lookup.py"), *arguments], + env=os.environ.copy(), capture_output=True, text=True, timeout=30, check=False, + ) + + +def test_should_keep_default_indexes_then_apply_and_restore_at_same_head( + alembic_engine: Engine, +) -> None: + config: Config = _build_config() + command.upgrade(config, "heads") + initial: dict[str, dict[str, str | bool]] = _inspect_indexes(alembic_engine) + assert len(initial) == 5 + assert _LEGACY_INDEX in initial + inspection: subprocess.CompletedProcess[str] = _run_operator() + assert inspection.returncode == 0, inspection.stderr + assert json.loads(inspection.stdout) == initial + + for _ in range(2): + result: subprocess.CompletedProcess[str] = _run_operator("--apply") + assert result.returncode == 0, result.stderr + assert json.loads(result.stdout) == { + name: state for name, state in initial.items() if name != _LEGACY_INDEX + } + with alembic_engine.connect() as connection: + assert connection.scalar(text("SELECT version_num FROM alembic_version")) == _REVISION + for _ in range(2): + result = _run_operator("--restore") + assert result.returncode == 0, result.stderr + assert json.loads(result.stdout) == initial + + +@pytest.mark.parametrize("is_external_transaction", (False, True)) +def test_should_gate_upgrade_and_restore_on_downgrade( + alembic_engine: Engine, is_external_transaction: bool, +) -> None: + config: Config = _build_config() + command.upgrade(config, _PARENT_REVISION) + initial: dict[str, dict[str, str | bool]] = _inspect_indexes(alembic_engine) + config.cmd_opts = argparse.Namespace(x=["retire_legacy_token_lookup=true"]) + if is_external_transaction: + with alembic_engine.begin() as connection: + config.attributes["connection"] = connection + command.upgrade(config, _REVISION) + config.attributes.pop("connection") + else: + command.upgrade(config, _REVISION) + assert _inspect_indexes(alembic_engine) == { + name: state for name, state in initial.items() if name != _LEGACY_INDEX + } + if is_external_transaction: + with alembic_engine.begin() as connection: + config.attributes["connection"] = connection + command.downgrade(config, _PARENT_REVISION) + else: + command.downgrade(config, _PARENT_REVISION) + assert _inspect_indexes(alembic_engine) == initial + + +@pytest.mark.parametrize( + ("index_name", "damage"), + ( + (_BINARY_INDEX, "missing"), + (_UNIT_LOOKUP_INDEX, "incompatible"), + (_BINARY_INDEX, "invalid"), + (_UNIT_LOOKUP_INDEX, "not_ready"), + ), +) +def test_should_refuse_retirement_without_compatible_replacements( + alembic_engine: Engine, index_name: str, damage: str, +) -> None: + command.upgrade(_build_config(), "heads") + with alembic_engine.begin() as connection: + if damage in {"missing", "incompatible"}: + connection.execute(text(f"DROP INDEX {index_name}")) + if damage == "incompatible": + connection.execute(text( + f"CREATE INDEX {index_name} ON document_map_unit_tokens (channel)" + )) + else: + column: str = "indisvalid" if damage == "invalid" else "indisready" + connection.execute(text( + f"UPDATE pg_index SET {column} = false " + "WHERE indexrelid = to_regclass(:name)" + ), {"name": index_name}) + before: dict[str, dict[str, str | bool]] = _inspect_indexes(alembic_engine) + result: subprocess.CompletedProcess[str] = _run_operator("--apply") + assert result.returncode == 1 + assert index_name in result.stderr + assert _inspect_indexes(alembic_engine) == before + + +def test_should_refuse_to_replace_incompatible_legacy_index( + alembic_engine: Engine, +) -> None: + command.upgrade(_build_config(), "heads") + with alembic_engine.begin() as connection: + connection.execute(text(f"DROP INDEX {_LEGACY_INDEX}")) + connection.execute(text( + f"CREATE INDEX {_LEGACY_INDEX} ON document_map_unit_tokens (channel)" + )) + before: dict[str, dict[str, str | bool]] = _inspect_indexes(alembic_engine) + for argument in ("--restore", "--apply"): + result: subprocess.CompletedProcess[str] = _run_operator(argument) + assert result.returncode == 1 + assert _LEGACY_INDEX in result.stderr + assert _inspect_indexes(alembic_engine) == before diff --git a/apps/worker/tests/contract/test_page_memory_retrieval_contract.py b/apps/worker/tests/contract/test_page_memory_retrieval_contract.py index bfebab1d1..f42f219b7 100644 --- a/apps/worker/tests/contract/test_page_memory_retrieval_contract.py +++ b/apps/worker/tests/contract/test_page_memory_retrieval_contract.py @@ -267,6 +267,25 @@ def test_result_storage_allows_page_citation_and_page_pdf_artifact_refs_not_debu ) +def test_result_storage_uses_validated_canonical_raw_prefix_without_job_fallback() -> None: + storage = JobResultStorage(results_bucket="test-results") + + assert ( + storage.build_raw_key( + job_id="job-materialization", + raw_prefix="results/demo-canonical/source-a/content-v1/", + relative_path="images/chart.png", + ) + == "results/demo-canonical/source-a/content-v1/images/chart.png" + ) + with pytest.raises(ValueError, match="under results"): + storage.build_raw_key( + job_id="job-materialization", + raw_prefix="uploads/shared/", + relative_path="images/chart.png", + ) + + def test_result_storage_upload_filters_to_referenced_artifacts(tmp_path) -> None: class FakeStorageAdapter: def __init__(self) -> None: diff --git a/docs/design/canonical-demo-bundle-reuse.md b/docs/design/canonical-demo-bundle-reuse.md new file mode 100644 index 000000000..98e3d8336 --- /dev/null +++ b/docs/design/canonical-demo-bundle-reuse.md @@ -0,0 +1,136 @@ +# Canonical demo bundle reuse + +## Scope + +Demo materialization reuses the same immutable ZIP and raw artifacts across +users and namespaces. Jobs, job results, documents, publication transactions, +authorization and retrieval isolation remain owner scoped. Normal Worker parsing +continues to store its result artifacts by job identifier. + +The previous demo path compressed and uploaded every artifact for each new job. +Two real local SpaceX materializations created 637 objects each and took 40.422 s +and 39.098 s from HTTP request to response. These are local end-to-end timings, +not production measurements or isolated publication transaction timings. +Each request duplicated 439,190,627 stored bytes, including a 215,559,924-byte ZIP. + +After the change, the first canonical upload took 40.211 s. Two subsequent +materializations into separate namespaces took **5.483 s** and **5.337 s**, +respectively (about 86% lower than the two legacy samples). Both added zero +objects and rewrote zero objects; all 638 canonical objects retained their +ETags, sizes and modification times. The extra object is the completion marker. +This removes about 439 MB of duplicate object writes per repeat materialization. + +Both repeat requests passed real HTTP checks of document image/table/page +assets, classic retrieval page/image assets, the public job ZIP URL, and archive +readback. Archiving one materialization left the other's shared assets readable. + +## Storage contract + +The content version hashes sorted relative file paths, lengths and file bytes, +including a bundle-format version. A bounded process-local memo avoids rereading +unchanged file bytes; changes to a file's size, mtime or ctime invalidate it. +Symbolic links and sources that change during hashing/upload are rejected. + +Keys preserve the existing result storage layout: + +```text +results/demo-canonical//.zip +results/demo-canonical/// +results/demo-canonical///.bundle-ready.json +``` + +The marker is written only after ZIP and raw uploads complete. A failed partial +upload is never reused. A subsequent request can retry it. A renewable Redis +lease serializes first uploads across API processes; release compares the owner +token and an abandoned lease expires. Compression, hashing and storage calls run +off the API event loop. Cancellation waits for an active upload thread before +releasing its lease. + +Reuse checks the completion marker and ZIP size without issuing one storage +request per raw artifact. If a raw object is deleted after completion, reuse +does not detect or repair it automatically. Recovery requires invalidating that +version's `.bundle-ready.json` and rerunning the preparation command against the +same source version. Protect the shared prefix from per-job lifecycle cleanup. + +`JobResult.result_s3_key` records the canonical ZIP. +`JobResult.document_metadata.result_raw_prefix` records the canonical raw prefix. +Artifact readers use that metadata when present and retain the job-key fallback +for older results and ordinary Worker output. No database migration is required. + +## DevOps enablement and rollback + +1. Deploy the writer and compatible artifact readers together. The independent + `DEMO_CANONICAL_BUNDLE_ENABLED` setting defaults to `false`. +2. Prepare bundles with the same deployed source files and result bucket, using + the [preparation command](canonical-demo-preparation.md). This moves the + initial ZIP compression and artifact upload ahead of user requests. A new + unprepared content version still needs its first upload. +3. Enable `DEMO_CANONICAL_BUNDLE_ENABLED=true` on API instances after checking + access to their configured result bucket and Redis. This does not change + `KNOWHERE_PUBLICATION_STRATEGY` or retire any database index. +4. Materialize the same source into two namespaces. Verify identical ZIP and + raw keys, and unchanged object timestamps/ETags on the second request. Fetch + image/table/page artifacts and the ZIP through their public response URLs. +5. Roll back new writes by setting `DEMO_CANONICAL_BUNDLE_ENABLED=false` and + restarting the API. Keep the compatible readers deployed: they continue to + resolve existing canonical results while new jobs use individual bundles. +6. Do not delete canonical objects on per-document archive or expiration. A + source/version can have references from multiple owners. Bucket lifecycle + rules must preserve this prefix for the full lifetime of its references. + +A binary rollback to a release without prefix-aware artifact readers requires +first restoring per-job raw objects for every canonical result, or retaining +the reader compatibility patch. Turning the switch off does not rewrite +existing result records. + +The separate `DEMO_PUBLICATION_PREPARATION_CACHE_ENABLED` flag also defaults to +`false`. It requires canonical bundles and reuses only pure lexical preparation +inside the current API process. See [cache boundaries and measurements](demo-publication-preparation-cache.md). + +[Legacy token lookup retirement](legacy-token-lookup-retirement.md) is an +independent database operation. Normal deployment preserves that index; enabling +either demo flag does not change it. Index restoration uses the operator command +at the same Alembic head and can take substantially longer than retirement. + +## Verification boundary + +Contract coverage targets reuse across namespaces, content changes, failed +uploads, concurrent creation and legacy artifact compatibility. Runtime +verification must also exercise actual HTTP, Redis/Celery Worker execution, +real parsing, publication, retrieval and archive. Local results do not establish +production latency or performance at a larger database size. + +The real HTTP verifier caught a missing raw-prefix projection in classic +`map_unit_discovery` hydration: materialization and document asset endpoints +succeeded, while a retrieved page image returned 404. The query now projects +the artifact prefix alongside the job identifier. Artifact location is also +carried through connected/reference hydration, small-corpus loading and the +agent read/grep/table tools; their storage reads share the same prefix support. + +The actual Worker Markdown regression used presigned local S3 upload and +confirm-upload, followed by Redis/Celery dispatch, the real parser, publication, +classic retrieval and archive. It completed in 3.240 s observed by the client; +the parser trace was 2.307 s. This small document verifies the ordinary Worker +path and is not a SpaceX parse benchmark. +After the canonical-reader changes, the Worker was restarted and the same real +queued path passed again in 3.278 s (job duration 2.906 s). All runtime test +documents were archived. Production configuration and indexes were untouched. + +After the preparation-cache changes, a fresh API and Worker processed another +HTTP-created Markdown job in 4.300 s client-observed time (3.814 s job duration). +The fresh Worker log confirmed ownership; two chunks, classic retrieval and +archive passed. The configured summarization provider returned HTTP 401, so this +run exercised the parser's fallback. It does not validate successful external +LLM/PDF parsing or full SpaceX parsing latency. + +Focused verification passed: six canonical/trace contracts, eight storage and +document lifecycle contracts, and 48 existing retrieval checks. The canonical +HTTP contract also exercises classic retrieval's shared asset URL after the +raw-SQL projection fix. Ruff, targeted Pyright and `git diff --check` passed. + +Run the canonical regression contracts with: + +```sh +uv run pytest -q apps/api/tests/contract/test_canonical_demo_bundle_contract.py \ + apps/api/tests/contract/test_demo_publication_trace_owner_contract.py +``` diff --git a/docs/design/canonical-demo-preparation.md b/docs/design/canonical-demo-preparation.md new file mode 100644 index 000000000..a97d8d827 --- /dev/null +++ b/docs/design/canonical-demo-preparation.md @@ -0,0 +1,27 @@ +# Prepare canonical demo bundles + +`prepare_demo_bundles.py` uploads the immutable ZIP and raw artifacts for +catalog demo sources ahead of user materialization. It uses the same content +version, Redis lease and completion marker as the API path. It does not create +jobs, documents, credits or publication state. + +Run from the repository root with the API's normal storage and Redis environment +configured (`S3_TYPE`, result bucket, endpoint and credentials, `REDIS_HOST` and +`REDIS_PORT`). The source checkout must include `apps/api/app/data/demo_documents`. + +```sh +uv run --package knowhere-api-app python apps/api/scripts/prepare_demo_bundles.py demo-spacex-s1 +uv run --package knowhere-api-app python apps/api/scripts/prepare_demo_bundles.py demo-spacex-s1 --upload +uv run --package knowhere-api-app python apps/api/scripts/prepare_demo_bundles.py --all --upload +``` + +The default invocation lists selected sources and makes no storage or Redis +calls. `--upload` is required for writes. A completed version is reported as +`reused` and uploads no objects. New versions are reported as `created` with +their exact ZIP key and raw prefix. A failure exits nonzero; rerunning is safe +because the completion marker is written only after all artifacts finish. + +Use the same configured result bucket as the API. The +`DEMO_CANONICAL_BUNDLE_ENABLED` API setting controls new materializations, not +this explicit preparation command. Keep the prefix-aware readers deployed for +existing canonical results. diff --git a/docs/design/demo-publication-preparation-cache.md b/docs/design/demo-publication-preparation-cache.md new file mode 100644 index 000000000..6c9039194 --- /dev/null +++ b/docs/design/demo-publication-preparation-cache.md @@ -0,0 +1,61 @@ +# Demo publication preparation cache + +## Purpose + +Canonical demo artifacts can be reused across user namespaces. Publishing each +copy still tokenizes the same parsed text and builds the same token-frequency +counts. A bounded process-local cache can reuse those pure results without +changing publication ownership or retrieval semantics. + +## Boundary + +The cache is enabled only when `DEMO_PUBLICATION_PREPARATION_CACHE_ENABLED=true` +and the canonical bundle supplies a content version. Its scope identifies the +source, content version, map-unit index format, and preparation algorithm +version. Other publication paths continue through the authoritative tokenizers. + +Cache values are immutable token sequences or frequency pairs. Callers receive +fresh mutable copies. Document IDs, job result IDs, section IDs, map-unit IDs, +namespace, transaction state, scores, and retrieval results never enter the +cache. Search-field construction and map-unit persistence remain in the shared +publication path and retain the same ordering and transaction boundary. + +The process cache is limited to 64 MiB and 4,096 entries. Eviction only causes +recomputation. A changed source content version, index format, or preparation +algorithm does not reuse an earlier entry. The default flag is off, so rollback +is a configuration change and existing persisted documents remain readable. + +## Verification + +The contract publishes the same canonical demo source in separate namespaces +with the flag off, on with an empty cache, on with a warm cache, and off again. +It compares persisted chunk search fields, map-unit semantics, token-frequency +rows, and IDF statistics after excluding generated database identifiers. It +also checks namespace-scoped retrieval and confirms no document crosses scope. + +This verifies semantic parity. Performance admission still requires a measured +publication stage improvement and end-to-end timing on production-shaped input. + +## Focused measurements + +The catalog's real SpaceX demo (227 chunks) was exercised through the real HTTP +materializer against the small local database, with the canonical bundle +already present. The separate 12.42-million-token clone experiment was not +repeated for this cache: + +| Mode | Request 1 | Request 2 | Bundle writes | +| --- | ---: | ---: | --- | +| Preparation cache off | 4.887 s | 5.051 s | 0 new or changed objects | +| Preparation cache on | 7.044 s (cold) | 2.967 s (warm) | 0 new or changed objects | + +Both modes used the same canonical ZIP and raw prefix. Page and image assets, +classic retrieval, and archival completed successfully. The warm sample is an +observed local improvement; the cold sample includes cache population and was +slower. These are focused observations, not a production p95 claim. + +The final Worker owner check used a small Markdown upload because the local +provider credentials do not establish successful external LLM or PDF parsing. +The Worker completed job `job_f5d5c6d4ed2b` in 3.814 s with a 236 ms publication +trace, and retrieval plus archival passed. The provider returned HTTP 401 and +the parser completed through its fallback; this validates the Worker publication boundary and +fallback behavior, not external model quality. diff --git a/docs/design/legacy-token-lookup-retirement.md b/docs/design/legacy-token-lookup-retirement.md new file mode 100644 index 000000000..b7ef66124 --- /dev/null +++ b/docs/design/legacy-token-lookup-retirement.md @@ -0,0 +1,122 @@ +# Legacy token lookup retirement + +## Purpose and evidence + +Retire only the nonunique `idx_document_map_unit_tokens_lookup` index to reduce +token persistence I/O. Keep the primary key, binary token lookup, unit index, +and unit-leading covering lookup. + +The production incident's token COPY samples waited on `IO:DataFileRead`. +In the disposable 12.42-million-row clone, removing this legacy index reduced +COPY buffer reads by approximately 53%, WAL by 47%, database COPY execution time +by 48%, and application token-stage time by 38%. Publication improved from +11.376 to 10.182 seconds in that pair; full HTTP time stayed around 54 seconds. +These are three targeted cold-cache publication observations, not production +latency guarantees. See [the experiment report](publication-scale-cold-cache-findings.md). + +Current classic discovery and recall use binary lookup. The remaining plain +token-hash reader in `scoring/map_lighting.py` also scopes by map-unit IDs and +channel, which the retained unit-leading index covers. Ten frozen query cases +returned identical counts and hashes before and after retirement; each used +the retained unit lookup. Historical application versions and every possible +query shape have not been validated. + +## Deployment behavior + +Revision `3c4d5e6f7a8b`, following `2b3c4d5e6f7a`, retains the legacy index by +default. No runtime application flag deletes or recreates this index. The +previous retirement migration `2f3g4h5i6j7` is unconditional and concerns two +different indexes; this change does not alter that historical migration. + +Before the new revision is recorded, an explicit migration invocation can opt in: + +```bash +cd apps/api +uv run alembic -x retire_legacy_token_lookup=true upgrade head +``` + +An ordinary deployment records the revision without retirement. Re-running +`upgrade head` with that argument later cannot execute an already-recorded +revision. Use the operator command below to activate at the same head, including +after later migrations. Do not downgrade the application schema for activation. + +## Operator procedure + +Supply `DATABASE_URL` through the protected environment. The command also honors +`DB_SSL_MODE`, `DB_SSL_CERT`, `DB_SSL_KEY`, and `DB_SSL_ROOT_CERT`; absent explicit +values, connection URL and libpq defaults apply. Do not put credentials in shell +arguments or reports. Run commands from the repository root with workspace +dependencies installed. + +1. Confirm the intended database and the application versions needed for rollback. + For production, obtain explicit approval for this index change before applying + it. The investigation and preparation do not authorize a production mutation. +2. Inspect the catalog without changing indexes: + + ```bash + uv run --package knowhere-api-app python scripts/retire_legacy_token_lookup.py + ``` + +3. Explicitly retire the legacy lookup: + + ```bash + uv run --package knowhere-api-app python scripts/retire_legacy_token_lookup.py --apply + ``` + +4. Review the returned catalog. The legacy lookup must be absent; the primary key, + binary lookup, unit index, and unit lookup must remain valid and ready. Check a + real publication and retrieval before deciding to retain the change. Observe + token persistence timing and database I/O; an HTTP success alone does not + establish an improvement. + +The helper verifies the binary and unit-leading lookup definitions and their +`indisvalid` and `indisready` states before removing anything. It also validates +the legacy definition before deletion. A same-name partial, incompatible, or +invalid index fails closed. Repeating apply is safe when the legacy index is +already absent and both replacements remain compatible. Only the named legacy +index is targeted; no table data or uniqueness constraint is removed. + +Run one maintenance operation at a time and coordinate other schema changes. +Independent migration and operator connections use concurrent DDL. The operator +sets a five-second lock timeout and a 30-minute statement timeout. Adjust the +latter with `--statement-timeout-seconds` if the rebuild needs a longer window. +Concurrent DDL can still wait for transactions and consume substantial I/O. +Externally owned migration transactions use regular DDL for compatibility with +contract fixtures; do not wrap production index maintenance in such a transaction. + +## Rollback + +Restore the exact old nonunique index independently of the Alembic head: + +```bash +uv run --package knowhere-api-app python scripts/retire_legacy_token_lookup.py --restore +``` + +This creates concurrently: + +```sql +CREATE INDEX idx_document_map_unit_tokens_lookup +ON public.document_map_unit_tokens USING btree (channel, token_hash, map_unit_id); +``` + +The command reads back and verifies the definition, validity, and readiness. +Repeating restore with a compatible valid index is a no-op. Downgrading the new +revision also restores the index, but is unnecessary for this operator rollback. +Application rollback alone does not restore a dropped index. + +Rebuilding scans the token table and consumes storage, CPU, I/O, and WAL; it may +take substantially longer than dropping the index. Allow space for the index and +temporary build work, and account for replica lag. If a concurrent operation is +interrupted, inspect the catalog. An invalid same-name index is deliberately +rejected, not silently reused or automatically deleted. Resolve that catalog +state as a separately reviewed operator action, then retry restoration. Until +the rebuild succeeds, the previous index configuration has not been restored. + +## Verification boundary + +Disposable PostgreSQL contracts cover default retention, opt-in migration, +independent and externally owned transactions, operator activation after the +revision is already recorded, idempotent apply/restore, exact restoration, and +refusal when replacement indexes are missing, incompatible, invalid, or not +ready. They also reject an incompatible legacy index. No production index is +changed by these tests or by adding this revision. diff --git a/docs/design/publication-scale-cold-cache-findings.md b/docs/design/publication-scale-cold-cache-findings.md new file mode 100644 index 000000000..95cb02753 --- /dev/null +++ b/docs/design/publication-scale-cold-cache-findings.md @@ -0,0 +1,173 @@ +# Publication at production scale with a cold database cache + +## Scope and environment + +This investigation follows the runtime repair in PR #448. We ran the actual +API application and Celery Worker against a disposable PostgreSQL clone on the +Mac benchmark host. Redis and S3 were isolated local services. Production was +inspected read-only; no production configuration, data, or indexes were changed. + +The clone was copied from a cleanly stopped template mounted read-only. We +applied current migrations and restored the two missing token indexes before +comparison. Historical catalog evidence from 2026-09-30 02:07:35 UTC confirms +that the five-index configuration also existed before the production failure. + +| Property | Clone configuration | +| --- | --- | +| Token rows, catalog estimate | 12,423,317 | +| Token heap | 2,385,043,456 bytes | +| Initial total token relation including indexes | 8,440,152,064 bytes | +| PostgreSQL | 15.19 | +| Container memory / swap limit | 1 GiB / no additional swap | +| CPU limit | One CPU | +| Shared buffers | 128 MiB | +| Work memory / maintenance work memory | 4 MiB / 64 MiB | +| Effective cache size | 512 MiB, planner estimate only | +| Maximum connections | 40 | +| Durability | fsync, full_page_writes, synchronous_commit enabled | +| I/O timing | Enabled | +| API statement / command timeout | 30 seconds / 30 seconds | + +Before each publication sample, we stopped only the disposable database, used +`POSIX_FADV_DONTNEED` on its data files, and restarted it. This clears PostgreSQL +buffers and requests eviction of the clone's Linux file cache. We did not clear +global host caches or restart the sealed source. macOS/Colima virtual-disk caches +remain uncontrolled. The API and Worker ran on the development machine through +an SSH database tunnel, so application timing includes network round trips. + +Each materialization used the same `demo-spacex-s1` input: 227 chunks and 10,778 +token rows. We captured `pg_stat_statements`, per-index read counters, publication +traces, and sampled wait events. Samples used separate namespaces; successfully +published documents were retrieved and archived. The same clone was reused, +with cache eviction between samples; earlier successful samples leave a small +amount of additional archived data. These are targeted observations, not +statistical latency percentiles or identical snapshots. + +## Batch repair comparison + +Both variants used candidate publication and the same five token indexes: +primary key, `lookup`, `token_lookup_binary`, `unit`, and `unit_lookup`. + +| Measurement | Before PR #448 | After PR #448 | +| --- | ---: | ---: | +| Complete HTTP materialization | 50.943 s | 54.149 s | +| Publication transaction | 10.416 s | 11.376 s | +| Token persistence stage | 3.652 s | 4.228 s | +| Token COPY calls | 1 | 11 | +| Longest token command, application measurement | 3.551 s | 0.482 s | +| Token COPY execution, database measurement | 3.552 s | 3.848 s | +| Token COPY buffer reads | 25,639 | 25,673 | +| Token COPY read time | 3.092 s | 3.148 s | +| Token COPY WAL bytes | 159,152,618 | 160,566,253 | + +Both requests succeeded with classic retrieval and archive. Neither reproduced +the production 30-second timeout. The first sample attributed 12 active samples +to token COPY waiting on `DataFileRead`, consistent with the database timing. + +The batching change bounded individual command duration but did not reduce +storage reads. In this pair, token-stage duration increased approximately 16% +and publication duration approximately 9%. Database execution increased by +approximately 0.30 seconds; additional connection round trips contribute to +the larger application-observed difference. A single pair cannot establish a +stable regression percentage or a production throughput estimate. + +This also corrects the interpretation of the earlier small-database run: +publication below ten seconds did not establish that latency with a large, +cold database. Full HTTP time includes preparation outside publication. + +## Attribution and a focused index experiment + +The first publication produced these per-index buffer-read deltas: + +| Token index | Read blocks | +| --- | ---: | +| Primary key | 148 | +| Legacy `lookup` | 13,435 | +| `token_lookup_binary` | 11,015 | +| `unit` | 400 | +| `unit_lookup` | 1,416 | + +These counters cover the publication observation window and are separate from +COPY's statement-level accounting. Both token-leading indexes dominate reads; +legacy `lookup` accounts for approximately half of observed index reads. + +Current classic discovery and `corpus.recall` use the binary token hash +predicate. Oversized outline/node filtering still has a normal-hash reader in +`scoring/map_lighting.py`, constrained by map-unit IDs and channel. The retained +`unit_lookup` covering index supports that reader. + +Ten read-only production EXPLAIN plans, spanning scopes of 1, 32, 227, 1,000, +and 2,000 map units and two token sets, selected `unit_lookup`. On the clone, +we froze equivalent fixtures and compared exact result counts and SHA-256 +digests before and after removing only the nonunique legacy `lookup` index. +All ten cases matched; every plan continued to use `unit_lookup`. These checks +do not exercise every historical application query or a full agentic episode. + +### Cold publication after removing legacy lookup + +The final sample retained the patched application and all other settings, and +removed only `idx_document_map_unit_tokens_lookup` in the disposable clone. + +| Measurement | Patched, five indexes | Patched, legacy lookup removed | +| --- | ---: | ---: | +| Complete HTTP materialization | 54.149 s | 54.177 s | +| Publication transaction | 11.376 s | 10.182 s | +| Token persistence stage | 4.228 s | 2.624 s | +| Token COPY execution, database measurement | 3.848 s | 1.989 s | +| Longest token command, application measurement | 0.482 s | 0.290 s | +| Token COPY buffer reads | 25,673 | 12,110 | +| Token COPY read time | 3.148 s | 1.618 s | +| Token COPY WAL bytes | 160,566,253 | 85,682,583 | + +Publication, chunk reads, classic retrieval, and archive all succeeded. Token +buffer reads fell approximately 53%, WAL approximately 47%, database COPY time +approximately 48%, and application token-stage duration approximately 38%. +Publication improved approximately 10.5% in this comparison, while full HTTP +latency remained essentially unchanged. Preparation and other stages still +matter, and the publication transaction remains slightly above ten seconds. + +This isolates a concrete write-cost reduction. It supports preparing legacy +lookup retirement as a separate index change while retaining binary lookup, +unit lookup, the unit index, and the primary key. It does not justify removing +binary lookup, changing retrieval semantics, or claiming production latency. +Before a production index change, account for application versions supported +by rollback and the time/storage required to rebuild the old index. The local +query checks do not cover every scope size or historical query shape. + +## Worker verification + +An HTTP-created Markdown ingestion job traversed real upload, confirmation, +Redis/Celery dispatch, gevent Worker startup, parsing, and publication against +the full five-index clone. It completed in 7.707 seconds end to end; Worker +processing reported 6.519 seconds. Both chunks were readable, classic retrieval +returned the document, and archive succeeded. This checks Worker compatibility +against the large database; the small input is not a full-PDF throughput test. +Copied job callback URLs were cleared in the disposable clone before Worker +startup to isolate external delivery. + +## Limits and evidence + +Production's observed token relation occupied approximately 15.54 GB after +recovery, versus 8.44 GB in the rebuilt clone. Row counts are comparable, but +index density, physical size, churn, and storage behavior are different. +The clone's first token COPY averaged approximately 0.121 ms per PostgreSQL +buffer read. Historical Aurora read latency was approximately 0.8 ms at one-minute +resolution; these are different measurement layers and cannot be substituted +into a precise production latency prediction. Aurora scale-up is not simulated. + +Local artifacts are under `/tmp/knowhere-publication-scale/`: before/after HTTP +results, database snapshots, trace summaries, `comparison.json`, index alignment +DDL, historical inventory evidence, and map-lighting parity results. Environment, +authentication, and frozen data fixtures are private local artifacts and are not +included in this document. No production deployment or index retirement is +performed by this investigation. + +## Experiment completion + +Only three large-document publication samples were run: old candidate batching, +patched batching, and patched batching without legacy lookup. There was one +small real Worker job and ten narrowly scoped query parity checks. No broad +benchmark suite or full PDF reparse was started. The disposable database is +retained stopped with the four-index experimental configuration for follow-up; +the source/template are unchanged. Runtime processes and the measurement tunnel +were stopped after capture. diff --git a/packages/shared-python/shared/core/config/storage.py b/packages/shared-python/shared/core/config/storage.py index 930ae06cd..a51193377 100644 --- a/packages/shared-python/shared/core/config/storage.py +++ b/packages/shared-python/shared/core/config/storage.py @@ -111,6 +111,20 @@ class StorageConfig(BaseModel): "ZIP bundles remain single-object uploads." ), ) + DEMO_CANONICAL_BUNDLE_ENABLED: bool = Field( + default=False, + description=( + "Reuse immutable demo ZIP and raw artifacts by content version. " + "Disabling only changes new uploads; existing canonical results remain readable." + ), + ) + DEMO_PUBLICATION_PREPARATION_CACHE_ENABLED: bool = Field( + default=False, + description=( + "Reuse pure tokenization and token frequencies for versioned canonical " + "demo publication, in a bounded process-local cache. Requires canonical bundles." + ), + ) SUPPORTED_EXTENSIONS: str = Field( default=".doc,.docx,.pdf,.txt,.xls,.xlsx,.pptx,.jpg,.jpeg,.png,.md,.html,.htm", description="Supported file extensions", diff --git a/packages/shared-python/shared/services/retrieval/agent_tools/tools/grep.py b/packages/shared-python/shared/services/retrieval/agent_tools/tools/grep.py index e86f2e04f..87c40fe49 100644 --- a/packages/shared-python/shared/services/retrieval/agent_tools/tools/grep.py +++ b/packages/shared-python/shared/services/retrieval/agent_tools/tools/grep.py @@ -223,6 +223,7 @@ async def grep(ctx: ToolContext, args: dict[str, Any]) -> ToolResult: matched.c.chunk_metadata, matched.c.job_result_id, JobResult.job_id, + JobResult.document_metadata["result_raw_prefix"].as_string(), matched.c.section_path, Document.source_file_name, func.count().over().label("total_matches"), @@ -253,6 +254,7 @@ async def grep(ctx: ToolContext, args: dict[str, Any]) -> ToolResult: chunk_metadata, job_result_id, job_id, + result_raw_prefix, section_path, source_file_name, _total_matches, @@ -275,6 +277,7 @@ async def grep(ctx: ToolContext, args: dict[str, Any]) -> ToolResult: "chunk_metadata": chunk_metadata or {}, "job_result_id": job_result_id, "job_id": job_id, + "result_raw_prefix": result_raw_prefix, } ) diff --git a/packages/shared-python/shared/services/retrieval/agent_tools/tools/query_table.py b/packages/shared-python/shared/services/retrieval/agent_tools/tools/query_table.py index 05927c704..69ce7f63c 100644 --- a/packages/shared-python/shared/services/retrieval/agent_tools/tools/query_table.py +++ b/packages/shared-python/shared/services/retrieval/agent_tools/tools/query_table.py @@ -68,7 +68,10 @@ async def query_table(ctx: ToolContext, args: dict[str, Any]) -> ToolResult: document = ( await ctx.db.execute( - select(Document, JobResult.job_id) + select( + Document, JobResult.job_id, + JobResult.document_metadata["result_raw_prefix"].as_string(), + ) .select_from(Document) .outerjoin(JobResult, JobResult.id == Document.current_job_result_id) .where(Document.document_id == document_id) @@ -80,7 +83,7 @@ async def query_table(ctx: ToolContext, args: dict[str, Any]) -> ToolResult: ).first() if document is None: return ToolResult(text="", error=f"unknown document_id: {document_id}") - doc, job_id = document + doc, job_id, result_raw_prefix = document if not doc.current_job_result_id: return ToolResult(text="", error=f"unknown document_id: {document_id}") @@ -107,6 +110,7 @@ async def query_table(ctx: ToolContext, args: dict[str, Any]) -> ToolResult: "content": chunk.content, "file_path": chunk.file_path, "job_id": job_id, + "result_raw_prefix": result_raw_prefix, } ) except TableDownloadError as exc: diff --git a/packages/shared-python/shared/services/retrieval/agent_tools/tools/read.py b/packages/shared-python/shared/services/retrieval/agent_tools/tools/read.py index 52f6db013..bd3b6632c 100644 --- a/packages/shared-python/shared/services/retrieval/agent_tools/tools/read.py +++ b/packages/shared-python/shared/services/retrieval/agent_tools/tools/read.py @@ -350,11 +350,18 @@ async def read(ctx: ToolContext, args: dict[str, Any]) -> ToolResult: source_file_name_by_doc = {d.document_id: d.source_file_name or "" for d in documents} job_result_ids = sorted(set(revision_by_doc.values())) job_id_by_revision: dict[str, str] = {} + raw_prefix_by_revision: dict[str, str | None] = {} if job_result_ids: job_rows = await ctx.db.execute( - select(JobResult.id, JobResult.job_id).where(JobResult.id.in_(job_result_ids)) + select( + JobResult.id, JobResult.job_id, + JobResult.document_metadata["result_raw_prefix"].as_string(), + ).where(JobResult.id.in_(job_result_ids)) ) - job_id_by_revision = {str(rid): str(jid) for rid, jid in job_rows.all() if rid and jid} + for revision_id, job_id, raw_prefix in job_rows.all(): + if revision_id and job_id: + job_id_by_revision[str(revision_id)] = str(job_id) + raw_prefix_by_revision[str(revision_id)] = raw_prefix # One entry per input ref (by index), in call order — the per-ref # ok/failed signal this tool now surfaces instead of a single joined @@ -414,6 +421,7 @@ async def read(ctx: ToolContext, args: dict[str, Any]) -> ToolResult: "document_id": document_id, "job_result_id": job_result_id, "job_id": job_id, + "result_raw_prefix": raw_prefix_by_revision.get(job_result_id), "source_file_name": source_file_name, "chunk_id": chunk.chunk_id, "section_id": chunk.section_id, @@ -453,6 +461,7 @@ async def read(ctx: ToolContext, args: dict[str, Any]) -> ToolResult: "resolved_path": resolved_path, "source_file_name": source_file_name, "job_id": job_id, + "result_raw_prefix": raw_prefix_by_revision.get(job_result_id), }, ) ) @@ -546,6 +555,7 @@ async def read(ctx: ToolContext, args: dict[str, Any]) -> ToolResult: "document_id": job["document_id"], "job_result_id": job["job_result_id"], "job_id": job["job_id"], + "result_raw_prefix": job.get("result_raw_prefix"), "source_file_name": job["source_file_name"], "chunk_id": chunk.chunk_id, "section_id": chunk.section_id, diff --git a/packages/shared-python/shared/services/retrieval/graph/query_service.py b/packages/shared-python/shared/services/retrieval/graph/query_service.py index 819d99469..f5d542b98 100644 --- a/packages/shared-python/shared/services/retrieval/graph/query_service.py +++ b/packages/shared-python/shared/services/retrieval/graph/query_service.py @@ -199,6 +199,10 @@ async def collect_candidate_chunks( 'chunk_metadata': chunk.chunk_metadata or {}, 'job_result_id': chunk.job_result_id, 'job_id': job_result.job_id if job_result else None, + 'result_raw_prefix': ( + (job_result.document_metadata or {}).get('result_raw_prefix') + if job_result else None + ), }) if len(rows) >= top_k: break diff --git a/packages/shared-python/shared/services/retrieval/hydration/assets.py b/packages/shared-python/shared/services/retrieval/hydration/assets.py index 8b679e5ea..bec7ea168 100644 --- a/packages/shared-python/shared/services/retrieval/hydration/assets.py +++ b/packages/shared-python/shared/services/retrieval/hydration/assets.py @@ -7,10 +7,11 @@ from shared.services.retrieval.hydration.row_utils import MEDIA_CHUNK_TYPES, normalize_chunk_type from shared.services.storage.page_pdf_crop import crop_source_pdf_pages +from shared.services.storage.raw_prefix_arguments import RawPrefixArguments from shared.services.storage.result_storage import get_result_storage AssetUrlValue = str -PagePdfRequestKey = tuple[str, tuple[int, ...]] +PagePdfRequestKey = tuple[str, tuple[int, ...], str] def _normalize_artifact_ref(asset_ref: object) -> str | None: @@ -107,7 +108,7 @@ def _resolve_page_pdf_request(row: dict[str, Any]) -> PagePdfRequestKey | None: normalized_pages = tuple(sorted({page for page in pages if page > 0})) if not normalized_pages: return None - return job_id, normalized_pages + return job_id, normalized_pages, str(row.get("result_raw_prefix") or "") async def _generate_retrieval_asset_url( @@ -122,6 +123,11 @@ async def _generate_retrieval_asset_url( job_id, artifact_ref = request try: return get_result_storage().generate_artifact_url( + **( + RawPrefixArguments(raw_prefix=str(row["result_raw_prefix"])) + if row.get("result_raw_prefix") + else RawPrefixArguments() + ), job_id=job_id, artifact_ref=artifact_ref, ) @@ -146,6 +152,11 @@ async def _generate_page_citation_asset_url( job_id, artifact_ref = request try: return get_result_storage().generate_artifact_url( + **( + RawPrefixArguments(raw_prefix=str(row["result_raw_prefix"])) + if row.get("result_raw_prefix") + else RawPrefixArguments() + ), job_id=job_id, artifact_ref=artifact_ref, ) @@ -176,6 +187,11 @@ async def _enrich_page_assets( if artifact_ref is not None: try: asset_url = get_result_storage().generate_artifact_url( + **( + RawPrefixArguments(raw_prefix=str(row["result_raw_prefix"])) + if row.get("result_raw_prefix") + else RawPrefixArguments() + ), job_id=job_id, artifact_ref=artifact_ref, ) @@ -208,10 +224,15 @@ async def _generate_page_pdf_asset_url_for_request( *, log_context: str, ) -> str | None: - job_id, pages = request + job_id, pages, raw_prefix = request try: return await asyncio.to_thread( crop_source_pdf_pages, + **( + RawPrefixArguments(raw_prefix=raw_prefix) + if raw_prefix + else RawPrefixArguments() + ), job_id=job_id, pages=list(pages), ) diff --git a/packages/shared-python/shared/services/retrieval/hydration/connected.py b/packages/shared-python/shared/services/retrieval/hydration/connected.py index b960946a8..93dc41305 100644 --- a/packages/shared-python/shared/services/retrieval/hydration/connected.py +++ b/packages/shared-python/shared/services/retrieval/hydration/connected.py @@ -67,11 +67,14 @@ async def hydrate_connected_target_rows( return [] stmt = ( - # Select only the job identifier needed for the public projection. + # Select only job and artifact location fields needed for hydration. # 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) + select( + Document, DocumentChunk, DocumentSection, JobResult.job_id, + JobResult.document_metadata["result_raw_prefix"].as_string(), + ) .join( DocumentChunk, ( @@ -100,7 +103,7 @@ async def hydrate_connected_target_rows( result = await db.execute(stmt) hydrated_rows: list[dict[str, Any]] = [] - for document, chunk, section, job_id in result.all(): + for document, chunk, section, job_id, result_raw_prefix in result.all(): section_path = section.section_path if section else None hydrated_rows.append( { @@ -116,6 +119,7 @@ async def hydrate_connected_target_rows( 'chunk_metadata': chunk.chunk_metadata or {}, 'job_result_id': chunk.job_result_id, 'job_id': job_id, + 'result_raw_prefix': result_raw_prefix, 'sort_order': chunk.sort_order, } ) diff --git a/packages/shared-python/shared/services/retrieval/hydration/evidence_compose.py b/packages/shared-python/shared/services/retrieval/hydration/evidence_compose.py index c5cb97463..0c1e4430d 100644 --- a/packages/shared-python/shared/services/retrieval/hydration/evidence_compose.py +++ b/packages/shared-python/shared/services/retrieval/hydration/evidence_compose.py @@ -26,6 +26,7 @@ TableDownloadError, load_table_html, ) +from shared.services.storage.raw_prefix_arguments import RawPrefixArguments from shared.services.storage.result_storage import get_result_storage _IMAGE_MEDIA_TYPES = { @@ -284,6 +285,11 @@ def _try_read_image_artifact( return None try: temp_path = storage.download_raw_to_temp( + **( + RawPrefixArguments(raw_prefix=str(row["result_raw_prefix"])) + if row.get("result_raw_prefix") + else RawPrefixArguments() + ), job_id=job_id, relative_path=normalized, suffix=Path(normalized).suffix or ".bin", diff --git a/packages/shared-python/shared/services/retrieval/hydration/path.py b/packages/shared-python/shared/services/retrieval/hydration/path.py index b33d193bd..5572f864b 100644 --- a/packages/shared-python/shared/services/retrieval/hydration/path.py +++ b/packages/shared-python/shared/services/retrieval/hydration/path.py @@ -228,6 +228,10 @@ async def _hydrate_chunk_paths( 'chunk_metadata': chunk.chunk_metadata or {}, 'job_result_id': chunk.job_result_id, 'job_id': job_result.job_id if job_result else None, + 'result_raw_prefix': ( + (job_result.document_metadata or {}).get('result_raw_prefix') + if job_result else None + ), 'source_chunk_path': chunk.source_chunk_path, 'sort_order': chunk.sort_order, 'hydrate_mode': path_mode, diff --git a/packages/shared-python/shared/services/retrieval/hydration/reference.py b/packages/shared-python/shared/services/retrieval/hydration/reference.py index 753fc66ff..5a8999e40 100644 --- a/packages/shared-python/shared/services/retrieval/hydration/reference.py +++ b/packages/shared-python/shared/services/retrieval/hydration/reference.py @@ -68,11 +68,14 @@ async def hydrate_referenced_chunk_rows( ) stmt = ( - # Select only the job identifier needed for the public projection. + # Select only job and artifact location fields needed for hydration. # Selecting the JobResult entity would trigger its ``chunks`` selectin # relationship, loading the entire legacy job-chunk collection for # every referenced revision during final hydration. - select(Document, DocumentChunk, DocumentSection, JobResult.job_id) + select( + Document, DocumentChunk, DocumentSection, JobResult.job_id, + JobResult.document_metadata["result_raw_prefix"].as_string(), + ) .join(DocumentChunk, chunk_join) .outerjoin( DocumentSection, DocumentSection.section_id == DocumentChunk.section_id @@ -94,7 +97,7 @@ async def hydrate_referenced_chunk_rows( rows_by_key: dict[ReferenceLookupKey, dict[str, Any]] = {} rows_by_base_key: dict[tuple[str, str], list[dict[str, Any]]] = {} - for document, chunk, section, job_id in result.all(): + for document, chunk, section, job_id, result_raw_prefix in result.all(): row = { "document_id": document.document_id, "chunk_id": chunk.chunk_id, @@ -115,6 +118,7 @@ async def hydrate_referenced_chunk_rows( "chunk_metadata": chunk.chunk_metadata or {}, "job_result_id": chunk.job_result_id, "job_id": job_id, + "result_raw_prefix": result_raw_prefix, "source_chunk_path": chunk.source_chunk_path, "sort_order": chunk.sort_order, } diff --git a/packages/shared-python/shared/services/retrieval/hydration/table_grid.py b/packages/shared-python/shared/services/retrieval/hydration/table_grid.py index 28c6159b4..668d041ea 100644 --- a/packages/shared-python/shared/services/retrieval/hydration/table_grid.py +++ b/packages/shared-python/shared/services/retrieval/hydration/table_grid.py @@ -21,6 +21,7 @@ LARGE_TABLE_AXIS, QUERY_TABLE_NAME, ) +from shared.services.storage.raw_prefix_arguments import RawPrefixArguments from shared.services.storage.result_storage import get_result_storage _TABLE_HTML_RE = re.compile(r")", re.IGNORECASE) @@ -58,6 +59,11 @@ def load_table_html(row: Mapping[str, Any]) -> str: return "" try: temp_path = storage.download_raw_to_temp( + **( + RawPrefixArguments(raw_prefix=str(row["result_raw_prefix"])) + if row.get("result_raw_prefix") + else RawPrefixArguments() + ), job_id=job_id, relative_path=normalized, suffix=".html", diff --git a/packages/shared-python/shared/services/retrieval/maintenance/legacy_token_lookup.py b/packages/shared-python/shared/services/retrieval/maintenance/legacy_token_lookup.py new file mode 100644 index 000000000..1857dd84a --- /dev/null +++ b/packages/shared-python/shared/services/retrieval/maintenance/legacy_token_lookup.py @@ -0,0 +1,101 @@ +"""Guarded retirement shared by the migration and the operator command.""" + +from __future__ import annotations + +from sqlalchemy import text +from sqlalchemy.engine import Connection + +_LEGACY_INDEX: str = "idx_document_map_unit_tokens_lookup" +_DEFINITIONS: dict[str, str] = { + _LEGACY_INDEX: "(channel, token_hash, map_unit_id)", + "idx_document_map_unit_tokens_token_lookup_binary": ( + "(channel, decode((token_hash)::text, 'hex'::text)) " + "INCLUDE (map_unit_id, token, frequency)" + ), + "idx_document_map_unit_tokens_unit_lookup": ( + "(map_unit_id, channel, token_hash) INCLUDE (token, frequency)" + ), +} + + +class LegacyTokenLookupMaintenance: + """Change only the legacy nonunique index after catalog validation. + + Concurrent operations require an AUTOCOMMIT connection. Transaction-owned + migration contracts may explicitly use regular DDL. Definitions are checked + conservatively so a same-name invalid, partial, or incompatible index cannot + silently satisfy a prerequisite. Keep these historical definitions stable. + """ + + def __init__(self, connection: Connection, *, is_concurrent: bool = True) -> None: + self._connection: Connection = connection + self._is_concurrent: bool = is_concurrent + + def inspect_indexes(self) -> dict[str, dict[str, str | bool]]: + rows = self._connection.execute( + text( + """ + SELECT relation.relname AS name, index.indisvalid AS is_valid, + index.indisready AS is_ready, + pg_get_indexdef(index.indexrelid) AS definition + FROM pg_index AS index + JOIN pg_class AS relation ON relation.oid = index.indexrelid + WHERE index.indrelid = to_regclass('public.document_map_unit_tokens') + ORDER BY relation.relname + """ + ) + ).mappings().all() + return { + str(row["name"]): { + "is_valid": bool(row["is_valid"]), + "is_ready": bool(row["is_ready"]), + "definition": str(row["definition"]), + } + for row in rows + } + + def _require_compatible_index( + self, indexes: dict[str, dict[str, str | bool]], name: str + ) -> None: + state: dict[str, str | bool] | None = indexes.get(name) + expected: str = ( + f"CREATE INDEX {name} ON public.document_map_unit_tokens " + f"USING btree {_DEFINITIONS[name]}" + ) + if ( + state is None + or not state["is_valid"] + or not state["is_ready"] + or state["definition"] != expected + ): + raise RuntimeError(f"Required token index is missing or incompatible: {name}") + + def apply_retirement(self) -> None: + indexes: dict[str, dict[str, str | bool]] = self.inspect_indexes() + for name in _DEFINITIONS: + if name != _LEGACY_INDEX: + self._require_compatible_index(indexes, name) + if _LEGACY_INDEX not in indexes: + return + self._require_compatible_index(indexes, _LEGACY_INDEX) + concurrent_clause: str = "CONCURRENTLY " if self._is_concurrent else "" + self._connection.execute( + text(f"DROP INDEX {concurrent_clause}IF EXISTS public.{_LEGACY_INDEX}") + ) + if _LEGACY_INDEX in self.inspect_indexes(): + raise RuntimeError("Legacy token lookup index remained after retirement") + + def restore_index(self) -> None: + indexes: dict[str, dict[str, str | bool]] = self.inspect_indexes() + if _LEGACY_INDEX in indexes: + self._require_compatible_index(indexes, _LEGACY_INDEX) + return + concurrent_clause: str = "CONCURRENTLY " if self._is_concurrent else "" + self._connection.execute( + text( + f"CREATE INDEX {concurrent_clause}{_LEGACY_INDEX} " + "ON public.document_map_unit_tokens USING btree " + f"{_DEFINITIONS[_LEGACY_INDEX]}" + ) + ) + self._require_compatible_index(self.inspect_indexes(), _LEGACY_INDEX) diff --git a/packages/shared-python/shared/services/retrieval/map_unit_index.py b/packages/shared-python/shared/services/retrieval/map_unit_index.py index 141a7a45b..6674c601b 100644 --- a/packages/shared-python/shared/services/retrieval/map_unit_index.py +++ b/packages/shared-python/shared/services/retrieval/map_unit_index.py @@ -31,6 +31,7 @@ ) from shared.services.retrieval.scoring.score_units import build_score_units from shared.services.retrieval.publication_models import DocumentPublicationScope +from shared.services.retrieval.publication_preparation_cache import PublicationPreparationCache from shared.services.retrieval.publication_strategy import resolve_publication_strategy from shared.services.retrieval.publication_token_copy import ( TokenCopyRow, @@ -130,16 +131,20 @@ def replace_document_map_units( if not unit_id or not section_id: continue map_unit_id = f"dmu_{uuid4().hex}" - path_tokens = str(unit.get("path_search_text") or "").split() - content_tokens = str(unit.get("content_search_text") or "").split() - if path_tokens: + path_frequencies = PublicationPreparationCache.count_tokens( + str(unit.get("path_search_text") or "") + ) + content_frequencies = PublicationPreparationCache.count_tokens( + str(unit.get("content_search_text") or "") + ) + path_token_count = sum(path_frequencies.values()) + content_token_count = sum(content_frequencies.values()) + if path_token_count: path_document_count += 1 - path_total_length += len(path_tokens) - if content_tokens: + path_total_length += path_token_count + if content_token_count: content_document_count += 1 - content_total_length += len(content_tokens) - path_frequencies = Counter(path_tokens) - content_frequencies = Counter(content_tokens) + content_total_length += content_token_count path_unit_df.update(path_frequencies.keys()) content_unit_df.update(content_frequencies.keys()) # ``provider.self_units`` already reflects root-asset remount (assets @@ -160,8 +165,8 @@ def replace_document_map_units( "unit_id": unit_id, "section_id": section_id, "unit_kind": str(unit.get("kind") or "leaf"), - "path_token_count": len(path_tokens), - "content_token_count": len(content_tokens), + "path_token_count": path_token_count, + "content_token_count": content_token_count, "has_image": "image" in section_types, "has_table": "table" in section_types, "sort_order": sort_order, diff --git a/packages/shared-python/shared/services/retrieval/publication_preparation_cache.py b/packages/shared-python/shared/services/retrieval/publication_preparation_cache.py new file mode 100644 index 000000000..529c7ef63 --- /dev/null +++ b/packages/shared-python/shared/services/retrieval/publication_preparation_cache.py @@ -0,0 +1,143 @@ +"""Bounded reuse of pure lexical preparation during canonical demo publication.""" + +from __future__ import annotations + +import sys +import threading +from collections import Counter, OrderedDict +from collections.abc import Iterator +from contextlib import contextmanager +from contextvars import ContextVar, Token +from dataclasses import dataclass +from hashlib import sha256 + +from shared.utils.text_utils import tokenize_for_retrieval + +_PREPARATION_FORMAT_VERSION = 1 +_MAX_BYTES = 64 * 1024 * 1024 +_MAX_ENTRIES = 4096 +_Scope = tuple[str, str, int, int] +_Key = tuple[_Scope, str, bool, bytes] +_active_scope: ContextVar[_Scope | None] = ContextVar( + "demo_publication_preparation_scope", default=None +) + + +@dataclass(frozen=True) +class _PreparedValue: + tokens: tuple[str, ...] = () + frequencies: tuple[tuple[str, int], ...] = () + size_bytes: int = 0 + + +class PublicationPreparationCache: + """Reuse token sequences and frequencies without storing owner-specific rows. + + Only the canonical materializer opens a scope. Ordinary ingestion, backfills + and retrieval call the same authoritative tokenizer without cache access. + The exact text digest and tokenization mode are part of every key; source + and algorithm versions isolate changed artifacts and preparation semantics. + """ + + _entries: OrderedDict[_Key, _PreparedValue] = OrderedDict() + _size_bytes: int = 0 + _lock: threading.RLock = threading.RLock() + + @classmethod + @contextmanager + def scope( + cls, *, source_id: str, content_version: str, index_format_version: int + ) -> Iterator[None]: + if not source_id or not content_version: + raise ValueError("Demo preparation reuse requires source and content versions") + scope: _Scope = ( + source_id, content_version, index_format_version, _PREPARATION_FORMAT_VERSION + ) + token: Token[_Scope | None] = _active_scope.set(scope) + try: + yield + finally: + _active_scope.reset(token) + + @classmethod + def tokenize(cls, text: str, *, dedupe: bool = False) -> list[str]: + key: _Key | None = cls._build_key(text, operation="tokens", dedupe=dedupe) + cached: _PreparedValue | None = cls._get(key) + if cached is not None: + return list(cached.tokens) + tokens: list[str] = tokenize_for_retrieval( + text, stopwords=[], dedupe=dedupe, min_token_length=2 + ) + if key is not None: + immutable_tokens: tuple[str, ...] = tuple(tokens) + size_bytes: int = sys.getsizeof(immutable_tokens) + sum( + sys.getsizeof(item) for item in immutable_tokens + ) + cls._store(key, _PreparedValue(tokens=immutable_tokens, size_bytes=size_bytes)) + return tokens + + @classmethod + def count_tokens(cls, text: str) -> Counter[str]: + key: _Key | None = cls._build_key(text, operation="frequencies", dedupe=False) + cached: _PreparedValue | None = cls._get(key) + if cached is not None: + return Counter(dict(cached.frequencies)) + frequencies: Counter[str] = Counter(text.split()) + if key is not None: + immutable_frequencies: tuple[tuple[str, int], ...] = tuple(frequencies.items()) + size_bytes: int = sys.getsizeof(immutable_frequencies) + sum( + sys.getsizeof(item) + sys.getsizeof(item[0]) + sys.getsizeof(item[1]) + for item in immutable_frequencies + ) + cls._store( + key, + _PreparedValue(frequencies=immutable_frequencies, size_bytes=size_bytes), + ) + return frequencies + + @classmethod + def clear(cls) -> None: + with cls._lock: + cls._entries.clear() + cls._size_bytes = 0 + + @staticmethod + def _build_key(text: str, *, operation: str, dedupe: bool) -> _Key | None: + scope: _Scope | None = _active_scope.get() + if scope is None: + return None + return scope, operation, dedupe, sha256(text.encode("utf-8")).digest() + + @classmethod + def _get(cls, key: _Key | None) -> _PreparedValue | None: + if key is None: + return None + with cls._lock: + value: _PreparedValue | None = cls._entries.get(key) + if value is not None: + cls._entries.move_to_end(key) + return value + + @classmethod + def _store(cls, key: _Key, value: _PreparedValue) -> None: + # Include conservative per-entry overhead for keys, scope and LRU nodes. + size_bytes: int = value.size_bytes + 1024 + sum( + sys.getsizeof(part) for part in key[0] + ) + if size_bytes > _MAX_BYTES: + return + value = _PreparedValue( + tokens=value.tokens, frequencies=value.frequencies, size_bytes=size_bytes + ) + with cls._lock: + previous: _PreparedValue | None = cls._entries.pop(key, None) + if previous is not None: + cls._size_bytes -= previous.size_bytes + while cls._entries and ( + cls._size_bytes + size_bytes > _MAX_BYTES + or len(cls._entries) >= _MAX_ENTRIES + ): + _, evicted = cls._entries.popitem(last=False) + cls._size_bytes -= evicted.size_bytes + cls._entries[key] = value + cls._size_bytes += size_bytes diff --git a/packages/shared-python/shared/services/retrieval/scoring/knowhere_hybrid.py b/packages/shared-python/shared/services/retrieval/scoring/knowhere_hybrid.py index ca73845ba..7e864e831 100644 --- a/packages/shared-python/shared/services/retrieval/scoring/knowhere_hybrid.py +++ b/packages/shared-python/shared/services/retrieval/scoring/knowhere_hybrid.py @@ -13,7 +13,7 @@ from dataclasses import dataclass from typing import Dict, List, Mapping, Optional, Sequence, Tuple -from shared.utils.text_utils import tokenize_for_retrieval as _tokenize_word_level +from shared.services.retrieval.publication_preparation_cache import PublicationPreparationCache RRF_K = 60 CHANNEL_WEIGHT_PATH = 1.0 @@ -30,12 +30,7 @@ def tokenize_for_retrieval(text: str, *, dedupe: bool = True) -> List[str]: ``min_token_length=2`` so map-unit indexes match ``document_chunks`` search text. Character-level regex tokenization was removed. """ - return _tokenize_word_level( - text, - stopwords=[], - dedupe=dedupe, - min_token_length=2, - ) + return PublicationPreparationCache.tokenize(text, dedupe=dedupe) def tokenize_query_for_ranker(query: str) -> List[str]: diff --git a/packages/shared-python/shared/services/retrieval/search/lexical_text.py b/packages/shared-python/shared/services/retrieval/search/lexical_text.py index b6647a16d..176dbf746 100644 --- a/packages/shared-python/shared/services/retrieval/search/lexical_text.py +++ b/packages/shared-python/shared/services/retrieval/search/lexical_text.py @@ -9,7 +9,7 @@ from shared.services.chunks.document_path import split_document_path from shared.services.chunks.path_segments import split_escaped_document_path -from shared.utils.text_utils import tokenize_contents_for_retrieval +from shared.services.retrieval.publication_preparation_cache import PublicationPreparationCache _SAME_AS_RE = re.compile(r"\[SAME-AS [^\]]+\]") @@ -44,8 +44,7 @@ def build_lexical_text(value: str) -> str: if not text: return "" - tokens = tokenize_contents_for_retrieval([text], stopwords=[], link_char=" ", dedupe=True) - token_text = tokens[0] if tokens else "" + token_text = " ".join(PublicationPreparationCache.tokenize(text, dedupe=True)) lexical_parts = [part for part in [text, token_text] if part] return "\n".join(lexical_parts) if lexical_parts else text @@ -136,8 +135,7 @@ def build_content_search_text( if section_summary and str(section_summary).strip(): parts.append(str(section_summary).strip()) raw = " ".join(parts) - tokens = tokenize_contents_for_retrieval([raw], stopwords=[], link_char=" ") - return tokens[0] if tokens else raw + return " ".join(PublicationPreparationCache.tokenize(raw)) def _chunk_summary_text(chunk: dict[str, Any]) -> str: @@ -179,8 +177,7 @@ def build_path_search_text( if not parts: return None raw = " ".join(parts) - tokens = tokenize_contents_for_retrieval([raw], stopwords=[], link_char=" ") - return tokens[0] if tokens else raw + return " ".join(PublicationPreparationCache.tokenize(raw)) def build_term_search_text( 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 aa5869ff2..d0ad50a2d 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 @@ -802,7 +802,8 @@ async def _hydrate_winning_units( "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 " + "jr.job_id, " + "jr.document_metadata ->> 'result_raw_prefix' AS result_raw_prefix " "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 " diff --git a/packages/shared-python/shared/services/retrieval/search/scoped_corpus.py b/packages/shared-python/shared/services/retrieval/search/scoped_corpus.py index 47f27d5a8..69c581fb0 100644 --- a/packages/shared-python/shared/services/retrieval/search/scoped_corpus.py +++ b/packages/shared-python/shared/services/retrieval/search/scoped_corpus.py @@ -190,6 +190,10 @@ async def load_all_scoped_chunks( 'chunk_metadata': chunk.chunk_metadata or {}, 'job_result_id': chunk.job_result_id, 'job_id': job_result.job_id if job_result else None, + 'result_raw_prefix': ( + (job_result.document_metadata or {}).get('result_raw_prefix') + if job_result else None + ), 'sort_order': chunk.sort_order, }) return rows diff --git a/packages/shared-python/shared/services/storage/page_pdf_crop.py b/packages/shared-python/shared/services/storage/page_pdf_crop.py index c3de8f514..9024f574e 100644 --- a/packages/shared-python/shared/services/storage/page_pdf_crop.py +++ b/packages/shared-python/shared/services/storage/page_pdf_crop.py @@ -16,6 +16,7 @@ def crop_source_pdf_pages( *, job_id: str, + raw_prefix: str | None = None, pages: list[int], storage: JobResultStorage | None = None, expires_in: int = 3600, @@ -27,22 +28,36 @@ def crop_source_pdf_pages( result_storage = storage or JobResultStorage() cache_ref = _cache_ref(normalized_pages) - if result_storage.verify_raw_exists(job_id=job_id, relative_path=cache_ref): - return result_storage.generate_artifact_url( + if _verify_raw_exists( + result_storage, + job_id=job_id, + raw_prefix=raw_prefix, + relative_path=cache_ref, + ): + return _generate_artifact_url( + result_storage, job_id=job_id, + raw_prefix=raw_prefix, artifact_ref=cache_ref, expires_in=expires_in, ) - if not result_storage.verify_raw_exists(job_id=job_id, relative_path=_SOURCE_PDF_REF): + if not _verify_raw_exists( + result_storage, + job_id=job_id, + raw_prefix=raw_prefix, + relative_path=_SOURCE_PDF_REF, + ): logger.warning("[page_pdf_crop] source.pdf missing for job_id={}", job_id) return None work_dir = temp_dir or tempfile.gettempdir() try: with tempfile.TemporaryDirectory(dir=work_dir) as local_dir: - source_path = result_storage.download_raw_to_temp( + source_path = _download_raw_to_temp( + result_storage, job_id=job_id, + raw_prefix=raw_prefix, relative_path=_SOURCE_PDF_REF, suffix=".pdf", temp_dir=local_dir, @@ -53,8 +68,10 @@ def crop_source_pdf_pages( output_path=cropped_path, pages=normalized_pages, ) - result_storage.upload_raw_file( + _upload_raw_file( + result_storage, job_id=job_id, + raw_prefix=raw_prefix, relative_path=cache_ref, local_file_path=cropped_path, ) @@ -67,13 +84,101 @@ def crop_source_pdf_pages( ) return None - return result_storage.generate_artifact_url( + return _generate_artifact_url( + result_storage, job_id=job_id, + raw_prefix=raw_prefix, artifact_ref=cache_ref, expires_in=expires_in, ) +def _verify_raw_exists( + storage: JobResultStorage, + *, + job_id: str, + raw_prefix: str | None, + relative_path: str, +) -> bool: + if raw_prefix is None: + return storage.verify_raw_exists(job_id=job_id, relative_path=relative_path) + return storage.verify_raw_exists( + job_id=job_id, + raw_prefix=raw_prefix, + relative_path=relative_path, + ) + + +def _generate_artifact_url( + storage: JobResultStorage, + *, + job_id: str, + raw_prefix: str | None, + artifact_ref: str, + expires_in: int, +) -> str | None: + if raw_prefix is None: + return storage.generate_artifact_url( + job_id=job_id, + artifact_ref=artifact_ref, + expires_in=expires_in, + ) + return storage.generate_artifact_url( + job_id=job_id, + raw_prefix=raw_prefix, + artifact_ref=artifact_ref, + expires_in=expires_in, + ) + + +def _download_raw_to_temp( + storage: JobResultStorage, + *, + job_id: str, + raw_prefix: str | None, + relative_path: str, + suffix: str, + temp_dir: str, +) -> str: + if raw_prefix is None: + return storage.download_raw_to_temp( + job_id=job_id, + relative_path=relative_path, + suffix=suffix, + temp_dir=temp_dir, + ) + return storage.download_raw_to_temp( + job_id=job_id, + raw_prefix=raw_prefix, + relative_path=relative_path, + suffix=suffix, + temp_dir=temp_dir, + ) + + +def _upload_raw_file( + storage: JobResultStorage, + *, + job_id: str, + raw_prefix: str | None, + relative_path: str, + local_file_path: str, +) -> None: + if raw_prefix is None: + storage.upload_raw_file( + job_id=job_id, + relative_path=relative_path, + local_file_path=local_file_path, + ) + return + storage.upload_raw_file( + job_id=job_id, + raw_prefix=raw_prefix, + relative_path=relative_path, + local_file_path=local_file_path, + ) + + def _normalize_pages(pages: list[int]) -> list[int]: normalized: set[int] = set() for page in pages: diff --git a/packages/shared-python/shared/services/storage/raw_prefix_arguments.py b/packages/shared-python/shared/services/storage/raw_prefix_arguments.py new file mode 100644 index 000000000..7766056c4 --- /dev/null +++ b/packages/shared-python/shared/services/storage/raw_prefix_arguments.py @@ -0,0 +1,7 @@ +"""Optional artifact location arguments for legacy and canonical result readers.""" + +from typing import TypedDict + + +class RawPrefixArguments(TypedDict, total=False): + raw_prefix: str diff --git a/packages/shared-python/shared/services/storage/result_storage.py b/packages/shared-python/shared/services/storage/result_storage.py index c68febab4..8583e19ee 100644 --- a/packages/shared-python/shared/services/storage/result_storage.py +++ b/packages/shared-python/shared/services/storage/result_storage.py @@ -42,16 +42,38 @@ def upload( raise NotImplementedError def generate_artifact_url( - self, *, job_id: str, artifact_ref: str, expires_in: int = 3600 + self, + *, + job_id: str, + artifact_ref: str, + raw_prefix: str | None = None, + expires_in: int = 3600, ) -> str | None: raise NotImplementedError def generate_raw_file_url( - self, *, job_id: str, relative_path: str, expires_in: int = 3600 + self, + *, + job_id: str, + relative_path: str, + raw_prefix: str | None = None, + expires_in: int = 3600, ) -> str | None: raise NotImplementedError - def verify_raw_exists(self, *, job_id: str, relative_path: str) -> bool: + def verify_raw_exists( + self, *, job_id: str, relative_path: str, raw_prefix: str | None = None + ) -> bool: + raise NotImplementedError + + def upload_raw_file( + self, + *, + job_id: str, + relative_path: str, + local_file_path: str, + raw_prefix: str | None = None, + ) -> None: raise NotImplementedError def normalize_artifact_ref(self, artifact_ref: str | None) -> str | None: @@ -62,6 +84,7 @@ def download_raw_to_temp( *, job_id: str, relative_path: str, + raw_prefix: str | None = None, suffix: str, temp_dir: str, ) -> str: @@ -92,11 +115,20 @@ def build_zip_key(self, *, job_id: str) -> str: def build_raw_prefix(self, *, job_id: str) -> str: return self._job_file_storage.build_result_raw_prefix(job_id=job_id) - def build_raw_key(self, *, job_id: str, relative_path: str) -> str: + def build_raw_key( + self, + *, + job_id: str, + relative_path: str, + raw_prefix: str | None = None, + ) -> str: normalized = self._normalize_raw_relative_path(relative_path) if not normalized: raise ValueError(f"Invalid result raw artifact path: {relative_path}") - return f"{self.build_raw_prefix(job_id=job_id)}{normalized}" + prefix = self._normalize_raw_prefix(raw_prefix) + if prefix is None: + prefix = self.build_raw_prefix(job_id=job_id) + return f"{prefix}{normalized}" def normalize_artifact_ref(self, artifact_ref: str | None) -> str | None: normalized = self._normalize_raw_relative_path(artifact_ref) @@ -205,16 +237,35 @@ def generate_url(self, *, storage_key: str, expires_in: int = 3600) -> str | Non expires_in=expires_in, )["download_url"] - def verify_raw_exists(self, *, job_id: str, relative_path: str) -> bool: - key = self.build_raw_key(job_id=job_id, relative_path=relative_path) + def verify_raw_exists( + self, + *, + job_id: str, + relative_path: str, + raw_prefix: str | None = None, + ) -> bool: + key = self.build_raw_key( + job_id=job_id, + relative_path=relative_path, + raw_prefix=raw_prefix, + ) result = self._job_file_storage.verify_exists(key, bucket=self.results_bucket) return bool(result.get("exists")) def generate_raw_file_url( - self, *, job_id: str, relative_path: str, expires_in: int = 3600 + self, + *, + job_id: str, + relative_path: str, + raw_prefix: str | None = None, + expires_in: int = 3600, ) -> str | None: return self.generate_url( - storage_key=self.build_raw_key(job_id=job_id, relative_path=relative_path), + storage_key=self.build_raw_key( + job_id=job_id, + relative_path=relative_path, + raw_prefix=raw_prefix, + ), expires_in=expires_in, ) @@ -223,11 +274,16 @@ def upload_raw_file( *, job_id: str, relative_path: str, + raw_prefix: str | None = None, local_file_path: str, ) -> None: self._job_file_storage.upload_local_file( local_file_path, - self.build_raw_key(job_id=job_id, relative_path=relative_path), + self.build_raw_key( + job_id=job_id, + relative_path=relative_path, + raw_prefix=raw_prefix, + ), bucket=self.results_bucket, ) @@ -236,24 +292,38 @@ def download_raw_to_temp( *, job_id: str, relative_path: str, + raw_prefix: str | None = None, suffix: str, temp_dir: str, ) -> str: return self._job_file_storage.download_to_temp( - self.build_raw_key(job_id=job_id, relative_path=relative_path), + self.build_raw_key( + job_id=job_id, + relative_path=relative_path, + raw_prefix=raw_prefix, + ), suffix=suffix, temp_dir=temp_dir, bucket=self.results_bucket, ) def generate_artifact_url( - self, *, job_id: str, artifact_ref: str, expires_in: int = 3600 + self, + *, + job_id: str, + artifact_ref: str, + raw_prefix: str | None = None, + expires_in: int = 3600, ) -> str | None: normalized_ref = self.normalize_artifact_ref(artifact_ref) if not normalized_ref: return None return self.generate_url( - storage_key=self.build_raw_key(job_id=job_id, relative_path=normalized_ref), + storage_key=self.build_raw_key( + job_id=job_id, + relative_path=normalized_ref, + raw_prefix=raw_prefix, + ), expires_in=expires_in, ) @@ -284,6 +354,21 @@ def _normalize_raw_relative_path(self, relative_path: str | None) -> str | None: return None return "/".join(parts) + def _normalize_raw_prefix(self, raw_prefix: str | None) -> str | None: + if raw_prefix is None: + return None + normalized = str(raw_prefix).strip().replace("\\", "/").lstrip("/") + if not normalized: + raise ValueError("Result raw prefix must not be empty") + if not normalized.startswith("results/"): + raise ValueError("Result raw prefix must be under results/") + parts = normalized.rstrip("/").split("/") + if not parts or any(not part or part in {".", ".."} for part in parts): + raise ValueError("Result raw prefix contains an invalid path component") + if any(self._is_excluded_dir(part) for part in parts): + raise ValueError("Result raw prefix contains an excluded directory") + return "/".join(parts) + "/" + def _normalize_artifact_refs(self, artifact_refs: set[str] | None) -> set[str] | None: if artifact_refs is None: return None diff --git a/packages/shared-python/shared/tests/test_agent_tools_discovery_read_contract.py b/packages/shared-python/shared/tests/test_agent_tools_discovery_read_contract.py index d667c7865..a74718c93 100644 --- a/packages/shared-python/shared/tests/test_agent_tools_discovery_read_contract.py +++ b/packages/shared-python/shared/tests/test_agent_tools_discovery_read_contract.py @@ -65,6 +65,7 @@ def all(self) -> list[tuple[object, ...]]: {"summary": "dose table"}, REV_ID, JOB_ID, + None, PATH_ROOT, FILE_NAME, 1, diff --git a/packages/shared-python/shared/tests/test_agent_tools_explore_table_image.py b/packages/shared-python/shared/tests/test_agent_tools_explore_table_image.py index 2c5ca25a6..845af133e 100644 --- a/packages/shared-python/shared/tests/test_agent_tools_explore_table_image.py +++ b/packages/shared-python/shared/tests/test_agent_tools_explore_table_image.py @@ -745,6 +745,7 @@ def all(self) -> list[tuple[object, ...]]: }, REV_ID, JOB_ID, + None, PATH_INTRO, FILE_NAME, 1, diff --git a/packages/shared-python/shared/tests/test_agent_tools_grep.py b/packages/shared-python/shared/tests/test_agent_tools_grep.py index 25e8f23a8..67437f1af 100644 --- a/packages/shared-python/shared/tests/test_agent_tools_grep.py +++ b/packages/shared-python/shared/tests/test_agent_tools_grep.py @@ -46,6 +46,7 @@ def _text_row( None, "jr_a", "job_a", + None, section_path, source_file_name, total, @@ -387,6 +388,7 @@ def all(self) -> list[tuple[object, ...]]: {"summary": "dose table summary 30 mg"}, "jr_a", "job_a", + None, "guide.pdf / Root", "guide.pdf", 1, @@ -556,6 +558,7 @@ def all(self) -> list[tuple[object, ...]]: None, "jr_a", "job_a", + None, "guide.pdf / Root", "guide.pdf", 1, diff --git a/scripts/retire_legacy_token_lookup.py b/scripts/retire_legacy_token_lookup.py new file mode 100644 index 000000000..4d709311b --- /dev/null +++ b/scripts/retire_legacy_token_lookup.py @@ -0,0 +1,74 @@ +"""Inspect, explicitly retire, or restore the legacy token lookup index.""" + +from __future__ import annotations + +import argparse +import json +import os +import sys + +from sqlalchemy import create_engine, text +from sqlalchemy.engine import Engine, URL, make_url +from sqlalchemy.pool import NullPool + +from shared.services.retrieval.maintenance.legacy_token_lookup import ( + LegacyTokenLookupMaintenance, +) + + +def main() -> int: + parser: argparse.ArgumentParser = argparse.ArgumentParser(description=__doc__) + actions: argparse._MutuallyExclusiveGroup = parser.add_mutually_exclusive_group() + actions.add_argument("--apply", action="store_true", help="Retire the legacy index") + actions.add_argument("--restore", action="store_true", help="Rebuild the legacy index") + parser.add_argument("--statement-timeout-seconds", type=int, default=1800) + arguments: argparse.Namespace = parser.parse_args() + database_url: str | None = os.environ.get("DATABASE_URL") + if not database_url: + parser.error("DATABASE_URL must be set in the environment") + if arguments.statement_timeout_seconds <= 0: + parser.error("--statement-timeout-seconds must be positive") + url: URL = make_url(database_url).set(drivername="postgresql+psycopg2") + ssl_variables: dict[str, str] = { + "DB_SSL_MODE": "sslmode", + "DB_SSL_CERT": "sslcert", + "DB_SSL_KEY": "sslkey", + "DB_SSL_ROOT_CERT": "sslrootcert", + } + connection_arguments: dict[str, str] = { + parameter: os.environ[variable] + for variable, parameter in ssl_variables.items() + if os.environ.get(variable) + } + engine: Engine = create_engine( + url, connect_args=connection_arguments, poolclass=NullPool, + isolation_level="AUTOCOMMIT", + ) + try: + with engine.connect() as connection: + connection.execute(text("SET lock_timeout = '5s'")) + connection.execute(text( + f"SET statement_timeout = '{arguments.statement_timeout_seconds}s'" + )) + maintenance: LegacyTokenLookupMaintenance = LegacyTokenLookupMaintenance( + connection + ) + if arguments.apply: + maintenance.apply_retirement() + elif arguments.restore: + maintenance.restore_index() + print(json.dumps(maintenance.inspect_indexes(), indent=2, sort_keys=True)) + except RuntimeError as error: + print(str(error), file=sys.stderr) + return 1 + except Exception as error: + # Connection exceptions may include credentials. Report only their type. + print(f"Index maintenance failed: {type(error).__name__}", file=sys.stderr) + return 1 + finally: + engine.dispose() + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) From 5c75c9cf94ff2ad7b24431d2a1088f833c20ca31 Mon Sep 17 00:00:00 2001 From: suguanYang Date: Wed, 30 Sep 2026 14:39:36 +0800 Subject: [PATCH 2/2] fix: await canonical bundle cancellation work --- .../api/app/services/demo/canonical_bundle.py | 4 +- .../test_canonical_demo_bundle_contract.py | 90 +++++++++++++++++++ 2 files changed, 92 insertions(+), 2 deletions(-) diff --git a/apps/api/app/services/demo/canonical_bundle.py b/apps/api/app/services/demo/canonical_bundle.py index 432c9c0ea..fa6923f22 100644 --- a/apps/api/app/services/demo/canonical_bundle.py +++ b/apps/api/app/services/demo/canonical_bundle.py @@ -104,7 +104,7 @@ async def ensure_bundle( uploaded_bundle: CanonicalDemoBundle = await asyncio.shield(upload) except asyncio.CancelledError: # A thread cannot be cancelled; retain the lease until it stops. - await upload + await asyncio.gather(upload) raise if renewal.done(): renewal.result() @@ -118,7 +118,7 @@ async def ensure_bundle( try: await asyncio.shield(marker_upload) except asyncio.CancelledError: - await marker_upload + await asyncio.gather(marker_upload) raise return uploaded_bundle finally: diff --git a/apps/api/tests/contract/test_canonical_demo_bundle_contract.py b/apps/api/tests/contract/test_canonical_demo_bundle_contract.py index 9a075609d..443d1647b 100644 --- a/apps/api/tests/contract/test_canonical_demo_bundle_contract.py +++ b/apps/api/tests/contract/test_canonical_demo_bundle_contract.py @@ -6,6 +6,8 @@ from collections.abc import Callable from contextlib import AbstractAsyncContextManager from pathlib import Path +from threading import Event +from typing import Literal import pytest from httpx import AsyncClient @@ -19,6 +21,94 @@ def _write_source(directory: Path, content: str) -> None: (directory / "source.md").write_text(content, encoding="utf-8") +@pytest.mark.asyncio +@pytest.mark.parametrize("stage", ("upload", "marker")) +async def test_cancellation_retains_lease_until_storage_work_finishes( + api_client_factory: Callable[[], AbstractAsyncContextManager[AsyncClient]], + tmp_path: Path, + monkeypatch: MonkeyPatch, + stage: Literal["upload", "marker"], +) -> None: + async with api_client_factory(): + from app.services.demo.canonical_bundle import ( + CanonicalDemoBundleStore, + _calculate_content_version, + _read_source_signature, + ) + from app.services.demo.canonical_bundle_result import CanonicalDemoBundle + from shared.services.redis import RedisServiceFactory + + directory: Path = tmp_path / stage + _write_source(directory, "Cancellation must retain the upload lease") + source_id: str = f"contract-cancel-{stage}" + version: str = _calculate_content_version( + directory, _read_source_signature(directory) + ) + lock_key: str = f"lock:demo_bundle:{source_id}:{version}" + redis_service = RedisServiceFactory.get_service() + store = CanonicalDemoBundleStore(redis_service=redis_service) + loop: asyncio.AbstractEventLoop = asyncio.get_running_loop() + started: asyncio.Event = asyncio.Event() + released: Event = Event() + finished: Event = Event() + original_upload = store._upload_bundle + original_marker = store._write_ready_marker + + def pause_storage() -> None: + loop.call_soon_threadsafe(started.set) + if not released.wait(timeout=10): + raise TimeoutError("Contract did not release paused storage work") + + def upload_bundle( + storage_id: str, + content_version: str, + source_directory: Path, + signature: tuple[tuple[str, int, int, int], ...], + ) -> CanonicalDemoBundle: + pause_storage() + try: + return original_upload( + storage_id, content_version, source_directory, signature + ) + finally: + finished.set() + + def write_marker(bundle: CanonicalDemoBundle) -> None: + pause_storage() + try: + original_marker(bundle) + finally: + finished.set() + + if stage == "upload": + monkeypatch.setattr(store, "_upload_bundle", upload_bundle) + else: + monkeypatch.setattr(store, "_write_ready_marker", write_marker) + request: asyncio.Task[CanonicalDemoBundle] = asyncio.create_task( + store.ensure_bundle(source_id=source_id, source_directory=directory) + ) + try: + await asyncio.wait_for(started.wait(), timeout=10) + owner: object = await redis_service.get(lock_key) + assert owner is not None + request.cancel() + await asyncio.sleep(0) + assert not request.done() + assert not finished.is_set() + assert await redis_service.get(lock_key) == owner + finally: + released.set() + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(request, timeout=10) + + assert finished.is_set() + assert not await redis_service.exists(lock_key) + recovered = await CanonicalDemoBundleStore( + redis_service=redis_service + ).ensure_bundle(source_id=source_id, source_directory=directory) + assert recovered.reused is (stage == "marker") + + @pytest.mark.asyncio async def test_canonical_bundle_reuses_and_versions_real_filesystem_objects( api_client_factory: Callable[[], AbstractAsyncContextManager[AsyncClient]],