From ceae239cf752918d3b8562308a94af36911f8691 Mon Sep 17 00:00:00 2001 From: suguanYang Date: Mon, 31 Aug 2026 11:05:14 +0800 Subject: [PATCH 01/10] Revert "Merge pull request #355 from Ontos-AI/perf/wangbinqi/retrieval-hydration-job-id" This reverts commit 5162d0ab0c92cc94c8bbe9b7155da52e40628a99, reversing changes made to e1b8c886118f432fdd7f0a6a4eed76f1293d29e9. --- ...etrieval_snapshot_large_corpus_contract.py | 11 ---- .../services/retrieval/execution/routes.py | 31 +++++++--- .../shared/services/retrieval/nav_snapshot.py | 61 ++++++++----------- .../retrieval/search/scoped_corpus.py | 13 +--- 4 files changed, 49 insertions(+), 67 deletions(-) diff --git a/apps/api/tests/contract/test_retrieval_snapshot_large_corpus_contract.py b/apps/api/tests/contract/test_retrieval_snapshot_large_corpus_contract.py index 0b45a703..a414e44b 100644 --- a/apps/api/tests/contract/test_retrieval_snapshot_large_corpus_contract.py +++ b/apps/api/tests/contract/test_retrieval_snapshot_large_corpus_contract.py @@ -6,7 +6,6 @@ from uuid import uuid4 from httpx import AsyncClient -import pytest from shared.models.database.document import Document, DocumentChunk, DocumentSection from shared.models.database.job_result import JobResult @@ -16,7 +15,6 @@ _REVISION_GROUP_SIZE, load_nav_snapshot, ) -import shared.services.retrieval.nav_snapshot as nav_snapshot_module from sqlalchemy import Executable, Result, select from sqlalchemy.engine import Row from sqlalchemy.sql.selectable import Select @@ -254,19 +252,10 @@ async def test_large_snapshot_keeps_all_retrieval_inputs_after_bounded_sql_load( developer_api_client_factory: Callable[ [], AbstractAsyncContextManager[AsyncClient] ], - monkeypatch: pytest.MonkeyPatch, ) -> None: namespace = f"large-corpus-{uuid4().hex[:8]}" async with developer_api_client_factory(): await _seed_large_retrieval_corpus(namespace) - async def unexpected_manifest_load(*_args: object, **_kwargs: object) -> None: - raise AssertionError("large snapshots must use normalized retrieval rows") - - monkeypatch.setattr( - nav_snapshot_module, - "_load_manifest_sections", - unexpected_manifest_load, - ) legacy_rows = await _load_legacy_rows(namespace) async with contract_db_session() as db: counting_db = _CountingSession(db) diff --git a/packages/shared-python/shared/services/retrieval/execution/routes.py b/packages/shared-python/shared/services/retrieval/execution/routes.py index 0998a8df..d2b6d211 100644 --- a/packages/shared-python/shared/services/retrieval/execution/routes.py +++ b/packages/shared-python/shared/services/retrieval/execution/routes.py @@ -23,6 +23,7 @@ ) from shared.services.retrieval.search.ranking import rank_retrieval_candidates from shared.services.retrieval.search.scoped_corpus import ( + count_manifest_chunks, count_scoped_chunks, load_all_scoped_chunks, ) @@ -57,15 +58,27 @@ async def _try_run_small_corpus_route( context: RetrievalRouteContext, ) -> RetrievalRouteOutcome | None: total_chunk_count: int | None = None - total_chunk_count = await count_scoped_chunks( - context.db, - user_id=context.user_id, - namespace=context.namespace, - exclude_document_ids=context.exclude_document_ids, - allowed_chunk_types=context.allowed_chunk_types, - revision_pins=context.revision_pins, - max_count=context.top_k + 1, - ) + if ( + context.use_agentic is not False + and context.revision_pins is not None + and not context.exclude_document_ids + and not context.exclude_sections + and context.allowed_chunk_types is None + and not context.signal_paths + ): + total_chunk_count = await count_manifest_chunks( + context.db, + revision_pins=context.revision_pins, + ) + if total_chunk_count is None: + total_chunk_count = await count_scoped_chunks( + context.db, + user_id=context.user_id, + namespace=context.namespace, + exclude_document_ids=context.exclude_document_ids, + allowed_chunk_types=context.allowed_chunk_types, + revision_pins=context.revision_pins, + ) logger.info(f"\n Total chunks in scope: {total_chunk_count}") if total_chunk_count > context.top_k: diff --git a/packages/shared-python/shared/services/retrieval/nav_snapshot.py b/packages/shared-python/shared/services/retrieval/nav_snapshot.py index 5d34e451..c68d6d8c 100644 --- a/packages/shared-python/shared/services/retrieval/nav_snapshot.py +++ b/packages/shared-python/shared/services/retrieval/nav_snapshot.py @@ -55,12 +55,6 @@ # Keep revision predicates bounded while reducing round trips for large # namespaces. Keyset paging still caps each payload query at 10,000 rows. _REVISION_GROUP_SIZE = 64 -# Compressed manifests are efficient for small namespaces, but transferring a -# large set of binary payloads can exceed the asyncpg statement timeout before -# PostgreSQL has done meaningful work. For larger snapshots, the normalized -# section/chunk index loaders below transfer only the fields map-nav needs and -# page them by revision. -_MANIFEST_MAX_REVISION_COUNT = 32 _logger = logging.getLogger(__name__) @@ -212,14 +206,12 @@ async def load_nav_snapshot( if job_result_id and job_id } - manifest_sections = None - if len(document_revisions) <= _MANIFEST_MAX_REVISION_COUNT: - manifest_sections = await _load_manifest_sections( - db, - document_revisions=document_revisions, - exclude_sections=excluded_secs, - job_id_by_result_id=job_id_by_result_id, - ) + manifest_sections = await _load_manifest_sections( + db, + document_revisions=document_revisions, + exclude_sections=excluded_secs, + job_id_by_result_id=job_id_by_result_id, + ) if manifest_sections is None: sections_by_doc, section_path_by_id = await _load_sections( db, @@ -366,31 +358,26 @@ async def _load_manifest_sections( if len(manifest_entries) != len(document_revisions): return None else: - manifest_entries = [] - for group_start in range(0, len(document_revisions), _REVISION_GROUP_SIZE): - revision_group = document_revisions[ - group_start : group_start + _REVISION_GROUP_SIZE - ] - statement = select( + statement = select( + RetrievalServingRevisionManifest.document_id, + RetrievalServingRevisionManifest.job_result_id, + RetrievalServingRevisionManifest.payload_zlib, + RetrievalServingRevisionManifest.checksum, + RetrievalServingRevisionManifest.format_version, + ).where( + tuple_( RetrievalServingRevisionManifest.document_id, RetrievalServingRevisionManifest.job_result_id, - RetrievalServingRevisionManifest.payload_zlib, - RetrievalServingRevisionManifest.checksum, - RetrievalServingRevisionManifest.format_version, - ).where( - tuple_( - RetrievalServingRevisionManifest.document_id, - RetrievalServingRevisionManifest.job_result_id, - ).in_(revision_group) - ) - try: - rows = (await db.execute(statement)).all() - except Exception: - await db.rollback() - return None - if len(rows) != len(revision_group): - return None - manifest_entries.extend(tuple(row) for row in rows) + ).in_(document_revisions) + ) + try: + rows = (await db.execute(statement)).all() + except Exception: + await db.rollback() + return None + if len(rows) != len(document_revisions): + return None + manifest_entries = [tuple(row) for row in rows] by_doc: dict[str, list[SectionRow]] = {} path_by_id: dict[str, str] = {} ids_by_doc: dict[str, list[str]] = {} 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 0e087549..ed6f2719 100644 --- a/packages/shared-python/shared/services/retrieval/search/scoped_corpus.py +++ b/packages/shared-python/shared/services/retrieval/search/scoped_corpus.py @@ -77,11 +77,10 @@ async def count_scoped_chunks( exclude_document_ids: list[str], allowed_chunk_types: set[str] | None, revision_pins: Mapping[str, str] | None = None, - max_count: int | None = None, ) -> int: if revision_pins is None: stmt = ( - select(DocumentChunk.id) + select(func.count(DocumentChunk.id)) .join( Document, (Document.document_id == DocumentChunk.document_id) @@ -93,7 +92,7 @@ async def count_scoped_chunks( ) else: stmt = ( - select(DocumentChunk.id) + select(func.count(DocumentChunk.id)) .join(Document, Document.document_id == DocumentChunk.document_id) .where(Document.user_id == user_id) .where(Document.namespace == namespace) @@ -107,13 +106,7 @@ async def count_scoped_chunks( stmt = stmt.where(Document.document_id.notin_(list(exclude_document_ids))) if allowed_chunk_types is not None: stmt = stmt.where(func.lower(DocumentChunk.chunk_type).in_(list(allowed_chunk_types))) - - if max_count is not None: - stmt = stmt.limit(max_count) - - result = await db.execute( - select(func.count()).select_from(stmt.order_by(None).subquery()) - ) + result = await db.execute(stmt) return result.scalar() or 0 From ff1ce7ce6618ab9c686f9b1e0ad20122613d7ddf Mon Sep 17 00:00:00 2001 From: suguanYang Date: Mon, 31 Aug 2026 11:06:31 +0800 Subject: [PATCH 02/10] Revert "Merge pull request #354 from Ontos-AI/perf/wangbinqi/retrieval-hydration-job-id" This reverts commit e1b8c886118f432fdd7f0a6a4eed76f1293d29e9, reversing changes made to 464ae7b4573fbdb4472836d983b60eaaf6cb34bd. --- .../services/retrieval/hydration/reference.py | 85 +++++++++---------- 1 file changed, 39 insertions(+), 46 deletions(-) diff --git a/packages/shared-python/shared/services/retrieval/hydration/reference.py b/packages/shared-python/shared/services/retrieval/hydration/reference.py index 753fc66f..cd240654 100644 --- a/packages/shared-python/shared/services/retrieval/hydration/reference.py +++ b/packages/shared-python/shared/services/retrieval/hydration/reference.py @@ -28,10 +28,10 @@ async def hydrate_referenced_chunk_rows( ref_keys = [ build_reference_lookup_key( - document_id=ref.get("document_id"), - chunk_id=ref.get("chunk_id"), - section_path=ref.get("section_path"), - file_path=ref.get("file_path"), + document_id=ref.get('document_id'), + chunk_id=ref.get('chunk_id'), + section_path=ref.get('section_path'), + file_path=ref.get('file_path'), ) for ref in refs ] @@ -42,16 +42,15 @@ async def hydrate_referenced_chunk_rows( document_ids = sorted({document_id for document_id, _, _, _ in ref_keys}) chunk_ids = sorted({chunk_id for _, chunk_id, _, _ in ref_keys}) pinned_document_ids = [ - document_id - for document_id in document_ids - if revision_pins and document_id in revision_pins + document_id for document_id in document_ids if revision_pins and document_id in revision_pins ] if revision_pins is not None and not pinned_document_ids: return [] if revision_pins is None: - chunk_join = (DocumentChunk.document_id == Document.document_id) & ( - DocumentChunk.job_result_id == Document.current_job_result_id + chunk_join = ( + (DocumentChunk.document_id == Document.document_id) + & (DocumentChunk.job_result_id == Document.current_job_result_id) ) else: chunk_join = and_( @@ -68,15 +67,9 @@ async def hydrate_referenced_chunk_rows( ) stmt = ( - # Select only the job identifier needed for the public projection. - # 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) .join(DocumentChunk, chunk_join) - .outerjoin( - DocumentSection, DocumentSection.section_id == DocumentChunk.section_id - ) + .outerjoin(DocumentSection, DocumentSection.section_id == DocumentChunk.section_id) .join(JobResult, JobResult.id == DocumentChunk.job_result_id) .where(Document.user_id == user_id) .where(Document.namespace == namespace) @@ -89,40 +82,40 @@ async def hydrate_referenced_chunk_rows( .order_by(DocumentChunk.sort_order) ) if revision_pins is None: - stmt = stmt.where(Document.status == "active") + stmt = stmt.where(Document.status == 'active') result = await db.execute(stmt) 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_result in result.all(): row = { - "document_id": document.document_id, - "chunk_id": chunk.chunk_id, - "section_id": chunk.section_id, - "section_path": section.section_path if section else None, - "source_file_name": document.source_file_name, - "chunk_type": chunk.chunk_type, - "content": chunk.content, + 'document_id': document.document_id, + 'chunk_id': chunk.chunk_id, + 'section_id': chunk.section_id, + 'section_path': section.section_path if section else None, + 'source_file_name': document.source_file_name, + 'chunk_type': chunk.chunk_type, + 'content': chunk.content, # Use the caller-supplied score when available (e.g. discovery RRF or # KG confidence). None signals "no score known" so consumers can # distinguish unscored chunks from genuinely high-scoring ones. - "score": ( + 'score': ( score_by_chunk_id.get(chunk.chunk_id) if score_by_chunk_id is not None else None ), - "file_path": chunk.file_path, - "chunk_metadata": chunk.chunk_metadata or {}, - "job_result_id": chunk.job_result_id, - "job_id": job_id, - "source_chunk_path": chunk.source_chunk_path, - "sort_order": chunk.sort_order, + 'file_path': chunk.file_path, + 'chunk_metadata': chunk.chunk_metadata or {}, + 'job_result_id': chunk.job_result_id, + 'job_id': job_result.job_id if job_result else None, + 'source_chunk_path': chunk.source_chunk_path, + 'sort_order': chunk.sort_order, } key = build_reference_lookup_key( - document_id=row["document_id"], - chunk_id=row["chunk_id"], - section_path=row["section_path"], - file_path=row["file_path"], + document_id=row['document_id'], + chunk_id=row['chunk_id'], + section_path=row['section_path'], + file_path=row['file_path'], ) rows_by_key[key] = row rows_by_base_key.setdefault((key[0], key[1]), []).append(row) @@ -138,7 +131,7 @@ async def hydrate_referenced_chunk_rows( candidate for candidate in candidates if key[2] - and str(candidate.get("section_path") or "").strip() == key[2] + and str(candidate.get('section_path') or '').strip() == key[2] ), None, ) @@ -148,10 +141,10 @@ async def hydrate_referenced_chunk_rows( candidate for candidate in candidates if build_reference_lookup_key( - document_id=candidate.get("document_id"), - chunk_id=candidate.get("chunk_id"), - section_path=candidate.get("section_path"), - file_path=candidate.get("file_path"), + document_id=candidate.get('document_id'), + chunk_id=candidate.get('chunk_id'), + section_path=candidate.get('section_path'), + file_path=candidate.get('file_path'), ) not in seen_keys ), @@ -159,10 +152,10 @@ async def hydrate_referenced_chunk_rows( ) if row is not None: row_key = build_reference_lookup_key( - document_id=row.get("document_id"), - chunk_id=row.get("chunk_id"), - section_path=row.get("section_path"), - file_path=row.get("file_path"), + document_id=row.get('document_id'), + chunk_id=row.get('chunk_id'), + section_path=row.get('section_path'), + file_path=row.get('file_path'), ) if row_key in seen_keys: continue From b365a59bd6b8426afd93e7731ced6a7f0c53ba12 Mon Sep 17 00:00:00 2001 From: suguanYang Date: Mon, 31 Aug 2026 11:06:43 +0800 Subject: [PATCH 03/10] Revert "Merge pull request #352 from Ontos-AI/perf/wangbinqi/mapnav-timing-instrumentation" This reverts commit 464ae7b4573fbdb4472836d983b60eaaf6cb34bd, reversing changes made to 9b355e5b7d0c66d68e11a49de118e3cd9d063589. --- CONTEXT.md | 70 --- ...b6c7d_add_map_unit_token_covering_index.py | 40 -- ...6c7d8e_add_section_snapshot_order_index.py | 41 -- ...d8e9f_add_retrieval_serving_generations.py | 171 ------ .../5b6c7d8e9f0a_add_term_trigram_indexes.py | 35 -- .../services/documents/lifecycle_service.py | 47 +- apps/api/scripts/backfill_map_unit_indexes.py | 62 +- .../test_bm25_fts_prefilter_contract.py | 38 +- .../tests/contract/test_documents_contract.py | 86 +-- ...etrieval_lazy_snapshot_quality_contract.py | 90 +-- .../test_retrieval_lazy_tree_contract.py | 33 -- .../test_retrieval_manifest_cache_contract.py | 56 -- .../test_retrieval_map_unit_index_contract.py | 152 +---- .../test_retrieval_mapnav_session_contract.py | 4 - ...test_retrieval_relit_map_cache_contract.py | 27 - .../test_retrieval_revision_races_contract.py | 108 ---- .../test_retrieval_rrf_duplicate_contract.py | 19 - ...est_retrieval_serving_manifest_contract.py | 51 -- ...retrieval_snapshot_consistency_contract.py | 54 -- .../test_retrieval_term_score_contract.py | 75 --- ...0005-stream-retrieval-progress-over-sse.md | 47 -- ...mically-publish-retrieval-serving-index.md | 6 - ...-coherent-retrieval-serving-generations.md | 6 - ...enance-window-for-serving-index-rollout.md | 6 - docs/adr/README.md | 5 +- docs/design/retrieval-serving-index-plan.md | 425 -------------- docs/design/retrieval-streaming-sse.md | 207 ------- .../shared/models/database/document.py | 178 ------ .../services/retrieval/execution/plan.py | 26 +- .../retrieval/execution/reference_resolver.py | 3 - .../retrieval/execution/revision_pins.py | 92 --- .../retrieval/execution/route_types.py | 3 - .../services/retrieval/execution/routes.py | 81 +-- .../services/retrieval/hydration/connected.py | 22 +- .../services/retrieval/hydration/reference.py | 44 +- .../retrieval/hydration/result_assembly.py | 3 - .../services/retrieval/manifest_cache.py | 47 -- .../services/retrieval/nav/nav_hierarchy.py | 18 +- .../services/retrieval/nav/nav_knowhere.py | 401 +++---------- .../services/retrieval/nav/nav_map_scores.py | 31 - .../services/retrieval/nav/nav_orchestrate.py | 3 - .../services/retrieval/nav/nav_types.py | 6 - .../shared/services/retrieval/nav_snapshot.py | 540 +++--------------- .../services/retrieval/publication_content.py | 3 - .../services/retrieval/publication_service.py | 41 -- .../services/retrieval/search/channels.py | 53 +- .../services/retrieval/search/discovery.py | 56 +- .../services/retrieval/search/ranking.py | 109 ++-- .../retrieval/search/scoped_corpus.py | 119 +--- .../services/retrieval/search/scoring.py | 10 +- .../services/retrieval/serving_generation.py | 60 -- .../services/retrieval/serving_manifest.py | 405 ------------- 52 files changed, 297 insertions(+), 4018 deletions(-) delete mode 100644 apps/api/alembic/versions/2e3f4a5b6c7d_add_map_unit_token_covering_index.py delete mode 100644 apps/api/alembic/versions/3f4a5b6c7d8e_add_section_snapshot_order_index.py delete mode 100644 apps/api/alembic/versions/4a5b6c7d8e9f_add_retrieval_serving_generations.py delete mode 100644 apps/api/alembic/versions/5b6c7d8e9f0a_add_term_trigram_indexes.py delete mode 100644 apps/api/tests/contract/test_retrieval_lazy_tree_contract.py delete mode 100644 apps/api/tests/contract/test_retrieval_manifest_cache_contract.py delete mode 100644 apps/api/tests/contract/test_retrieval_relit_map_cache_contract.py delete mode 100644 apps/api/tests/contract/test_retrieval_revision_races_contract.py delete mode 100644 apps/api/tests/contract/test_retrieval_rrf_duplicate_contract.py delete mode 100644 apps/api/tests/contract/test_retrieval_serving_manifest_contract.py delete mode 100644 apps/api/tests/contract/test_retrieval_term_score_contract.py delete mode 100644 docs/adr/0005-stream-retrieval-progress-over-sse.md delete mode 100644 docs/adr/0006-atomically-publish-retrieval-serving-index.md delete mode 100644 docs/adr/0007-use-coherent-retrieval-serving-generations.md delete mode 100644 docs/adr/0008-use-a-maintenance-window-for-serving-index-rollout.md delete mode 100644 docs/design/retrieval-serving-index-plan.md delete mode 100644 docs/design/retrieval-streaming-sse.md delete mode 100644 packages/shared-python/shared/services/retrieval/execution/revision_pins.py delete mode 100644 packages/shared-python/shared/services/retrieval/manifest_cache.py delete mode 100644 packages/shared-python/shared/services/retrieval/serving_generation.py delete mode 100644 packages/shared-python/shared/services/retrieval/serving_manifest.py diff --git a/CONTEXT.md b/CONTEXT.md index b67e16d0..c8903a58 100644 --- a/CONTEXT.md +++ b/CONTEXT.md @@ -173,76 +173,6 @@ The query workflow that returns cited evidence from published documents. The typed retrieval request that owns cache-shaping fields and route policy: scope, filters, data type, channels, ranking options, and agentic toggle. -### Retrieval Run - -One execution of a Retrieval Query, recorded across classic, map-nav, -small-corpus, cache-hit, failed, and cancelled outcomes with route, timing, -and terminal-status metadata. - -### Retrieval Progress Event - -A safe, user-facing update about the current phase of a Retrieval run. It uses -the fixed phases `started`, `planning`, `searching`, `reviewing_sources`, and -`finalizing`; it never contains chain-of-thought or raw planner output. - -### Retrieval Stream - -The server-to-client event stream for a Retrieval Query. It carries -Retrieval Progress Events during execution and one authoritative final result -or terminal failure, while leaving answer generation to downstream clients. - -### Retrieval Duration - -The end-to-end server time for a Retrieval Query, measured from retrieval -execution start through final public-result assembly. It includes cache lookup -and excludes authentication, network transfer, SSE delivery time, and -downstream answer generation. - -### Retrieval Non-LLM Work - -The database and retrieval-engine work for a Retrieval Query: snapshot or -serving-index loading, lexical scoring, ranking, result hydration, citation -assembly, and asset-reference resolution. It excludes planner, harvest, -control, and answer-generation model time, which are measured separately. - -### Retrieval Serving Index - -The publication-derived read model used to load retrieval structure and -scoring inputs without rebuilding them from the full document corpus for each -query. It is revision-pinned and complete before its document revision becomes -active. - -### Retrieval Serving Fallback - -The exact legacy retrieval path used when a serving index is missing, -incomplete, or inconsistent. It preserves retrieval quality while sacrificing -the serving-index latency target until the derived data is repaired. - -### Retrieval Serving Generation - -The namespace-scoped version that identifies one coherent set of active -document revisions and their serving-index statistics. Retrieval captures one -generation and retries or falls back if publication changes it during capture. - -### Retrieval Semantic Parity - -The compatibility requirement that a serving-index retrieval returns the same -selected chunk IDs, ordering, rounded scores, citations, and asset references -as the legacy retrieval path for the same request. - -### Retrieval Revision Pin - -The set of document revision IDs captured at retrieval start and used for the -entire retrieval run, including lazy content and asset resolution. A later -publication affects subsequent runs, not the run already in progress. - -### Online Retrieval Serving Rollout - -The additive rollout of retrieval-serving schema and derived data while -retrieval and document publication remain available. Incomplete or -inconsistent revisions use the exact legacy retrieval path until backfill and -validation finish. - ### Workflow Run Request The agentic Retrieval request passed through planning and step execution. It diff --git a/apps/api/alembic/versions/2e3f4a5b6c7d_add_map_unit_token_covering_index.py b/apps/api/alembic/versions/2e3f4a5b6c7d_add_map_unit_token_covering_index.py deleted file mode 100644 index 61250639..00000000 --- a/apps/api/alembic/versions/2e3f4a5b6c7d_add_map_unit_token_covering_index.py +++ /dev/null @@ -1,40 +0,0 @@ -"""Add a covering index for persisted map-nav token lookups.""" - -from __future__ import annotations - -from alembic import op - - -revision = "2e3f4a5b6c7d" -down_revision = "1d2e3f4a5b6c" -branch_labels = None -depends_on = None - -_INDEX_NAME = "idx_document_map_unit_tokens_unit_lookup" - - -def upgrade() -> None: - external_transaction = bool( - op.get_context().opts.get("knowhere_external_transaction", False) - ) - statement = ( - f"CREATE INDEX {{concurrently}}IF NOT EXISTS {_INDEX_NAME} " - "ON document_map_unit_tokens (map_unit_id, channel, token_hash) " - "INCLUDE (token, frequency)" - ) - if external_transaction: - op.execute(statement.format(concurrently="")) - return - with op.get_context().autocommit_block(): - op.execute(statement.format(concurrently="CONCURRENTLY ")) - - -def downgrade() -> None: - external_transaction = bool( - op.get_context().opts.get("knowhere_external_transaction", False) - ) - if external_transaction: - op.execute(f"DROP INDEX IF EXISTS {_INDEX_NAME}") - return - with op.get_context().autocommit_block(): - op.execute(f"DROP INDEX CONCURRENTLY IF EXISTS {_INDEX_NAME}") diff --git a/apps/api/alembic/versions/3f4a5b6c7d8e_add_section_snapshot_order_index.py b/apps/api/alembic/versions/3f4a5b6c7d8e_add_section_snapshot_order_index.py deleted file mode 100644 index 4dcc5011..00000000 --- a/apps/api/alembic/versions/3f4a5b6c7d8e_add_section_snapshot_order_index.py +++ /dev/null @@ -1,41 +0,0 @@ -"""Add the index used by lazy map-nav section pagination.""" - -from __future__ import annotations - -from alembic import op - - -revision = "3f4a5b6c7d8e" -down_revision = "2e3f4a5b6c7d" -branch_labels = None -depends_on = None - -_INDEX_NAME = "idx_document_sections_revision_snapshot_order" - - -def upgrade() -> None: - external_transaction = bool( - op.get_context().opts.get("knowhere_external_transaction", False) - ) - statement = ( - f"CREATE INDEX {{concurrently}}IF NOT EXISTS {_INDEX_NAME} " - "ON document_sections " - "(document_id, job_result_id, sort_order, section_id)" - ) - if external_transaction: - op.execute(statement.format(concurrently="")) - return - with op.get_context().autocommit_block(): - op.execute(statement.format(concurrently="CONCURRENTLY ")) - - -def downgrade() -> None: - external_transaction = bool( - op.get_context().opts.get("knowhere_external_transaction", False) - ) - statement = f"DROP INDEX {{concurrently}}IF EXISTS {_INDEX_NAME}" - if external_transaction: - op.execute(statement.format(concurrently="")) - return - with op.get_context().autocommit_block(): - op.execute(statement.format(concurrently="CONCURRENTLY ")) diff --git a/apps/api/alembic/versions/4a5b6c7d8e9f_add_retrieval_serving_generations.py b/apps/api/alembic/versions/4a5b6c7d8e9f_add_retrieval_serving_generations.py deleted file mode 100644 index 426ff6e0..00000000 --- a/apps/api/alembic/versions/4a5b6c7d8e9f_add_retrieval_serving_generations.py +++ /dev/null @@ -1,171 +0,0 @@ -"""Add revision-pinned serving manifests and namespace statistics.""" - -from __future__ import annotations - -from collections.abc import Sequence - -from alembic import op -import sqlalchemy as sa - - -revision: str = "4a5b6c7d8e9f" -down_revision: str | None = "3f4a5b6c7d8e" -branch_labels: Sequence[str] | None = None -depends_on: Sequence[str] | None = None - - -def upgrade() -> None: - if not sa.inspect(op.get_bind()).has_table("retrieval_namespace_generations"): - op.create_table( - "retrieval_namespace_generations", - sa.Column("id", sa.String(length=100), nullable=False), - sa.Column("user_id", sa.Text(), nullable=False), - sa.Column("namespace", sa.String(length=255), nullable=False), - sa.Column( - "generation", sa.BigInteger(), nullable=False, server_default="0" - ), - sa.Column("updated_at", sa.DateTime(), nullable=False), - sa.PrimaryKeyConstraint("id"), - sa.UniqueConstraint( - "user_id", - "namespace", - name="uq_retrieval_namespace_generations_scope", - ), - ) - if not sa.inspect(op.get_bind()).has_table("retrieval_serving_revision_manifests"): - op.create_table( - "retrieval_serving_revision_manifests", - sa.Column("id", sa.String(length=100), nullable=False), - sa.Column("user_id", sa.Text(), nullable=False), - sa.Column("namespace", sa.String(length=255), nullable=False), - sa.Column("document_id", sa.String(length=36), nullable=False), - sa.Column("job_result_id", sa.String(length=36), nullable=False), - sa.Column("format_version", sa.Integer(), nullable=False), - sa.Column("payload_zlib", sa.LargeBinary(), nullable=False), - sa.Column("checksum", sa.String(length=64), nullable=False), - sa.Column("created_at", sa.DateTime(), nullable=False), - sa.ForeignKeyConstraint( - ["document_id"], ["documents.document_id"], ondelete="CASCADE" - ), - sa.ForeignKeyConstraint( - ["job_result_id"], ["job_results.id"], ondelete="CASCADE" - ), - sa.PrimaryKeyConstraint("id"), - sa.UniqueConstraint( - "document_id", - "job_result_id", - name="uq_retrieval_serving_revision_manifests_revision", - ), - ) - if not _index_exists("idx_retrieval_serving_revision_manifests_scope"): - op.create_index( - "idx_retrieval_serving_revision_manifests_scope", - "retrieval_serving_revision_manifests", - ["user_id", "namespace", "document_id", "job_result_id"], - ) - if not sa.inspect(op.get_bind()).has_table("retrieval_serving_revision_stats"): - op.create_table( - "retrieval_serving_revision_stats", - sa.Column("id", sa.String(length=100), nullable=False), - sa.Column("user_id", sa.Text(), nullable=False), - sa.Column("namespace", sa.String(length=255), nullable=False), - sa.Column("document_id", sa.String(length=36), nullable=False), - sa.Column("job_result_id", sa.String(length=36), nullable=False), - sa.Column("format_version", sa.Integer(), nullable=False), - sa.Column("payload_zlib", sa.LargeBinary(), nullable=False), - sa.Column("checksum", sa.String(length=64), nullable=False), - sa.Column("created_at", sa.DateTime(), nullable=False), - sa.ForeignKeyConstraint( - ["document_id"], ["documents.document_id"], ondelete="CASCADE" - ), - sa.ForeignKeyConstraint( - ["job_result_id"], ["job_results.id"], ondelete="CASCADE" - ), - sa.PrimaryKeyConstraint("id"), - sa.UniqueConstraint( - "document_id", - "job_result_id", - name="uq_retrieval_serving_revision_stats_revision", - ), - ) - if not _index_exists("idx_retrieval_serving_revision_stats_scope"): - op.create_index( - "idx_retrieval_serving_revision_stats_scope", - "retrieval_serving_revision_stats", - ["user_id", "namespace", "document_id", "job_result_id"], - ) - if not sa.inspect(op.get_bind()).has_table("retrieval_namespace_stats"): - op.create_table( - "retrieval_namespace_stats", - sa.Column("id", sa.String(length=100), nullable=False), - sa.Column("user_id", sa.Text(), nullable=False), - sa.Column("namespace", sa.String(length=255), nullable=False), - sa.Column("generation", sa.BigInteger(), nullable=False), - sa.Column("payload_zlib", sa.LargeBinary(), nullable=False), - sa.Column("checksum", sa.String(length=64), nullable=False), - sa.Column("updated_at", sa.DateTime(), nullable=False), - sa.PrimaryKeyConstraint("id"), - sa.UniqueConstraint( - "user_id", "namespace", name="uq_retrieval_namespace_stats_scope" - ), - ) - if not sa.inspect(op.get_bind()).has_table("retrieval_namespace_token_stats"): - op.create_table( - "retrieval_namespace_token_stats", - sa.Column("id", sa.String(length=100), nullable=False), - sa.Column("user_id", sa.Text(), nullable=False), - sa.Column("namespace", sa.String(length=255), nullable=False), - sa.Column("generation", sa.BigInteger(), nullable=False), - sa.Column("channel", sa.String(length=32), nullable=False), - sa.Column("token_hash", sa.String(length=64), nullable=False), - sa.Column("document_frequency", sa.Integer(), nullable=False), - sa.PrimaryKeyConstraint("id"), - sa.UniqueConstraint( - "user_id", - "namespace", - "channel", - "token_hash", - name="uq_retrieval_namespace_token_stats_key", - ), - ) - if not _index_exists("idx_retrieval_namespace_token_stats_lookup"): - op.create_index( - "idx_retrieval_namespace_token_stats_lookup", - "retrieval_namespace_token_stats", - ["user_id", "namespace", "generation", "channel", "token_hash"], - ) - - -def _index_exists(index_name: str) -> bool: - bind = op.get_bind() - inspector = sa.inspect(bind) - for table_name in inspector.get_table_names(): - if any( - index.get("name") == index_name - for index in inspector.get_indexes(table_name) - ): - return True - return False - - -def downgrade() -> None: - op.drop_index( - "idx_retrieval_namespace_token_stats_lookup", - table_name="retrieval_namespace_token_stats", - if_exists=True, - ) - op.drop_table("retrieval_namespace_token_stats", if_exists=True) - op.drop_table("retrieval_namespace_stats", if_exists=True) - op.drop_index( - "idx_retrieval_serving_revision_stats_scope", - table_name="retrieval_serving_revision_stats", - if_exists=True, - ) - op.drop_table("retrieval_serving_revision_stats", if_exists=True) - op.drop_index( - "idx_retrieval_serving_revision_manifests_scope", - table_name="retrieval_serving_revision_manifests", - if_exists=True, - ) - op.drop_table("retrieval_serving_revision_manifests", if_exists=True) - op.drop_table("retrieval_namespace_generations", if_exists=True) diff --git a/apps/api/alembic/versions/5b6c7d8e9f0a_add_term_trigram_indexes.py b/apps/api/alembic/versions/5b6c7d8e9f0a_add_term_trigram_indexes.py deleted file mode 100644 index 71720ff0..00000000 --- a/apps/api/alembic/versions/5b6c7d8e9f0a_add_term_trigram_indexes.py +++ /dev/null @@ -1,35 +0,0 @@ -"""Add trigram acceleration for exact term-channel candidate discovery.""" - -from __future__ import annotations - -from collections.abc import Sequence - -from alembic import op - - -revision: str = "5b6c7d8e9f0a" -down_revision: str | None = "4a5b6c7d8e9f" -branch_labels: Sequence[str] | None = None -depends_on: Sequence[str] | None = None - -_MAP_UNIT_INDEX = "idx_document_map_units_term_trgm" -_CHUNK_INDEX = "idx_document_chunks_term_trgm" - - -def upgrade() -> None: - op.execute("CREATE EXTENSION IF NOT EXISTS pg_trgm") - op.execute( - f"CREATE INDEX IF NOT EXISTS {_MAP_UNIT_INDEX} " - "ON document_map_units USING gin " - "(term_search_text_lower gin_trgm_ops)" - ) - op.execute( - f"CREATE INDEX IF NOT EXISTS {_CHUNK_INDEX} " - "ON document_chunks USING gin " - "(lower(COALESCE(term_search_text, '')) gin_trgm_ops)" - ) - - -def downgrade() -> None: - op.execute(f"DROP INDEX IF EXISTS {_CHUNK_INDEX}") - op.execute(f"DROP INDEX IF EXISTS {_MAP_UNIT_INDEX}") diff --git a/apps/api/app/services/documents/lifecycle_service.py b/apps/api/app/services/documents/lifecycle_service.py index 6b39507d..1bd5ae3a 100644 --- a/apps/api/app/services/documents/lifecycle_service.py +++ b/apps/api/app/services/documents/lifecycle_service.py @@ -8,25 +8,13 @@ from app.repositories.document_repository import DocumentRepository from loguru import logger -from sqlalchemy import delete from sqlalchemy.ext.asyncio import AsyncSession -from shared.models.database.document import ( - DocumentChunk, - DocumentSection, - RetrievalServingRevisionStat, -) +from shared.models.database.document import DocumentChunk, DocumentSection from shared.services.retrieval.cache_service import ( invalidate_retrieval_cache_namespaces, ) from shared.services.retrieval.graph.service import DocumentGraphService, GraphScope -from shared.services.retrieval.serving_generation import ( - advance_namespace_generation, - lock_namespace_generation, -) -from shared.services.retrieval.serving_manifest import ( - rebuild_namespace_serving_statistics, -) from shared.services.storage.result_storage import ResultStorage, get_result_storage _DOCUMENT_CHUNK_ASSET_URL_EXPIRES_SECONDS = 7 * 24 * 60 * 60 @@ -115,7 +103,7 @@ def _normalize_page_asset(raw_asset: dict[str, Any]) -> dict[str, Any] | None: "content_type": content_type, "source": source, } - if asset_url := str(raw_asset.get("asset_url") or "").strip(): + if (asset_url := str(raw_asset.get("asset_url") or "").strip()): asset["asset_url"] = asset_url if (width := _positive_int(raw_asset.get("width"))) is not None: asset["width"] = width @@ -458,38 +446,7 @@ async def archive_document( return document_payload(document) previous_namespace = document.namespace - await db.run_sync( - lambda sync_db: lock_namespace_generation( - sync_db, - user_id=user_id, - namespace=previous_namespace, - ) - ) await self._repository.archive_document(db, document=document) - current_revision = document.current_job_result_id - if current_revision: - await db.run_sync( - lambda sync_db: sync_db.execute( - delete(RetrievalServingRevisionStat).where( - RetrievalServingRevisionStat.document_id == document_id, - RetrievalServingRevisionStat.job_result_id == current_revision, - ) - ) - ) - await db.run_sync( - lambda sync_db: rebuild_namespace_serving_statistics( - sync_db, - user_id=user_id, - namespace=previous_namespace, - ) - ) - await db.run_sync( - lambda sync_db: advance_namespace_generation( - sync_db, - user_id=user_id, - namespace=previous_namespace, - ) - ) await db.run_sync( lambda sync_db: self._graph_service.remove_document_graph( sync_db, diff --git a/apps/api/scripts/backfill_map_unit_indexes.py b/apps/api/scripts/backfill_map_unit_indexes.py index 2bfd4177..c0301d38 100644 --- a/apps/api/scripts/backfill_map_unit_indexes.py +++ b/apps/api/scripts/backfill_map_unit_indexes.py @@ -45,14 +45,6 @@ def _bootstrap_python_path() -> None: from shared.models.database.document import Document from shared.services.retrieval.map_unit_index import replace_document_map_units from shared.services.retrieval.publication_models import DocumentPublicationScope -from shared.services.retrieval.serving_generation import ( - advance_namespace_generation, - lock_namespace_generation, -) -from shared.services.retrieval.serving_manifest import ( - persist_revision_serving_state, - rebuild_namespace_serving_statistics, -) def _build_parser() -> argparse.ArgumentParser: @@ -73,11 +65,7 @@ def _build_parser() -> argparse.ArgumentParser: def _load_documents(document_id: str) -> list[Document]: session_factory = get_sync_session_factory() with session_factory() as db: - statement = ( - select(Document) - .where(Document.status == "active") - .where(Document.current_job_result_id.is_not(None)) - ) + statement = select(Document).where(Document.current_job_result_id.is_not(None)) normalized_document_id = document_id.strip() if normalized_document_id: statement = statement.where(Document.document_id == normalized_document_id) @@ -98,49 +86,15 @@ def backfill_map_unit_indexes(*, apply: bool, document_id: str = "") -> int: job_result_id = document.current_job_result_id if not job_result_id: continue + scope = DocumentPublicationScope( + user_id=document.user_id, + namespace=document.namespace, + document_id=document.document_id, + job_result_id=job_result_id, + source_file_name=str(document.source_file_name or ""), + ) with session_factory() as db: - lock_namespace_generation( - db, - user_id=document.user_id, - namespace=document.namespace, - ) - locked_document = db.execute( - select(Document) - .where(Document.document_id == document.document_id) - .with_for_update() - ).scalar_one_or_none() - if ( - locked_document is None - or locked_document.status != "active" - or locked_document.current_job_result_id != job_result_id - or locked_document.user_id != document.user_id - or locked_document.namespace != document.namespace - ): - db.rollback() - print( - f"skipped stale or inactive document={document.document_id} " - f"revision={job_result_id}" - ) - continue - scope = DocumentPublicationScope( - user_id=locked_document.user_id, - namespace=locked_document.namespace, - document_id=locked_document.document_id, - job_result_id=job_result_id, - source_file_name=str(locked_document.source_file_name or ""), - ) replace_document_map_units(db, scope=scope) - persist_revision_serving_state(db, scope=scope) - rebuild_namespace_serving_statistics( - db, - user_id=scope.user_id, - namespace=scope.namespace, - ) - advance_namespace_generation( - db, - user_id=scope.user_id, - namespace=scope.namespace, - ) db.commit() print(f"backfilled document={document.document_id} revision={job_result_id}") return len(documents) diff --git a/apps/api/tests/contract/test_bm25_fts_prefilter_contract.py b/apps/api/tests/contract/test_bm25_fts_prefilter_contract.py index 56add8e5..6b958329 100644 --- a/apps/api/tests/contract/test_bm25_fts_prefilter_contract.py +++ b/apps/api/tests/contract/test_bm25_fts_prefilter_contract.py @@ -23,10 +23,10 @@ user_id TEXT, namespace TEXT, status TEXT, - current_job_result_id TEXT, + current_job_result_id INTEGER, source_file_name TEXT ); -CREATE TABLE job_results (id TEXT PRIMARY KEY, job_id TEXT); +CREATE TABLE job_results (id INTEGER PRIMARY KEY, job_id TEXT); CREATE TABLE document_sections (section_id TEXT PRIMARY KEY, section_path TEXT); CREATE TABLE document_chunks ( id SERIAL PRIMARY KEY, @@ -38,7 +38,7 @@ source_chunk_path TEXT, file_path TEXT, chunk_metadata JSONB, - job_result_id TEXT, + job_result_id INTEGER, sort_order INTEGER, content_search_text TEXT, content_search_tsv TSVECTOR GENERATED ALWAYS AS @@ -197,35 +197,3 @@ async def test_exclusions_still_apply_under_the_prefilter( exclude_sections=[], ) assert rows == [] - - -@pytest.mark.asyncio -async def test_content_channel_uses_the_requested_revision_pin( - seeded_session: AsyncSession, -) -> None: - await seeded_session.execute(text("INSERT INTO job_results VALUES (2, 'job2')")) - await seeded_session.execute( - text( - "INSERT INTO document_chunks " - "(chunk_id, document_id, section_id, chunk_type, content, " - " job_result_id, sort_order, content_search_text, path_search_text) " - "VALUES ('new-hit', 'd1', 's1', 'text', 'new body', 2, 1, " - " 'alpha replacement', 'new path')" - ) - ) - await seeded_session.execute( - text("UPDATE documents SET current_job_result_id = 2 WHERE document_id = 'd1'") - ) - - rows = await content_channel( - seeded_session, - user_id="u1", - namespace="ns1", - query="alpha", - top_k=50, - exclude_document_ids=[], - exclude_sections=[], - revision_pins={"d1": "1"}, - ) - - assert [str(row["chunk_id"]) for row in rows] == ["hit-en"] diff --git a/apps/api/tests/contract/test_documents_contract.py b/apps/api/tests/contract/test_documents_contract.py index a71d4029..6b8e33b0 100644 --- a/apps/api/tests/contract/test_documents_contract.py +++ b/apps/api/tests/contract/test_documents_contract.py @@ -13,7 +13,6 @@ from tests.support.contract_database import ContractDatabase from shared.testing.contract_runtime import get_contract_database_url -from shared.services.retrieval.serving_manifest import encode_serving_manifest async def _create_contract_engine() -> AsyncEngine: @@ -1020,7 +1019,9 @@ async def test_should_include_media_asset_urls_in_document_chunk_list_when_reque assert chunks[1]["asset_url"] == expected_asset_url assert default_response.status_code == 200 - default_chunks = cast(list[dict[str, object]], default_response.json()["chunks"]) + default_chunks = cast( + list[dict[str, object]], default_response.json()["chunks"] + ) assert default_chunks[1]["asset_url"] is None @@ -1333,84 +1334,3 @@ async def test_should_archive_a_document_via_the_legacy_archive_route( assert response_json["archived_at"] assert persisted_document["status"] == "archived" assert persisted_document["archived_at"] is not None - - -@pytest.mark.asyncio -async def test_archive_removes_revision_serving_stats_and_advances_generation( - developer_api_client_factory: Callable[ - [], AbstractAsyncContextManager[AsyncClient] - ], -) -> None: - document_id = f"doc_{uuid4().hex[:12]}" - namespace = f"archive-serving-{uuid4().hex[:8]}" - async with developer_api_client_factory() as api_client: - revision = await _insert_document_revision_with_chunks( - document_id=document_id, - namespace=namespace, - chunks=[ - { - "id": f"dchk_{uuid4().hex[:12]}", - "chunk_id": "archive-serving-chunk", - "chunk_type": "text", - "content": "serving contribution", - "source_chunk_path": "Archive/Serving", - "metadata": {}, - } - ], - ) - payload_bytes, checksum, version = encode_serving_manifest( - { - "document_id": document_id, - "job_result_id": revision["job_result_id"], - "unit_count": 1, - "path_token_count": 1, - "content_token_count": 2, - "token_frequencies": { - "path": {"archive": 1}, - "content": {"serving": 1}, - }, - } - ) - await ContractDatabase.execute( - """ - INSERT INTO retrieval_serving_revision_stats ( - id, user_id, namespace, document_id, job_result_id, - format_version, payload_zlib, checksum, created_at - ) VALUES ( - :id, :user_id, :namespace, :document_id, :job_result_id, - :format_version, :payload_zlib, :checksum, NOW() - ) - """, - { - "id": f"rss_{uuid4().hex[:12]}", - "user_id": "local-dev-user", - "namespace": namespace, - "document_id": document_id, - "job_result_id": revision["job_result_id"], - "format_version": version, - "payload_zlib": payload_bytes, - "checksum": checksum, - }, - ) - response = await api_client.post(f"/api/v1/documents/{document_id}/archive") - - assert response.status_code == 200 - remaining_revision_stats = await ContractDatabase.fetch_one( - """ - SELECT id - FROM retrieval_serving_revision_stats - WHERE document_id = :document_id AND job_result_id = :job_result_id - """, - {"document_id": document_id, "job_result_id": revision["job_result_id"]}, - ) - namespace_stats = await ContractDatabase.fetch_one( - """ - SELECT generation - FROM retrieval_namespace_stats - WHERE user_id = :user_id AND namespace = :namespace - """, - {"user_id": "local-dev-user", "namespace": namespace}, - ) - assert remaining_revision_stats is None - assert namespace_stats is not None - assert int(namespace_stats["generation"]) >= 1 diff --git a/apps/api/tests/contract/test_retrieval_lazy_snapshot_quality_contract.py b/apps/api/tests/contract/test_retrieval_lazy_snapshot_quality_contract.py index 0752a3b9..ae6d78ed 100644 --- a/apps/api/tests/contract/test_retrieval_lazy_snapshot_quality_contract.py +++ b/apps/api/tests/contract/test_retrieval_lazy_snapshot_quality_contract.py @@ -2,8 +2,6 @@ from dataclasses import dataclass from collections.abc import Mapping, Sequence -from collections import Counter -import math from typing import Any from shared.services.retrieval.nav.nav_hierarchy import ProviderToolSpace @@ -22,11 +20,7 @@ ) from shared.services.retrieval.nav.knowhere_hybrid import ( ScoreUnitRow, - PersistedBm25Stats, - PersistedScoreCorpus, - PersistedScoreUnit, score_rows_hybrid_all, - score_persisted_corpus_many, score_unit_stream_hybrid_all, score_unit_stream_hybrid_many, ) @@ -310,7 +304,8 @@ def test_streaming_scorer_scores_multiple_queries_with_one_corpus_read() -> None ] queries: list[str] = ["alpha evidence", "beta evidence"] expected = { - query: score_unit_stream_hybrid_all(lambda: rows, query) for query in queries + query: score_unit_stream_hybrid_all(lambda: rows, query) + for query in queries } read_count: int = 0 @@ -323,87 +318,6 @@ def unit_factory() -> Sequence[ScoreUnitRow]: assert read_count == 1 -def test_persisted_score_projection_preserves_exact_eager_scores() -> None: - rows: list[ScoreUnitRow] = [ - { - "chunk_id": "unit-a", - "path_search_text": "root alpha", - "content_search_text": "common alpha alpha evidence", - "term_search_text": "common alpha alpha evidence root", - }, - { - "chunk_id": "unit-b", - "path_search_text": "root beta", - "content_search_text": "common beta evidence", - "term_search_text": "common beta evidence root", - }, - { - "chunk_id": "unit-c", - "path_search_text": "root common", - "content_search_text": "common evidence", - "term_search_text": "common evidence root", - }, - ] - queries = ["common alpha", "beta evidence"] - expected = { - query: score_unit_stream_hybrid_all(lambda: rows, query) for query in queries - } - query_tokens = {token for query in queries for token in query.split()} - - def build_stats(search_field: str) -> PersistedBm25Stats: - token_rows = [str(row[search_field]).split() for row in rows] - document_frequency = Counter( - token for tokens in token_rows for token in set(tokens) - ) - document_count = len(token_rows) - raw_idfs = [ - math.log(document_count - frequency + 0.5) - math.log(frequency + 0.5) - for frequency in document_frequency.values() - ] - return PersistedBm25Stats( - document_count=document_count, - total_length=sum(len(tokens) for tokens in token_rows), - document_frequency={ - token: document_frequency[token] for token in query_tokens - }, - average_idf=sum(raw_idfs) / len(raw_idfs), - ) - - corpus = PersistedScoreCorpus( - units=[ - PersistedScoreUnit( - unit_id=str(row["chunk_id"]), - path_length=len(str(row["path_search_text"]).split()), - content_length=len(str(row["content_search_text"]).split()), - path_frequencies={ - token: str(row["path_search_text"]).split().count(token) - for token in query_tokens - }, - content_frequencies={ - token: str(row["content_search_text"]).split().count(token) - for token in query_tokens - }, - term_scores=tuple( - 100.0 - if query in str(row["term_search_text"]) - else float( - sum( - token in str(row["term_search_text"]) - for token in query.split() - ) - ) - for query in queries - ), - ) - for row in rows - ], - path_stats=build_stats("path_search_text"), - content_stats=build_stats("content_search_text"), - ) - - assert score_persisted_corpus_many(corpus, queries) == expected - - def test_corpus_map_scores_multiple_queries_with_one_lazy_load() -> None: eager, lazy, store = _providers() queries: list[str] = ["alpha retrieval", "supporting image"] diff --git a/apps/api/tests/contract/test_retrieval_lazy_tree_contract.py b/apps/api/tests/contract/test_retrieval_lazy_tree_contract.py deleted file mode 100644 index 30994dad..00000000 --- a/apps/api/tests/contract/test_retrieval_lazy_tree_contract.py +++ /dev/null @@ -1,33 +0,0 @@ -"""Contract coverage for lazy MAP-NAV tree traversal.""" - -from __future__ import annotations - -from shared.services.retrieval.nav.nav_hierarchy import NodeMeta, ProviderToolSpace -from shared.services.retrieval.nav.nav_knowhere import KnowhereProvider, SectionRow -from shared.services.retrieval.nav.nav_map_scores import _walk_tree - - -class _MetadataForbiddenProvider(KnowhereProvider): - def node_meta(self, section_id: str) -> NodeMeta: - raise AssertionError(f"tree traversal materialized metadata for {section_id}") - - -def test_tree_walk_reads_children_and_titles_without_materializing_metadata() -> None: - provider = _MetadataForbiddenProvider( - doc_id="doc", - sections=[ - SectionRow("root", None, "Root", "Root", 0, "", 0), - SectionRow("child", "root", "Root / Child", "Child", 1, "", 1), - ], - units=(), - ) - - children, leaves, titles = _walk_tree( - ProviderToolSpace(provider), - "doc", - ["root"], - ) - - assert children == {"root": ["child"], "child": []} - assert leaves == {"child"} - assert titles == {"root": "Root", "child": "Child"} diff --git a/apps/api/tests/contract/test_retrieval_manifest_cache_contract.py b/apps/api/tests/contract/test_retrieval_manifest_cache_contract.py deleted file mode 100644 index e2f73e47..00000000 --- a/apps/api/tests/contract/test_retrieval_manifest_cache_contract.py +++ /dev/null @@ -1,56 +0,0 @@ -from __future__ import annotations - -import pytest -from sqlalchemy.exc import SQLAlchemyError - -from shared.services.retrieval.manifest_cache import ( - cache_manifest_payloads, - get_cached_manifest_payloads, -) -from shared.services.retrieval.search.scoped_corpus import count_manifest_chunks - - -class _SessionWithInfo: - def __init__(self) -> None: - self.info: dict[str, object] = {} - - -def test_manifest_payload_cache_is_scoped_to_revision_pin_set() -> None: - session = _SessionWithInfo() - revisions = {"doc-a": "result-a", "doc-b": "result-b"} - payloads = { - ("doc-a", "result-a"): {"chunks": [{"chunk_id": "chunk-a"}]}, - ("doc-b", "result-b"): {"chunks": [{"chunk_id": "chunk-b"}]}, - } - - cache_manifest_payloads(session, revisions=revisions, payloads=payloads) - - assert get_cached_manifest_payloads(session, revisions=revisions) == payloads - assert get_cached_manifest_payloads( - session, - revisions={"doc-a": "different-result"}, - ) is None - - -class _UnavailableManifestSession: - def __init__(self) -> None: - self.rollback_count = 0 - - async def execute(self, _statement: object) -> object: - raise SQLAlchemyError("serving manifest table is unavailable") - - async def rollback(self) -> None: - self.rollback_count += 1 - - -@pytest.mark.asyncio -async def test_manifest_count_falls_back_after_derived_table_error() -> None: - session = _UnavailableManifestSession() - - result = await count_manifest_chunks( - session, # type: ignore[arg-type] - revision_pins={"doc-a": "result-a"}, - ) - - assert result is None - assert session.rollback_count == 1 diff --git a/apps/api/tests/contract/test_retrieval_map_unit_index_contract.py b/apps/api/tests/contract/test_retrieval_map_unit_index_contract.py index eb124fb2..2f2027b0 100644 --- a/apps/api/tests/contract/test_retrieval_map_unit_index_contract.py +++ b/apps/api/tests/contract/test_retrieval_map_unit_index_contract.py @@ -5,7 +5,7 @@ from uuid import uuid4 from httpx import AsyncClient -from sqlalchemy import delete, select, text +from sqlalchemy import delete, select from shared.models.database.document import ( DocumentMapUnit, @@ -13,11 +13,9 @@ DocumentMapUnitToken, ) from shared.services.retrieval.nav.nav_hierarchy import ProviderToolSpace -from shared.services.retrieval.nav._compat import Chunk, EpisodeResult from shared.services.retrieval.nav.nav_map_scores import ( build_score_units, compute_corpus_map_and_unit_scores, - select_map_highlights, ) from shared.services.retrieval.nav.nav_knowhere import ( KnowhereProvider, @@ -31,7 +29,6 @@ from shared.services.retrieval.publication_content import ( replace_document_revision_content, ) -from shared.services.retrieval.nav_bridge import build_referenced_chunks from shared.services.retrieval.publication_models import DocumentPublicationScope from tests.support.contract_database import ContractDatabase from tests.support.retrieval_snapshot_support import contract_db_session @@ -182,7 +179,6 @@ async def test_published_map_units_preserve_scores_without_chunk_payload_reads( doc_ids=[document_id], query="common alpha", ) - expected_highlights = select_map_highlights(expected_scores[1], k=3) async with contract_db_session() as db: index = ( @@ -212,18 +208,6 @@ async def test_published_map_units_preserve_scores_without_chunk_payload_reads( ) ).scalars() ) - index_names = { - str(row[0]) - for row in ( - await db.execute( - text( - "SELECT indexname FROM pg_indexes " - "WHERE schemaname = current_schema() " - "AND tablename = 'document_map_unit_tokens'" - ) - ) - ).all() - } lazy_snapshot = await load_nav_snapshot( db, user_id=_USER_ID, @@ -236,7 +220,6 @@ async def test_published_map_units_preserve_scores_without_chunk_payload_reads( str(unit["chunk_id"]) for unit in expected_units ] assert persisted_tokens - assert "idx_document_map_unit_tokens_lookup" in index_names def reject_payload_read( _store: ReadOnlyChunkStore, @@ -274,11 +257,6 @@ def reject_payload_read( await db.execute( delete(DocumentMapUnitToken).where(DocumentMapUnitToken.id == token_id) ) - await db.execute( - delete(DocumentMapUnitIndex).where( - DocumentMapUnitIndex.document_id == document_id - ) - ) await db.commit() incomplete_snapshot = await load_nav_snapshot( db, @@ -296,137 +274,9 @@ def reject_payload_read( eager_snapshot.close() assert actual_scores == expected_scores - assert select_map_highlights(actual_scores[1], k=3) == expected_highlights assert fallback_scores == expected_scores -async def test_lazy_snapshot_defers_selected_asset_reference_metadata( - developer_api_client_factory: Callable[ - [], AbstractAsyncContextManager[AsyncClient] - ], - monkeypatch, -) -> None: - identifier = uuid4().hex[:8] - namespace = f"lazy-ref-{identifier}" - document_id = f"doc_ref_{identifier}" - job_id = f"job_ref_{identifier}" - job_result_id = f"result_ref_{identifier}" - async with developer_api_client_factory(): - await _seed_revision( - namespace=namespace, - document_id=document_id, - job_id=job_id, - job_result_id=job_result_id, - ) - scope = DocumentPublicationScope( - user_id=_USER_ID, - namespace=namespace, - document_id=document_id, - job_result_id=job_result_id, - source_file_name="refs.pdf", - ) - chunks = [ - { - "chunk_id": "body", - "type": "text", - "content": "body evidence", - "path": "refs.pdf/Root/Section/body", - "order": 1, - "metadata": {"connect_to": [{"target": "asset"}]}, - }, - { - "chunk_id": "asset", - "type": "image", - "content": "image description", - "path": "refs.pdf/Root/image", - "order": 2, - "file_path": "images/asset.png", - "metadata": {}, - }, - ] - async with contract_db_session() as db: - await db.run_sync( - lambda sync_db: replace_document_revision_content( - sync_db, - scope=scope, - chunks=chunks, - ) - ) - await db.commit() - - calls: list[tuple[str, str]] = [] - original = ReadOnlyChunkStore.load_chunk_reference_metadata - - def record_reference_load( - store: ReadOnlyChunkStore, - document: str, - chunk: str, - ) -> Mapping[str, object] | None: - calls.append((document, chunk)) - return original(store, document, chunk) - - monkeypatch.setattr( - ReadOnlyChunkStore, - "load_chunk_reference_metadata", - record_reference_load, - ) - async with contract_db_session() as db: - snapshot = await load_nav_snapshot( - db, - user_id=_USER_ID, - namespace=namespace, - lazy=True, - ) - - assert calls == [] - assert snapshot.chunk_ref_index[f"{document_id}:asset"]["file_path"] == ( - "images/asset.png" - ) - assert calls == [(document_id, "asset")] - episode = EpisodeResult( - representation="", - steps=[], - scored_chunks=[ - ( - Chunk( - node_id="asset", - doc_id=document_id, - text="image description", - line_ids=(2,), - section_id="root", - ), - 0.75, - ) - ], - kept_chunks=[ - Chunk( - node_id="asset", - doc_id=document_id, - text="image description", - line_ids=(2,), - section_id="root", - ) - ], - evidence_text="image description", - evidence_chars_actual=17, - retrieved_nodes=["asset"], - ) - refs, scores = build_referenced_chunks(episode, snapshot) - assert refs == [ - { - "chunk_id": "asset", - "document_id": document_id, - "chunk_type": "image", - "section_path": "Root / image", - "file_path": "images/asset.png", - "job_id": job_id, - "score": 0.75, - } - ] - assert scores == {"asset": 0.75} - snapshot.close() - - def test_incomplete_index_falls_back_for_duplicate_unit_ids() -> None: first_sections = [ SectionRow("root-a", None, "Root A", "Root A", 0, "", 0), diff --git a/apps/api/tests/contract/test_retrieval_mapnav_session_contract.py b/apps/api/tests/contract/test_retrieval_mapnav_session_contract.py index d267941b..a277c4c0 100644 --- a/apps/api/tests/contract/test_retrieval_mapnav_session_contract.py +++ b/apps/api/tests/contract/test_retrieval_mapnav_session_contract.py @@ -135,14 +135,12 @@ async def fake_resolve_workflow_references( namespace: str, refs: list[RouteRow], score_by_chunk_id: dict[str, float] | None = None, - revision_pins: dict[str, str] | None = None, ) -> ResolvedWorkflowReferences: assert db is fresh_db assert user_id == "contract-user" assert namespace == "contract-namespace" assert refs assert score_by_chunk_id is not None - assert revision_pins is None events.append("resolve_references") row = { "document_id": "doc_contract", @@ -161,13 +159,11 @@ async def fake_assemble_retrieval_results( exclude_document_ids: list[str], exclude_sections: list[dict[str, str]], allowed_chunk_types: set[str] | None, - revision_pins: dict[str, str] | None = None, ) -> list[RouteRow]: assert db is fresh_db assert exclude_document_ids == [] assert exclude_sections == [] assert allowed_chunk_types is None - assert revision_pins is None events.append("assemble_results") return rows diff --git a/apps/api/tests/contract/test_retrieval_relit_map_cache_contract.py b/apps/api/tests/contract/test_retrieval_relit_map_cache_contract.py deleted file mode 100644 index b8257c9c..00000000 --- a/apps/api/tests/contract/test_retrieval_relit_map_cache_contract.py +++ /dev/null @@ -1,27 +0,0 @@ -from __future__ import annotations - -from shared.services.retrieval.nav.nav_orchestrate import _relit_map -from shared.services.retrieval.nav.nav_types import NavState - - -def test_relit_map_reuses_same_query_within_episode(monkeypatch) -> None: - calls: list[str] = [] - - def fake_relight_map_for_query(*_args, **kwargs): - calls.append(str(kwargs["query"])) - return {"section": 1.0}, {"unit": 2.0}, ["section"] - - monkeypatch.setattr( - "shared.services.retrieval.nav.nav_map_scores.relight_map_for_query", - fake_relight_map_for_query, - ) - state = NavState(doc_id="", query="retrieval") - config = type("Config", (), {"collect_top_k": 6})() - - with _relit_map(None, state, config, query="retrieval"): - pass - with _relit_map(None, state, config, query="retrieval"): - pass - - assert calls == ["retrieval"] - assert state.relit_map_cache["retrieval"][0] == {"section": 1.0} diff --git a/apps/api/tests/contract/test_retrieval_revision_races_contract.py b/apps/api/tests/contract/test_retrieval_revision_races_contract.py deleted file mode 100644 index 3ebb0232..00000000 --- a/apps/api/tests/contract/test_retrieval_revision_races_contract.py +++ /dev/null @@ -1,108 +0,0 @@ -"""Deterministic contracts for revision and channel-session coherence.""" - -from __future__ import annotations - -import importlib -from collections.abc import AsyncGenerator -from contextlib import asynccontextmanager -from typing import Any, cast - -import pytest -from sqlalchemy.ext.asyncio import AsyncSession - -from shared.services.retrieval.execution.revision_pins import ( - RetrievalRevisionPins, - is_revision_generation_stable, -) -from shared.services.retrieval.search import discovery - - -@pytest.mark.asyncio -async def test_classic_channels_share_pins_but_use_distinct_sessions( - monkeypatch: pytest.MonkeyPatch, -) -> None: - pins = RetrievalRevisionPins( - revisions={"doc-1": "revision-1"}, - generation=7, - ) - sessions: list[object] = [] - observed: list[tuple[object, object]] = [] - - @asynccontextmanager - async def fake_context() -> AsyncGenerator[object, None]: - session = object() - sessions.append(session) - yield session - - async def fake_channel( - db: AsyncSession, - **kwargs: Any, - ) -> list[dict[str, Any]]: - observed.append((db, kwargs["revision_pins"])) - return [] - - # Import the active module explicitly. Some API contract fixtures reload - # shared.core.database between tests, leaving the package attribute pointed - # at an old module object; dotted-string patching can then miss the module - # used by discovery's lazy import. - database_module = importlib.import_module("shared.core.database") - monkeypatch.setattr(database_module, "get_db_context", fake_context) - monkeypatch.setattr(discovery, "path_channel", fake_channel) - monkeypatch.setattr(discovery, "content_channel", fake_channel) - monkeypatch.setattr(discovery, "term_channel", fake_channel) - - result = await discovery.bottom_discovery( - cast(AsyncSession, object()), - user_id="user-1", - namespace="namespace-1", - query="coherent query", - top_k=3, - exclude_document_ids=[], - exclude_sections=[], - revision_pins=pins, - ) - - assert result.status == "discovery_done" - assert len(sessions) == 3 - assert len({id(session) for session in sessions}) == 3 - assert len(observed) == 3 - assert {id(session) for session, _pins in observed} == { - id(session) for session in sessions - } - assert all(observed_pins is pins for _session, observed_pins in observed) - - -class _GenerationResult: - def __init__(self, value: int | None) -> None: - self._value = value - - def scalar_one_or_none(self) -> int | None: - return self._value - - -class _GenerationSession: - def __init__(self, values: list[int | None]) -> None: - self._values = iter(values) - - async def execute(self, _statement: object) -> _GenerationResult: - return _GenerationResult(next(self._values)) - - -@pytest.mark.asyncio -async def test_generation_change_is_detected_before_scoring() -> None: - pins = RetrievalRevisionPins(revisions={"doc-1": "revision-1"}, generation=7) - stable = await is_revision_generation_stable( - cast(AsyncSession, _GenerationSession([7])), - user_id="user-1", - namespace="namespace-1", - pins=pins, - ) - changed = await is_revision_generation_stable( - cast(AsyncSession, _GenerationSession([8])), - user_id="user-1", - namespace="namespace-1", - pins=pins, - ) - - assert stable is True - assert changed is False diff --git a/apps/api/tests/contract/test_retrieval_rrf_duplicate_contract.py b/apps/api/tests/contract/test_retrieval_rrf_duplicate_contract.py deleted file mode 100644 index e056af6c..00000000 --- a/apps/api/tests/contract/test_retrieval_rrf_duplicate_contract.py +++ /dev/null @@ -1,19 +0,0 @@ -"""Contracts for duplicate chunk handling in reciprocal-rank fusion.""" - -from __future__ import annotations - -from shared.services.retrieval.search.scoring import merge_channels_rrf - - -def test_rrf_counts_each_chunk_once_per_channel() -> None: - rows = [ - {"chunk_id": "shared", "document_id": "doc-a"}, - {"chunk_id": "shared", "document_id": "doc-b"}, - {"chunk_id": "other", "document_id": "doc-c"}, - ] - - result = merge_channels_rrf([rows], [1.0], top_k=3) - - assert [row["chunk_id"] for row in result] == ["shared", "other"] - assert result[0]["score"] == round(1.0 / 61, 6) - assert result[1]["score"] == round(1.0 / 62, 6) diff --git a/apps/api/tests/contract/test_retrieval_serving_manifest_contract.py b/apps/api/tests/contract/test_retrieval_serving_manifest_contract.py deleted file mode 100644 index 5bb68c06..00000000 --- a/apps/api/tests/contract/test_retrieval_serving_manifest_contract.py +++ /dev/null @@ -1,51 +0,0 @@ -"""Contract tests for serving-manifest integrity and version handling.""" - -from __future__ import annotations - -import pytest - -from shared.services.retrieval.serving_manifest import ( - SERVING_MANIFEST_FORMAT_VERSION, - decode_serving_manifest, - encode_serving_manifest, -) - - -def test_serving_manifest_round_trip_preserves_canonical_payload() -> None: - payload = { - "document_id": "doc_contract", - "job_result_id": "result_contract", - "sections": [{"section_id": "sec_1", "sort_order": 0}], - "chunks": [{"chunk_id": "chunk_1", "connect_to": []}], - } - - compressed, checksum, version = encode_serving_manifest(payload) - - assert version == SERVING_MANIFEST_FORMAT_VERSION - assert decode_serving_manifest( - compressed, - checksum=checksum, - format_version=version, - ) == payload - - -def test_serving_manifest_rejects_checksum_mismatch() -> None: - compressed, _, version = encode_serving_manifest({"document_id": "doc"}) - - with pytest.raises(ValueError, match="checksum mismatch"): - decode_serving_manifest( - compressed, - checksum="0" * 64, - format_version=version, - ) - - -def test_serving_manifest_rejects_unknown_version() -> None: - compressed, checksum, _ = encode_serving_manifest({"document_id": "doc"}) - - with pytest.raises(ValueError, match="unsupported serving manifest version"): - decode_serving_manifest( - compressed, - checksum=checksum, - format_version=SERVING_MANIFEST_FORMAT_VERSION + 1, - ) diff --git a/apps/api/tests/contract/test_retrieval_snapshot_consistency_contract.py b/apps/api/tests/contract/test_retrieval_snapshot_consistency_contract.py index bbe03e05..64aaf4c0 100644 --- a/apps/api/tests/contract/test_retrieval_snapshot_consistency_contract.py +++ b/apps/api/tests/contract/test_retrieval_snapshot_consistency_contract.py @@ -6,12 +6,8 @@ from uuid import uuid4 from httpx import AsyncClient -import pytest from sqlalchemy import Executable, Result -from shared.services.retrieval.execution.reference_resolver import ( - resolve_workflow_references, -) from shared.services.retrieval.nav_snapshot import SnapshotSession, load_nav_snapshot from tests.support.retrieval_snapshot_support import contract_db_session from tests.support.contract_database import ContractDatabase @@ -184,53 +180,3 @@ async def publish_new_revision() -> None: assert snapshot.chunk_ref_index[chunks[0].chunk_id]["section_path"] == ( "republished.pdf/old" ) - - -@pytest.mark.asyncio -async def test_reference_hydration_keeps_the_captured_revision_after_republish( - developer_api_client_factory: Callable[ - [], AbstractAsyncContextManager[AsyncClient] - ], -) -> None: - namespace = f"revision-hydration-race-{uuid4().hex[:8]}" - async with developer_api_client_factory(): - document_id, new_result_id = await _seed_republished_document(namespace) - revision_rows = await ContractDatabase.fetch_all( - """ - SELECT job_result_id, chunk_id - FROM document_chunks - WHERE document_id = :document_id - ORDER BY job_result_id - """, - {"document_id": document_id}, - ) - old_revision = next( - row for row in revision_rows if row["job_result_id"] != new_result_id - ) - await ContractDatabase.execute( - """ - UPDATE documents - SET current_job_result_id = :new_result_id - WHERE document_id = :document_id - """, - {"new_result_id": new_result_id, "document_id": document_id}, - ) - - async with contract_db_session() as db: - resolved = await resolve_workflow_references( - db=db, - user_id=_USER_ID, - namespace=namespace, - refs=[ - { - "document_id": document_id, - "chunk_id": old_revision["chunk_id"], - } - ], - revision_pins={document_id: old_revision["job_result_id"]}, - ) - - assert [row["content"] for row in resolved.rows] == ["old content"] - assert [row["job_result_id"] for row in resolved.rows] == [ - old_revision["job_result_id"] - ] diff --git a/apps/api/tests/contract/test_retrieval_term_score_contract.py b/apps/api/tests/contract/test_retrieval_term_score_contract.py deleted file mode 100644 index 425f4de0..00000000 --- a/apps/api/tests/contract/test_retrieval_term_score_contract.py +++ /dev/null @@ -1,75 +0,0 @@ -"""Contracts for persisted map-unit term scoring.""" - -from __future__ import annotations - -from shared.services.retrieval.nav.nav_knowhere import ReadOnlyChunkStore - - -class _FakeCursor: - def __init__(self, rows: list[tuple[str, str]]) -> None: - self.rows = rows - self.statement = "" - self.parameters: object = None - - def execute(self, statement: str, parameters: object) -> None: - self.statement = statement - self.parameters = parameters - - def fetchall(self) -> list[tuple[str, str]]: - return self.rows - - -def _store() -> ReadOnlyChunkStore: - return ReadOnlyChunkStore.__new__(ReadOnlyChunkStore) - - -def test_term_scores_keep_literal_substring_and_token_hit_semantics() -> None: - cursor = _FakeCursor( - [ - ("unit-full", "prefix alpha beta suffix"), - ("unit-token", "prefix alpha gamma suffix"), - ] - ) - queries = ["alpha beta", "", "alpha"] - query_tokens = { - "alpha beta": ["alpha", "beta"], - "": [], - "alpha": ["alpha"], - } - - scores = _store()._load_term_scores( - cursor, # type: ignore[arg-type] - map_unit_ids=["unit-full", "unit-token", "unit-miss"], - queries=queries, - query_tokens_by_query=query_tokens, - ) - - assert scores == { - "unit-full": (100.0, 0.0, 100.0), - "unit-token": (1.0, 0.0, 100.0), - } - assert "LIKE ANY" in cursor.statement - assert "POSITION" not in cursor.statement - parameters = cursor.parameters - assert isinstance(parameters, tuple) - assert parameters[0] == ["unit-full", "unit-token", "unit-miss"] - assert parameters[1] == ["%alpha beta%", "%alpha%", "%beta%"] - - -def test_long_query_uses_constant_shape_candidate_sql() -> None: - cursor = _FakeCursor([]) - tokens = [f"token-{index}" for index in range(300)] - query = " ".join(tokens) - - _store()._load_term_scores( - cursor, # type: ignore[arg-type] - map_unit_ids=["unit-1"], - queries=[query], - query_tokens_by_query={query: tokens}, - ) - - assert cursor.statement.count("POSITION") == 0 - assert "LIKE ANY" not in cursor.statement - parameters = cursor.parameters - assert isinstance(parameters, tuple) - assert parameters == (["unit-1"],) diff --git a/docs/adr/0005-stream-retrieval-progress-over-sse.md b/docs/adr/0005-stream-retrieval-progress-over-sse.md deleted file mode 100644 index dfaafbd8..00000000 --- a/docs/adr/0005-stream-retrieval-progress-over-sse.md +++ /dev/null @@ -1,47 +0,0 @@ -# Stream Retrieval Progress Over SSE - -## Status - -Accepted - -## Context - -Online Brain users currently wait for the complete Retrieval response while -map-nav planning, searching, and source review run. Knowhere owns retrieval and -evidence, while answer generation belongs to downstream clients. A streaming -contract must improve perceived latency without exposing chain-of-thought or -making partial citations authoritative. - -## Decision - -Add a v2-only `POST /v2/retrieval/query/stream` endpoint using Server-Sent -Events. The endpoint accepts the full `RetrievalQueryRequest` shape and is -live-only: disconnecting cooperatively cancels the retrieval run, and retries -start a new run. In-progress events use a fixed, route-aware phase vocabulary -(`started`, `planning`, `searching`, `reviewing_sources`, `finalizing`) and may -include only safe aggregate counts. The stream terminates with a versioned -envelope containing either the existing retrieval response as the authoritative -result or a typed, user-safe failure (`failed`, `cancelled`, or `no_results`). - -The synchronous map-nav engine remains intact and publishes sanitized progress -through an optional step callback bridged to the SSE route by an async queue. -Answer-token streaming remains downstream. The endpoint sends heartbeats, -disables proxy buffering, and uses per-connection event IDs without promising -replay in the first version. - -Retrieval duration is measured from the execution plan's start through final -public-result assembly. The same duration definition is used for persisted -`retrieval_runs.latency_ms` and latency aggregates; cache hits are recorded too. -Trace persistence must receive the execution start timestamp rather than -starting its own timer after navigation and hydration. Timing and terminal -outcomes must cover classic, map-nav, small-corpus, cache-hit, failed, and -cancelled routes. `retrieval_runs` is the ledger for all of those routes and -stores explicit route, cache, latency, and terminal-status fields. - -## Consequences - -The existing JSON retrieval endpoint remains backward compatible, while SDKs -and Online Brain clients need a new streaming adapter and UI state model. A -future resumable stream would require durable event replay and is deliberately -out of scope. Partial evidence and provisional citations are also deferred -until their grounding and revision semantics are defined. diff --git a/docs/adr/0006-atomically-publish-retrieval-serving-index.md b/docs/adr/0006-atomically-publish-retrieval-serving-index.md deleted file mode 100644 index 6f0b0800..00000000 --- a/docs/adr/0006-atomically-publish-retrieval-serving-index.md +++ /dev/null @@ -1,6 +0,0 @@ -# Atomically publish the retrieval-serving index - -- Status: Accepted -- Context: Retrieval will use a persistent derived serving index to avoid rebuilding a large namespace on every first request. A document revision without a complete index would have unpredictable latency and could produce inconsistent scoring metadata. -- Decision: Build the serving manifest and scoring statistics in the same database transaction as the document revision. Write the completeness marker last. If serving-index construction fails, roll back the publication and retry the job; do not expose an active revision with a partial serving index. -- Consequences: Active revisions have a simple completeness invariant and predictable first-request behavior. Publication takes more work and storage, and an index failure can delay publication, but retrieval can retain a guarded legacy fallback for migrations or already-existing incomplete revisions. diff --git a/docs/adr/0007-use-coherent-retrieval-serving-generations.md b/docs/adr/0007-use-coherent-retrieval-serving-generations.md deleted file mode 100644 index 89d2ca05..00000000 --- a/docs/adr/0007-use-coherent-retrieval-serving-generations.md +++ /dev/null @@ -1,6 +0,0 @@ -# Use coherent retrieval-serving generations - -- Status: Accepted -- Context: A namespace can contain many active document revisions, and publication can replace them while a retrieval request is loading serving metadata and scoring statistics. -- Decision: Assign each namespace a serving generation. Retrieval captures one generation and verifies it across serving reads; if it changes, retry once and use the exact legacy path if consistency cannot be established. -- Consequences: Retrieval never combines incompatible revision metadata and scoring statistics. Publication and retrieval need a small amount of generation bookkeeping, and rare concurrent updates may cause a retry or slower fallback. diff --git a/docs/adr/0008-use-a-maintenance-window-for-serving-index-rollout.md b/docs/adr/0008-use-a-maintenance-window-for-serving-index-rollout.md deleted file mode 100644 index e7db06bf..00000000 --- a/docs/adr/0008-use-a-maintenance-window-for-serving-index-rollout.md +++ /dev/null @@ -1,6 +0,0 @@ -# Roll out the serving index online - -- Status: Accepted -- Context: The server and document publication must remain available while serving-index schema changes and backfill are introduced. -- Decision: Use additive online migrations and bounded idempotent backfill. The retrieval reader automatically uses the serving index only when a revision is complete and consistent; otherwise it uses the exact legacy reader. New publication continues online and builds serving data atomically before activating a revision. Do not expose partial serving data. -- Consequences: There is no planned retrieval or publication downtime. Backfill consumes bounded database resources and some revisions remain on the slower legacy path until complete. Generation checks and stale-revision guards are required while backfill and publication run concurrently. diff --git a/docs/adr/README.md b/docs/adr/README.md index 00d7b9da..68e15b22 100644 --- a/docs/adr/README.md +++ b/docs/adr/README.md @@ -19,7 +19,4 @@ Use this shape: | [0002](0002-use-typed-workflow-outcomes.md) | Use typed workflow outcomes | | [0003](0003-keep-retrieval-workflow-policy-explicit.md) | Keep retrieval workflow policy explicit | | [0004](0004-anonymous-self-hosted-telemetry.md) | Anonymous self-hosted telemetry | -| [0005](0005-stream-retrieval-progress-over-sse.md) | Stream retrieval progress over SSE | -| [0006](0006-atomically-publish-retrieval-serving-index.md) | Atomically publish the retrieval-serving index | -| [0007](0007-use-coherent-retrieval-serving-generations.md) | Use coherent retrieval-serving generations | -| [0008](0008-use-a-maintenance-window-for-serving-index-rollout.md) | Roll out the serving index online | +| \ No newline at end of file diff --git a/docs/design/retrieval-serving-index-plan.md b/docs/design/retrieval-serving-index-plan.md deleted file mode 100644 index 5e66f6d6..00000000 --- a/docs/design/retrieval-serving-index-plan.md +++ /dev/null @@ -1,425 +0,0 @@ -# Retrieval-serving index: online performance plan - -**Status:** Proposed for review -**Reviewed against:** current `knowhere` retrieval and publication code, 2026-08-29 -**Scope:** first-request retrieval performance; LLM planner/harvest/control time is excluded - -## 1. Goal and non-goals - -The target is predictable, bounded **Retrieval Non-LLM Work** on the current -production-sized namespace (about 643 active documents, 50k sections, and 60k -chunks). The measured baseline is roughly 160-170 seconds before map-nav's LLM -episode begins. We do not require a one-second absolute target in this phase; -we require that the request path avoid repeated full-corpus work and have a -clear linear/bounded complexity profile. - -The target includes: - -- snapshot or serving-index loading; -- classic or map-nav lexical scoring; -- ranking; -- selected-result hydration; -- citation and asset-reference assembly. - -It excludes planner, harvest, control, and answer-generation model time. Those -stages must continue to have separate timings. - -The plan does not change prompts, models, tokenization, BM25 formulas, RRF -weights, cache semantics, citation rules, or public HTTP response shapes. -It does not add LLM-response caching, process-wide serving caches, or startup -prewarming. The acceptance benchmark is a cold, uncached retrieval request, -and planner/harvest/control model time remains a separately reported -dependency. Episode-local reuse is allowed only while one request is active. - -## 2. Verified current behavior - -The current code has two retrieval routes in -`shared/services/retrieval/execution/routes.py`: - -- `use_agentic=false` runs `bottom_discovery()` and then ranking/hydration. -- The default map-nav route calls `load_nav_snapshot(..., lazy=True)`, runs the - synchronous navigation episode, then opens a fresh database context for - reference resolution and result assembly. - -`load_nav_snapshot()` currently: - -1. Reads active documents and their current job-result IDs. -2. Loads all matching sections into memory. -3. Loads all chunk identities and `connect_to` metadata into memory. -4. Uses a lazy store for selected chunk content and asset paths. - -The current persisted map scorer in -`shared/services/retrieval/nav/nav_knowhere.py` still loads all eligible map -units, then reads query frequencies and term scores from the persisted tables. -The scorer itself is fast; the broad database projection is not. - -`bottom_discovery()` currently executes path, content, and term channels -sequentially. `term_channel()` uses substring predicates over lowercased text, -without a trigram index. - -Publication currently writes sections/chunks and then calls -`replace_document_map_units()` in the same SQLAlchemy transaction. The existing -`document_map_unit_indexes` row is a per-document-revision completeness marker. -The existing backfill script only rebuilds that map-unit index; it does not yet -build the proposed serving manifest or namespace statistics. - -The working tree already contains an uncommitted keyset-pagination change and -the `3f4a5b6c7d8e` section-order migration. Keep those changes separate from -the serving-index implementation and from unrelated documentation edits. - -## 3. Architecture decision - -Use PostgreSQL as the source of truth and add a persistent, revision-pinned -serving read model. Do not add OpenSearch, Elasticsearch, Tantivy, or another -search service in this version. PostgreSQL's existing FTS plus `pg_trgm` keeps -the current scoring and tie-breaking behavior directly testable. - -### 3.1 Revision serving manifest - -Add one compressed manifest row per `(document_id, job_result_id)`. The payload -contains ordered metadata only: - -- document, revision, source filename, and job identity; -- section IDs, parent IDs, paths, titles, levels, summaries, and sort order; -- chunk IDs, section IDs, types, sort order, and `connect_to` target IDs; -- map-unit row IDs, unit IDs, unit kinds, token lengths, and sort order; -- root-asset IDs and remounted asset owners. - -It must not contain full chunk content or asset file paths. Those remain in the -canonical tables and are loaded lazily for selected evidence. - -Store the payload as canonical JSON compressed with the standard-library zlib -implementation, with a format version and checksum. The serving loader must -reject an unknown version, checksum mismatch, or incomplete payload. - -### 3.2 Serving generations and statistics - -Add namespace-scoped generation metadata and persistent scoring statistics: - -- `retrieval_namespace_generations`: current generation per user/namespace; -- `retrieval_serving_revision_stats`: compressed per-revision contributions; -- `retrieval_namespace_stats`: aggregate unit counts, total lengths, vocabulary - frequency histograms, and generation; -- `retrieval_namespace_token_stats`: queryable document frequency per channel - and token hash. - -The generation is a consistency marker, not a replacement for revision IDs. -Retrieval captures active revision IDs and one generation. If generation changes -while the snapshot is being captured or before scoring starts, retry once; if -consistency still cannot be proven, use the exact legacy reader. - -Every retrieval route must carry that capture as an immutable revision pin set: -`{document_id, namespace, job_result_id}` plus the captured generation. The pin -set is the source of truth for the request after capture. Downstream queries -must constrain sections, map units, chunks, connected assets, ranking lookups, -reference resolution, and result assembly by the pinned `job_result_id`; they -must not re-join through the live `Document.current_job_result_id`. A generation -change after scoring has started must never cause a mix of old and new rows: -finish against the captured pins (or return an exact legacy result), and only -retry before work that depends on the snapshot begins. -Snapshot admission is therefore decided at capture time. If a later archive -must suppress an in-flight result, discard/retry the whole request; do not -replace its pinned revision with the document's new current revision. - -Cache hits occur before route execution, so every operation that changes the -serving generation (publication, republish, archive, or namespace move) must -also advance the namespace retrieval-cache version, or store the generation in -the cache entry and reject mismatches. This keeps cached responses from -outliving the generation they represent without changing the public response -shape. - -### 3.3 Indexes - -Additive migrations should provide: - -- a token-first covering index for map-unit token candidates. The existing - `idx_document_map_unit_tokens_lookup` is token-first but does not cover the - selected columns; the pending `2e3f4a5b6c7d` migration adds a unit-first - covering index for a different access pattern, so the serving reader may - need one additional token-first covering index; -- a revision/section lookup index for map units; -- a trigram GIN index on `document_map_units.term_search_text_lower`; -- a generated lowercased term field and trigram GIN index for - `document_chunks.term_search_text`; -- the existing chunk ordering index plus the pending token-covering and - section-order migrations (`2e3f4a5b6c7d` and `3f4a5b6c7d8e`). - -Enable PostgreSQL's built-in `pg_trgm` extension. No separate search service is -required. - -## 4. Retrieval changes - -Capture the revision pin set and generation at retrieval-route entry, before -the small-corpus count or route selection. Pass that capture into whichever -route is selected; a route-local capture is allowed only when it is performed -as the same snapshot transaction. This prevents the count/load pair in the -small-corpus optimization from straddling a publication. - -### 4.1 Fast map-nav snapshot - -Extend `load_nav_snapshot()` to try the serving manifest first: - -1. Capture active documents, current revision IDs, and namespace generation in a - short read-only transaction, returning the immutable revision pin set with - the snapshot. -2. Fetch one manifest row per active revision. -3. Decode and validate manifests. -4. Apply the existing document and section exclusion predicates. -5. Build the current `LazyKnowhereProvider` and `LazyChunkRefIndex` from the - decoded metadata. -6. Pin the lazy chunk store to the captured revision IDs. -7. Verify generation stability before returning the snapshot. - -For an unfiltered map-nav request, route selection may count chunks from these -same validated manifests instead of scanning `document_chunks`; filtered and -classic requests retain the exact SQL counter. The count shortcut must fall -back when any manifest is missing or invalid. - -If any manifest is absent or invalid, use the existing legacy snapshot loader. -The legacy loader must return the same revision pin set and apply the same -downstream predicates. This fallback is automatic and exact; it is not a public -feature flag. - -The serving path must preserve current ordering, duplicate bare/document-scoped -reference keys, root-asset remounting, section filtering, and revision pinning. -Keep the pin set available through the complete map-nav request. After the LLM -episode, either materialize selected rows (including connected assets) from the -pinned lazy store before closing it, or pass the pin set to -`resolve_workflow_references()` and `assemble_retrieval_results()`. Their SQL -must select the captured `(document_id, job_result_id)` rows directly, so a -republish during the episode cannot make final citations resolve against the -new current revision. - -### 4.2 Exact persisted map scoring - -Keep `PersistedScoreCorpus` and the existing scorer unchanged wherever possible. -Replace only the data-loading strategy: - -- Prepare the immutable, revision-pinned unit projection and namespace scoring - statistics once per retrieval episode. Checklist relight waves must reuse that - projection; they may fetch or compute only query-specific postings/scores. - A wave must not issue another full-namespace unit/statistics load for the same - pin set. Instrument the loader call count and include it in the benchmark - report so repeated projection loads cannot hide behind separate wave timings. -- use manifest map-unit metadata to represent all units, including zero-score - units; the serving reader should not re-query `document_map_units` for these - IDs, lengths, or section membership; -- query token postings only for tokens in the request; -- filter postings by captured revisions and allowed sections; -- discover term-channel candidates through the trigram index while retaining the - current exact substring/token-hit scoring. The candidate predicate must use - the trigram-indexed `LIKE '%term%'` form (with the same lowercased query and - tokens), then apply the existing exact score expression; do not scan every - unit's term text in Python. -- obtain normal-corpus lengths, document frequencies, and IDF-flooring data from - persistent statistics; -- preserve the existing lexical sort key and RRF ranking. - -Queries with document or section exclusions must remain exact. If adjusted -statistics cannot be calculated with certainty, use the legacy scorer for that -request rather than approximating them. - -### 4.3 Classic retrieval — one pinned revision snapshot - -Keep the existing channel implementations and result projection. In -`bottom_discovery()`: - -- capture one revision pin set and generation before starting any channel; -- execute enabled channels concurrently; -- give each channel its own short-lived database session; -- pass the same pin set to every channel and constrain every channel query to - those revisions; -- preserve channel limits, Python BM25, term scoring, RRF merge, score - normalization, and all-or-error behavior; -- use the new trigram index only to narrow term candidates. - -Do not share one `AsyncSession` across concurrent channel tasks. -Ranking lookups, duplicate suppression, connected-target hydration, and final -assembly must receive the same pin set as discovery. The classic result must -therefore contain rows from one revision per document even if publication -replaces a document while one of the channel sessions is running. The -small-corpus optimization must use this same captured snapshot/pin contract (or -the exact legacy equivalent), rather than loading all rows through live current -revision joins. - -## 5. Publication and lifecycle behavior - -Refactor publication so the same build pass produces: - -- canonical sections/chunks; -- existing map-unit rows and completeness marker; -- the revision serving manifest; -- revision statistics and namespace-statistics deltas. - -All of this happens synchronously in the existing publication transaction. The -completeness marker and generation update are written last. If serving-index -construction fails, the publication transaction rolls back. - -New publication remains online during backfill. First publication, republish, -archive, and namespace-move paths must update statistics under the same -namespace generation row lock. The lock covers the active revision set, -namespace membership, revision contributions, and the generation increment, so -readers and writers have one lifecycle ordering. - -Backfill must rebuild the complete derived serving state (map units, manifest, -revision contribution, and namespace-statistics delta), not only the existing -map-unit index. It must select only documents with `status = 'active'`, a non-null -`current_job_result_id`, and the intended user/namespace. Immediately before -writing a contribution, it must hold the namespace lock and re-read the -document, then require all of the following to remain true: active status, -unchanged user/namespace, and `current_job_result_id` equal to the captured -revision. Otherwise it skips that revision without adding statistics. This -active-status predicate is required in the selector as well as in the -commit-time guard; update `apps/api/scripts/backfill_map_unit_indexes.py` to -include it in the existing selector. - -Archiving must atomically remove or invalidate that document revision's serving -statistics contribution while holding the same lock and advance the namespace -generation. `archive` currently changes `status` without clearing -`current_job_result_id`, so checking the revision pointer alone is insufficient -and would allow an in-flight backfill to re-add an archived revision. - -## 6. Online rollout - -There is no planned downtime, runtime feature flag, or production shadow-read -mode. - -1. Deploy additive schema/index migrations, beginning with the pending - `2e3f4a5b6c7d` and `3f4a5b6c7d8e` migrations. -2. Deploy code that automatically uses the serving reader only for complete, - valid revisions and otherwise uses the legacy reader. -3. Run an explicit, idempotent, bounded backfill for existing active revisions. -4. Keep retrieval and publication online while backfill runs. -5. Verify manifest checksums, revision coverage, namespace statistics, and - generation consistency. -6. Run strict legacy-versus-serving differential checks before considering the - rollout complete. - -If backfill is incomplete, affected revisions continue on the exact legacy -path. If online serving data is corrupted, reject it, alert, repair it with the -backfill/rebuild script, and do not serve partial data. - -Before enabling the serving reader for a namespace, record an inventory of -active `(document_id, current_job_result_id)` pairs, manifest completeness, and -expected per-revision and aggregate unit counts. After backfill, reconcile those -same values and verify that every aggregate includes only active, namespace- -member revisions. Abort the fast-path rollout on any missing/extra revision, -checksum failure, count mismatch, archived contribution, or generation -discontinuity. - -Any migration, backfill, or other database write—especially against -production—requires explicit approval immediately before execution. Read-only -inspection and benchmarking may proceed without that approval. - -## 6.1 DevOps operations runbook - -DevOps owns the production rollout mechanics; application code does not run a -startup backfill or create serving tables implicitly. Execute the following in -order: - -1. **Preflight (read-only):** confirm the target account, database, migration - head, available disk, connection headroom, and a recent rollback point. Record - the active `(document_id, current_job_result_id)` inventory for each namespace - that will be backfilled. -2. **Schema rollout:** with explicit approval immediately beforehand, apply the - additive migrations in dependency order: `2e3f4a5b6c7d`, - `3f4a5b6c7d8e`, `4a5b6c7d8e9f`, then `5b6c7d8e9f0a`. Run the trigram-index - migration during a low-traffic window and monitor for blocking locks. -3. **Application rollout:** deploy the API and worker versions containing the - serving reader and atomic publication changes. Verify health, error rate, and - legacy fallback before starting the backfill. -4. **Bounded backfill:** with separate approval, run - `uv run python apps/api/scripts/backfill_map_unit_indexes.py --apply` from a - controlled operator environment. Limit concurrency, pause on database - saturation, and resume safely; the operation is idempotent and stale or - inactive revisions must be skipped. -5. **Reconciliation:** compare the preflight inventory with serving manifests, - checksums, per-revision unit counts, namespace aggregates, and generation - values. Confirm aggregates contain only active documents still belonging to - the namespace. Investigate every missing, extra, stale, or invalid revision. -6. **Acceptance:** run the production read-only legacy-versus-serving - differential harness and record latency, selected IDs, order, scores, - citations, section paths, asset references, and fallback behavior. Declare - the rollout complete only after zero semantic mismatches. - -If migration or backfill must be stopped, leave the serving tables in place and -stop the operator job. The reader will continue using the exact legacy path for -incomplete revisions. Roll back application code first if necessary; do not -drop serving tables or indexes as an emergency rollback action. Repair a failed -revision by rerunning the bounded backfill after the cause is understood. - -## 7. Contract tests and benchmarks - -Use contract tests only. Add contracts for: - -- manifest round-trip, checksum, version, and revision pinning; -- eager versus serving snapshot equivalence; -- exclusions, duplicate chunk IDs, document-scoped references, and root assets; -- exact Latin/CJK, empty, no-hit, phrase, token-only, and negative-IDF cases; -- incomplete serving data falling back to legacy; -- publication replacement, archive deltas, concurrent generation changes, and - stale backfill protection; -- cache invalidation racing with a generation change, proving an old cached - response is not returned for a newer serving generation; -- map-nav republish during the LLM episode, proving final hydration and - connected-asset resolution stay on the captured revisions; -- classic publication replacement during concurrent channels, ranking, and - final assembly, proving every returned row shares the channel's pin set; -- archive/backfill races proving archived or namespace-moved revisions never - contribute to serving statistics; -- concurrent classic channels preserving IDs, order, scores, citations, and - fallback behavior. - -Race tests must use barriers or an equivalent deterministic hook to force a -republish during the map-nav episode, a publication between classic channel -sessions, an archive during backfill, and a namespace move during backfill. -Each test must assert both the returned evidence and the persisted statistics, -not merely that the request completed. - -The validation harness must also inspect the generated SQL/query plans (or an -equivalent query-boundary assertion) to prove pinned reads do not use live -`Document.current_job_result_id` joins. Run cache-version/generation races and -verify an old cached response is rejected after a lifecycle change. - -Run a differential harness against the production read-only database and -compare selected IDs, ordering, rounded scores, citations, section paths, -asset references, and fallback behavior. - -Benchmark fresh processes and uncached queries. Report separately: - -- serving capture/decode; -- map index projection and scoring; -- episode-local projection reuse (number of full projection loads and per-wave - query-only scoring time); -- classic discovery; -- ranking; -- hydration/assembly; -- total Retrieval Non-LLM Work; -- planner/harvest/control LLM time. - -The complexity check is explicit: one request may perform one full pinned -snapshot/projection pass, relight work should scale with query postings rather -than reloading the corpus, and hydration should scale with selected evidence -(`top_k`/references), not namespace size. Navigation wave count must not -multiply full-corpus database loads. - -Record peak resident memory for a fresh worker during the same benchmark and -repeat it with the expected concurrent-request level. Memory is reported as an -operational trade-off rather than a latency acceptance gate for this phase; -before production rollout, any episode-local or process-local reuse still needs -an explicit byte/item budget and an agreed worker ceiling. - -The fast path is accepted only after zero semantic mismatches and evidence that -the cold request performs one bounded serving projection, does not repeat -full-corpus loads per navigation wave, and meets an agreed latency budget for -the current production-sized corpus. - -## 8. Main tradeoffs and risks - -- Publication becomes slower and uses more storage because derived data is built - synchronously. -- Existing documents need an explicit backfill before they use the fast path. -- A serving-index inconsistency causes a slower legacy request, not approximate - evidence. -- PostgreSQL remains a scaling dependency; a future search-engine migration - would require a new semantic-parity review. diff --git a/docs/design/retrieval-streaming-sse.md b/docs/design/retrieval-streaming-sse.md deleted file mode 100644 index b0b5a73e..00000000 --- a/docs/design/retrieval-streaming-sse.md +++ /dev/null @@ -1,207 +0,0 @@ -# Retrieval Streaming over SSE - -**Status:** Accepted design -**Related issue:** [#330](https://github.com/Ontos-AI/knowhere/issues/330) -**Related ADR:** [0005](../adr/0005-stream-retrieval-progress-over-sse.md) - -## Purpose - -Online Brain users currently wait for a complete retrieval response while -map-nav planning, searching, source review, and final hydration run. This -design makes that work visible without exposing chain-of-thought or changing -who owns answer generation. - -## Ownership boundary - -Knowhere owns retrieval, safe progress, evidence, and authoritative citations. -The downstream Online Brain client owns answer synthesis and answer-token -streaming. Knowhere does not begin generating the final answer as part of this -feature. - -## Public API - -Add `POST /v2/retrieval/query/stream` with `Content-Type: text/event-stream`. -The request body is the full existing -`RetrievalQueryRequest`; streaming changes delivery, not retrieval semantics. -The existing JSON endpoints remain unchanged as the fallback. - -The stream is live-only. A disconnect cooperatively cancels the run; retrying -starts a new run. Event IDs are monotonic per connection but do not imply -replay or resumability. - -The route remains behind the normal authenticated-user dependency and v2 -route-admission policy. The guest-key allowlist, system-limit configuration, -OpenAPI registration, and any CORS policy must explicitly include the stream -route. v2 BYOK `llm_config` is accepted exactly as on the JSON route and is -never copied into an event payload or log message. - -## Event contract - -Each SSE frame has an event name (`progress`, `heartbeat`, or `terminal`) and a -JSON payload with `schema_version`, `stream_id`, `sequence`, and server -`elapsed_ms`. Progress payloads use the fixed, route-aware phases: - -```text -started → planning? → searching → reviewing_sources → finalizing -``` - -Classic and small-corpus routes omit `planning`; they must not emit phases that -did not occur. In-progress events contain only safe aggregate counts such as -`candidate_source_count` and `reviewed_source_count`. They do not contain -document names, chunk content, query rewrites, raw planner output, citations, -or chain-of-thought. - -The terminal event is a versioned envelope containing the existing retrieval -response as the authoritative result: - -```json -{ - "schema_version": 1, - "stream_id": "rst_...", - "sequence": 7, - "elapsed_ms": 2410, - "status": "completed", - "response": { "namespace": "default", "query": "...", "router_used": "mapnav", "evidence_text": "...", "referenced_chunks": [], "results": [] } -} -``` - -Failure terminals use `failed`, `cancelled`, or `no_results`. Failures expose a -stable user-safe `code` and `message`; detailed provider, database, -authentication, and planner errors remain in server logs. - -Wire requirements are part of the contract: frames use UTF-8 SSE `id`, -`event`, and one `data` line followed by a blank line; a heartbeat is an SSE -comment or named heartbeat event and carries no retrieval data; exactly one -terminal event is sent before the connection closes; no `retry` directive is -promised because the stream is not resumable. Progress events may be -coalesced when a bounded queue is full, but terminal events must never be -dropped. - -Cache hits still emit `started` and a terminal `completed` event, with a safe -`cache_hit: true` indicator. They omit phases for work that did not run. -`no_results` is reserved for a successful retrieval execution that produces no -evidence; provider, timeout, validation, and internal failures remain -`failed`. - -HTTP authentication, request validation, and route-admission failures happen -before the stream starts and use ordinary HTTP error responses. Once a `200` -SSE response has been opened, execution failures must be represented by a -terminal SSE event because the HTTP status can no longer be changed. - -## Internal implementation seam - -Keep the synchronous map-nav implementation. Add an optional callback that -receives a sanitized progress projection after each completed planner, -search, or review step. The SSE route bridges this callback to an -`asyncio.Queue` using a thread-safe loop handoff while retrieval continues in -its existing worker thread. - -The queue is bounded and owns cleanup of the worker task, callback, heartbeat -task, and database session. The callback must not touch an `AsyncSession` from -the worker thread. The route polls request disconnect state and propagates a -cancellation token; all producer tasks are joined or cancelled in a `finally` -block so abandoned streams cannot leak threads or connections. - -Add cooperative cancellation checks between steps. An in-flight synchronous -provider call may finish before cancellation takes effect. Cancelled runs do -not perform final hydration when cancellation is observed in time. - -Phase ownership is explicit: the route emits `started`; the map-nav adapter -emits `planning` before `plan_query` and `searching` before navigation or -classic discovery; the route emits `reviewing_sources` after retrieval -selection and before reference hydration; and it emits `finalizing` before -public projection. Counts are sourced from existing snapshot, reference, and -assembled-result counts and are omitted when not yet known. - -## Correct duration accounting - -The execution plan already starts a monotonic timer before cache lookup and -logs elapsed time after the route. However, `TraceRecorder` currently starts a -second timer in its constructor, and the map-nav route constructs it only -after navigation, reference resolution, and result assembly. Persisted -`retrieval_runs.latency_ms` and its aggregates therefore under-report retrieval -latency and mostly measure trace flush time. - -The canonical `Retrieval Duration` is the execution-plan timer from retrieval -start through final public-result assembly. It includes cache lookup and -applies to cache hits and misses. It excludes authentication, network/SSE -delivery, and downstream answer generation. - -Required changes: - -- pass the execution start timestamp into `TraceRecorder`; -- set `retrieval_runs.latency_ms` from that timestamp; -- record cache-hit runs with the same definition; -- ensure classic, map-nav, small-corpus, cache-hit, failed, and cancelled - retrievals all have an explicit timing/observability outcome; -- expose separate `time_to_first_event_ms`, `retrieval_latency_ms`, and - downstream `time_to_first_token_ms` measurements; -- retain per-step `elapsed_ms` as step latency, not total request latency. - -`retrieval_runs` is the ledger for every retrieval execution, not only -map-nav. Each row records the route type, `agentic_enabled`, `cache_hit`, -canonical latency, and terminal status for classic, map-nav, small-corpus, -cache-hit, failed, and cancelled runs. Add a backward-compatible status field -and migration rather than overloading free-form error text. - -The execution timer must end after public response projection, not merely when -the internal route outcome is assembled. If trace persistence is best-effort, -its latency update must still use the captured execution timestamps and must -not extend the user-visible retrieval duration with an unbounded database -flush. - -Because the map-nav route deliberately rolls back its request session before -the synchronous LLM episode, a trace row must not be created in that session -before the rollback. Capture the execution start immediately, then create or -update the trace record at the terminal persistence point with the captured -start and an explicit end/duration supplied by the execution plan (or use a -separate trace session). `TraceRecorder.complete()` must not silently choose a -later local constructor time or include its own flush duration. - -## Operational requirements - -- heartbeat every 15 seconds while active; -- `Cache-Control: no-cache` and `X-Accel-Buffering: no`; -- flush after every event; -- use `fetch`-style clients where POST authorization headers are required; -- propagate disconnects to the cancellation token. -- bound maximum stream lifetime and enforce the same request/rate-limit policy - as the JSON endpoint; -- document worker-thread, database-connection, and concurrent-stream limits. -- count one stream request as one retrieval request under the existing user and - system limits; retries count as new requests and cannot bypass quota. - -## Delivery slices - -1. **Knowhere contract and timing:** typed event models, stream route, callback - bridge, cancellation, route admission, cache semantics, corrected - `RetrievalRun` timing across all route types, and API contract tests. -2. **SDK adapters:** typed Python and Node stream consumers with fallback to - the existing JSON query; parse named SSE events, expose abort/error - handling, and preserve the full terminal response shape. -3. **Online Brain UX:** phase state model, progress view, safe error states, - and integration with existing downstream answer-token streaming. -4. **Production verification:** proxy-path e2e tests and latency/buffering - instrumentation before setting hard p50/p95 targets. - -## Verification gates - -- correct phase order for map-nav, classic, and small-corpus routes; -- cache-hit streams emit only applicable phases and identify the cache hit; -- no sensitive planner or evidence data before the terminal event; -- terminal citations match the existing JSON endpoint; -- cancellation, timeout, no-results, provider failure, and disconnect are - distinguishable; -- persisted latency includes the full retrieval path for every supported route - and matches API timing within an expected tolerance; -- cache-hit latency and route type are visible in observability data; -- heartbeats pass through the deployed proxy without buffering; -- first-event and first-token times are measured separately. - -## Explicit non-goals - -- v1 API streaming endpoint; -- resumable/replayed streams; -- partial authoritative evidence or provisional citation revision; -- moving answer generation into Knowhere; -- invented latency SLAs before a production baseline exists. diff --git a/packages/shared-python/shared/models/database/document.py b/packages/shared-python/shared/models/database/document.py index c396283d..33dda2f9 100644 --- a/packages/shared-python/shared/models/database/document.py +++ b/packages/shared-python/shared/models/database/document.py @@ -12,10 +12,8 @@ DateTime, Float, ForeignKey, - BigInteger, Index, Integer, - LargeBinary, String, Text, UniqueConstraint, @@ -132,13 +130,6 @@ class DocumentSection(Base): ), Index("idx_document_sections_scope", "user_id", "namespace"), Index("idx_document_sections_doc_revision", "document_id", "job_result_id"), - Index( - "idx_document_sections_revision_snapshot_order", - "document_id", - "job_result_id", - "sort_order", - "section_id", - ), ) @@ -310,13 +301,6 @@ class DocumentMapUnitToken(Base): "token_hash", "map_unit_id", ), - Index( - "idx_document_map_unit_tokens_unit_lookup", - "map_unit_id", - "channel", - "token_hash", - postgresql_include=["token", "frequency"], - ), Index("idx_document_map_unit_tokens_unit", "map_unit_id", "channel"), ) @@ -356,168 +340,6 @@ class DocumentMapUnitIndex(Base): ) -class RetrievalNamespaceGeneration(Base): - """Monotonic serving generation for one user-owned namespace.""" - - __tablename__ = "retrieval_namespace_generations" - - id: Mapped[str] = mapped_column( - String(100), primary_key=True, default=lambda: f"rng_{uuid4().hex}" - ) - user_id: Mapped[str] = mapped_column(Text, nullable=False) - namespace: Mapped[str] = mapped_column(String(255), nullable=False) - generation: Mapped[int] = mapped_column(BigInteger, nullable=False, default=0) - updated_at: Mapped[datetime] = mapped_column( - DateTime, default=utc_now_naive, onupdate=utc_now_naive, nullable=False - ) - - __table_args__ = ( - UniqueConstraint( - "user_id", - "namespace", - name="uq_retrieval_namespace_generations_scope", - ), - ) - - -class RetrievalServingRevisionManifest(Base): - """Compressed ordered metadata for one document revision.""" - - __tablename__ = "retrieval_serving_revision_manifests" - - id: Mapped[str] = mapped_column( - String(100), primary_key=True, default=lambda: f"rsm_{uuid4().hex}" - ) - user_id: Mapped[str] = mapped_column(Text, nullable=False) - namespace: Mapped[str] = mapped_column(String(255), nullable=False) - document_id: Mapped[str] = mapped_column( - String(36), ForeignKey("documents.document_id", ondelete="CASCADE"), nullable=False - ) - job_result_id: Mapped[str] = mapped_column( - String(36), ForeignKey("job_results.id", ondelete="CASCADE"), nullable=False - ) - format_version: Mapped[int] = mapped_column(Integer, nullable=False) - payload_zlib: Mapped[bytes] = mapped_column(LargeBinary, nullable=False) - checksum: Mapped[str] = mapped_column(String(64), nullable=False) - created_at: Mapped[datetime] = mapped_column( - DateTime, default=utc_now_naive, nullable=False - ) - - __table_args__ = ( - UniqueConstraint( - "document_id", - "job_result_id", - name="uq_retrieval_serving_revision_manifests_revision", - ), - Index( - "idx_retrieval_serving_revision_manifests_scope", - "user_id", - "namespace", - "document_id", - "job_result_id", - ), - ) - - -class RetrievalServingRevisionStat(Base): - """Compressed scoring contribution for one document revision.""" - - __tablename__ = "retrieval_serving_revision_stats" - - id: Mapped[str] = mapped_column( - String(100), primary_key=True, default=lambda: f"rss_{uuid4().hex}" - ) - user_id: Mapped[str] = mapped_column(Text, nullable=False) - namespace: Mapped[str] = mapped_column(String(255), nullable=False) - document_id: Mapped[str] = mapped_column( - String(36), ForeignKey("documents.document_id", ondelete="CASCADE"), nullable=False - ) - job_result_id: Mapped[str] = mapped_column( - String(36), ForeignKey("job_results.id", ondelete="CASCADE"), nullable=False - ) - format_version: Mapped[int] = mapped_column(Integer, nullable=False) - payload_zlib: Mapped[bytes] = mapped_column(LargeBinary, nullable=False) - checksum: Mapped[str] = mapped_column(String(64), nullable=False) - created_at: Mapped[datetime] = mapped_column( - DateTime, default=utc_now_naive, nullable=False - ) - - __table_args__ = ( - UniqueConstraint( - "document_id", - "job_result_id", - name="uq_retrieval_serving_revision_stats_revision", - ), - Index( - "idx_retrieval_serving_revision_stats_scope", - "user_id", - "namespace", - "document_id", - "job_result_id", - ), - ) - - -class RetrievalNamespaceStat(Base): - """Compressed aggregate scoring statistics for one namespace generation.""" - - __tablename__ = "retrieval_namespace_stats" - - id: Mapped[str] = mapped_column( - String(100), primary_key=True, default=lambda: f"rns_{uuid4().hex}" - ) - user_id: Mapped[str] = mapped_column(Text, nullable=False) - namespace: Mapped[str] = mapped_column(String(255), nullable=False) - generation: Mapped[int] = mapped_column(BigInteger, nullable=False) - payload_zlib: Mapped[bytes] = mapped_column(LargeBinary, nullable=False) - checksum: Mapped[str] = mapped_column(String(64), nullable=False) - updated_at: Mapped[datetime] = mapped_column( - DateTime, default=utc_now_naive, onupdate=utc_now_naive, nullable=False - ) - - __table_args__ = ( - UniqueConstraint( - "user_id", - "namespace", - name="uq_retrieval_namespace_stats_scope", - ), - ) - - -class RetrievalNamespaceTokenStat(Base): - """Document frequency for one token/channel in a namespace generation.""" - - __tablename__ = "retrieval_namespace_token_stats" - - id: Mapped[str] = mapped_column( - String(100), primary_key=True, default=lambda: f"rnt_{uuid4().hex}" - ) - user_id: Mapped[str] = mapped_column(Text, nullable=False) - namespace: Mapped[str] = mapped_column(String(255), nullable=False) - generation: Mapped[int] = mapped_column(BigInteger, nullable=False) - channel: Mapped[str] = mapped_column(String(32), nullable=False) - token_hash: Mapped[str] = mapped_column(String(64), nullable=False) - document_frequency: Mapped[int] = mapped_column(Integer, nullable=False) - - __table_args__ = ( - UniqueConstraint( - "user_id", - "namespace", - "channel", - "token_hash", - name="uq_retrieval_namespace_token_stats_key", - ), - Index( - "idx_retrieval_namespace_token_stats_lookup", - "user_id", - "namespace", - "generation", - "channel", - "token_hash", - ), - ) - - class GraphNode(Base): """Persisted derived graph node used for routing and expansion.""" diff --git a/packages/shared-python/shared/services/retrieval/execution/plan.py b/packages/shared-python/shared/services/retrieval/execution/plan.py index adb2a6da..8d05753a 100644 --- a/packages/shared-python/shared/services/retrieval/execution/plan.py +++ b/packages/shared-python/shared/services/retrieval/execution/plan.py @@ -1,7 +1,6 @@ from __future__ import annotations import time -from dataclasses import replace from typing import Any from loguru import logger @@ -17,10 +16,6 @@ set_cached_retrieval_query_result, ) from shared.services.retrieval.execution.routes import run_retrieval_route -from shared.services.retrieval.execution.revision_pins import ( - capture_revision_pins, - is_revision_generation_stable, -) from shared.services.retrieval.stats.recorder import ( schedule_retrieval_hit_stats_update, ) @@ -147,26 +142,7 @@ async def _execute_with_overrides(self, request: RetrievalQuery) -> dict[str, An logger.debug(f" 📦 Cache miss (version={cache_version}), running full pipeline") - route_context = request.build_route_context() - revision_pins = await capture_revision_pins( - request.db, - user_id=request.user_id, - namespace=request.namespace, - ) - if not await is_revision_generation_stable( - request.db, - user_id=request.user_id, - namespace=request.namespace, - pins=revision_pins, - ): - revision_pins = await capture_revision_pins( - request.db, - user_id=request.user_id, - namespace=request.namespace, - ) - outcome = await run_retrieval_route( - replace(route_context, revision_pins=revision_pins) - ) + outcome = await run_retrieval_route(request.build_route_context()) if cache_version is not None: await _write_cached_response( diff --git a/packages/shared-python/shared/services/retrieval/execution/reference_resolver.py b/packages/shared-python/shared/services/retrieval/execution/reference_resolver.py index 459ad273..802fa278 100644 --- a/packages/shared-python/shared/services/retrieval/execution/reference_resolver.py +++ b/packages/shared-python/shared/services/retrieval/execution/reference_resolver.py @@ -1,7 +1,6 @@ from __future__ import annotations from dataclasses import dataclass -from collections.abc import Mapping from typing import Any from sqlalchemy.ext.asyncio import AsyncSession @@ -29,7 +28,6 @@ async def resolve_workflow_references( namespace: str, refs: list[dict[str, Any]], score_by_chunk_id: dict[str, float] | None = None, - revision_pins: Mapping[str, str] | None = None, ) -> ResolvedWorkflowReferences: hydrated_rows = await hydrate_referenced_chunk_rows( db=db, @@ -37,7 +35,6 @@ async def resolve_workflow_references( namespace=namespace, refs=refs, score_by_chunk_id=score_by_chunk_id, - revision_pins=revision_pins, ) resolved = _select_matching_references(refs, hydrated_rows) enriched_rows = await enrich_referenced_chunks_with_asset_url(resolved.rows) diff --git a/packages/shared-python/shared/services/retrieval/execution/revision_pins.py b/packages/shared-python/shared/services/retrieval/execution/revision_pins.py deleted file mode 100644 index e7e9efcc..00000000 --- a/packages/shared-python/shared/services/retrieval/execution/revision_pins.py +++ /dev/null @@ -1,92 +0,0 @@ -"""Capture and carry one immutable revision set through a retrieval request.""" - -from __future__ import annotations - -from collections.abc import Iterator, Mapping -from dataclasses import dataclass -from types import MappingProxyType - -from sqlalchemy import select -from sqlalchemy.exc import SQLAlchemyError -from sqlalchemy.ext.asyncio import AsyncSession - -from shared.models.database.document import Document -from shared.models.database.document import RetrievalNamespaceGeneration - - -@dataclass(frozen=True) -class RetrievalRevisionPins(Mapping[str, str]): - """The active document revisions admitted to one retrieval request.""" - - revisions: Mapping[str, str] - generation: int | None = None - - def __post_init__(self) -> None: - object.__setattr__(self, "revisions", MappingProxyType(dict(self.revisions))) - - def __getitem__(self, document_id: str) -> str: - return self.revisions[document_id] - - def __iter__(self) -> Iterator[str]: - return iter(self.revisions) - - def __len__(self) -> int: - return len(self.revisions) - - -async def capture_revision_pins( - db: AsyncSession, - *, - user_id: str, - namespace: str, -) -> RetrievalRevisionPins: - """Capture active document revisions in one database read transaction.""" - statement = ( - select(Document.document_id, Document.current_job_result_id) - .where(Document.user_id == user_id) - .where(Document.namespace == namespace) - .where(Document.status == "active") - .where(Document.current_job_result_id.is_not(None)) - .order_by(Document.document_id) - ) - try: - generation_result = await db.execute( - select(RetrievalNamespaceGeneration.generation) - .where(RetrievalNamespaceGeneration.user_id == user_id) - .where(RetrievalNamespaceGeneration.namespace == namespace) - ) - generation_row = generation_result.scalar_one_or_none() - except SQLAlchemyError: - await db.rollback() - generation_row = None - rows = (await db.execute(statement)).all() - revisions = { - str(document_id): str(job_result_id) - for document_id, job_result_id in rows - if document_id and job_result_id - } - return RetrievalRevisionPins( - revisions=revisions, - generation=int(generation_row) if generation_row is not None else 0, - ) - - -async def is_revision_generation_stable( - db: AsyncSession, - *, - user_id: str, - namespace: str, - pins: RetrievalRevisionPins, -) -> bool: - """Return whether the namespace generation is unchanged since capture.""" - try: - result = await db.execute( - select(RetrievalNamespaceGeneration.generation) - .where(RetrievalNamespaceGeneration.user_id == user_id) - .where(RetrievalNamespaceGeneration.namespace == namespace) - ) - current_generation = result.scalar_one_or_none() - except SQLAlchemyError: - await db.rollback() - return True - return int(current_generation or 0) == int(pins.generation or 0) diff --git a/packages/shared-python/shared/services/retrieval/execution/route_types.py b/packages/shared-python/shared/services/retrieval/execution/route_types.py index 13c21c46..59c173e4 100644 --- a/packages/shared-python/shared/services/retrieval/execution/route_types.py +++ b/packages/shared-python/shared/services/retrieval/execution/route_types.py @@ -5,8 +5,6 @@ from sqlalchemy.ext.asyncio import AsyncSession -from shared.services.retrieval.execution.revision_pins import RetrievalRevisionPins - @dataclass(frozen=True) class RetrievalRouteContext: @@ -28,7 +26,6 @@ class RetrievalRouteContext: internal_recall_k: int | None effective_recall_k: int use_agentic: bool | None - revision_pins: RetrievalRevisionPins | None = None @dataclass(frozen=True) diff --git a/packages/shared-python/shared/services/retrieval/execution/routes.py b/packages/shared-python/shared/services/retrieval/execution/routes.py index d2b6d211..97961074 100644 --- a/packages/shared-python/shared/services/retrieval/execution/routes.py +++ b/packages/shared-python/shared/services/retrieval/execution/routes.py @@ -8,29 +8,18 @@ from sqlalchemy.ext.asyncio import AsyncSession from shared.services.retrieval.search.discovery import bottom_discovery -from shared.services.retrieval.execution.reference_resolver import ( - resolve_workflow_references, -) -from shared.services.retrieval.hydration.result_assembly import ( - assemble_retrieval_results, -) -from shared.services.retrieval.hydration.legacy_evidence import ( - render_legacy_evidence_text, -) +from shared.services.retrieval.execution.reference_resolver import resolve_workflow_references +from shared.services.retrieval.hydration.result_assembly import assemble_retrieval_results +from shared.services.retrieval.hydration.legacy_evidence import render_legacy_evidence_text from shared.services.retrieval.execution.route_types import ( RetrievalRouteContext, RetrievalRouteOutcome, ) from shared.services.retrieval.search.ranking import rank_retrieval_candidates from shared.services.retrieval.search.scoped_corpus import ( - count_manifest_chunks, count_scoped_chunks, load_all_scoped_chunks, ) -from shared.services.retrieval.execution.revision_pins import ( - capture_revision_pins, - is_revision_generation_stable, -) def open_fresh_database_context() -> AbstractAsyncContextManager[AsyncSession]: @@ -57,28 +46,13 @@ async def run_retrieval_route( async def _try_run_small_corpus_route( context: RetrievalRouteContext, ) -> RetrievalRouteOutcome | None: - total_chunk_count: int | None = None - if ( - context.use_agentic is not False - and context.revision_pins is not None - and not context.exclude_document_ids - and not context.exclude_sections - and context.allowed_chunk_types is None - and not context.signal_paths - ): - total_chunk_count = await count_manifest_chunks( - context.db, - revision_pins=context.revision_pins, - ) - if total_chunk_count is None: - total_chunk_count = await count_scoped_chunks( - context.db, - user_id=context.user_id, - namespace=context.namespace, - exclude_document_ids=context.exclude_document_ids, - allowed_chunk_types=context.allowed_chunk_types, - revision_pins=context.revision_pins, - ) + total_chunk_count = await count_scoped_chunks( + context.db, + user_id=context.user_id, + namespace=context.namespace, + exclude_document_ids=context.exclude_document_ids, + allowed_chunk_types=context.allowed_chunk_types, + ) logger.info(f"\n Total chunks in scope: {total_chunk_count}") if total_chunk_count > context.top_k: @@ -97,7 +71,6 @@ async def _try_run_small_corpus_route( allowed_chunk_types=context.allowed_chunk_types, signal_paths=context.signal_paths or [], filter_mode=context.filter_mode, - revision_pins=context.revision_pins, ) logger.info( f" small_corpus load: loaded={len(all_rows)} rows after signal/exclude filters" @@ -108,7 +81,6 @@ async def _try_run_small_corpus_route( exclude_document_ids=context.exclude_document_ids, exclude_sections=context.exclude_sections, allowed_chunk_types=context.allowed_chunk_types, - revision_pins=context.revision_pins, ) results = assembled_rows response = { @@ -145,7 +117,6 @@ async def _run_classic_topk_route( channels=context.channels, channel_weights=context.channel_weights, internal_recall_k=context.internal_recall_k, - revision_pins=context.revision_pins, ) fused_rows = ( @@ -161,7 +132,6 @@ async def _run_classic_topk_route( discovery_rows=fused_rows, routed_rows=[], top_k=context.top_k, - revision_pins=context.revision_pins, ) assembled_rows = await assemble_retrieval_results( @@ -170,7 +140,6 @@ async def _run_classic_topk_route( exclude_document_ids=context.exclude_document_ids, exclude_sections=context.exclude_sections, allowed_chunk_types=context.allowed_chunk_types, - revision_pins=context.revision_pins, ) results = assembled_rows response = { @@ -212,9 +181,7 @@ async def _run_mapnav_route( episode_token_count, episode_workflow_plan, ) - snapshot_started = time.perf_counter() - snapshot_pins = context.revision_pins snapshot = await load_nav_snapshot( context.db, user_id=context.user_id, @@ -222,29 +189,7 @@ async def _run_mapnav_route( exclude_document_ids=context.exclude_document_ids, exclude_sections=context.exclude_sections, lazy=True, - revision_pins=snapshot_pins, ) - if snapshot_pins is not None and not await is_revision_generation_stable( - context.db, - user_id=context.user_id, - namespace=context.namespace, - pins=snapshot_pins, - ): - snapshot.close() - snapshot_pins = await capture_revision_pins( - context.db, - user_id=context.user_id, - namespace=context.namespace, - ) - snapshot = await load_nav_snapshot( - context.db, - user_id=context.user_id, - namespace=context.namespace, - exclude_document_ids=context.exclude_document_ids, - exclude_sections=context.exclude_sections, - lazy=True, - revision_pins=snapshot_pins, - ) snapshot_seconds = time.perf_counter() - snapshot_started logger.info( "retrieval mapnav stage=snapshot_load seconds={:.3f} documents={} refs={}".format( @@ -294,7 +239,6 @@ async def _run_mapnav_route( namespace=context.namespace, refs=refs, score_by_chunk_id=score_by_chunk_id or None, - revision_pins=snapshot.document_revisions, ) assembled_rows = await assemble_retrieval_results( db=final_db, @@ -302,7 +246,6 @@ async def _run_mapnav_route( exclude_document_ids=context.exclude_document_ids, exclude_sections=context.exclude_sections, allowed_chunk_types=context.allowed_chunk_types, - revision_pins=snapshot.document_revisions, ) decision_steps = build_decision_trace( @@ -356,7 +299,9 @@ async def _run_mapnav_route( "decision_trace": decision_trace, } - completion_detail = f"chunks | evidence={len(evidence_text)} chars | router=mapnav" + completion_detail = ( + f"chunks | evidence={len(evidence_text)} chars | router=mapnav" + ) return RetrievalRouteOutcome( response=response, hit_stats_results=resolved.refs, diff --git a/packages/shared-python/shared/services/retrieval/hydration/connected.py b/packages/shared-python/shared/services/retrieval/hydration/connected.py index 2a1f3127..fcbc647b 100644 --- a/packages/shared-python/shared/services/retrieval/hydration/connected.py +++ b/packages/shared-python/shared/services/retrieval/hydration/connected.py @@ -1,6 +1,5 @@ from __future__ import annotations -from collections.abc import Mapping from typing import Any from sqlalchemy import and_, or_, select @@ -21,7 +20,6 @@ async def hydrate_connected_target_rows( rows: list[dict[str, Any]], exclude_document_ids: list[str], exclude_sections: list[dict[str, str]], - revision_pins: Mapping[str, str] | None = None, ) -> list[dict[str, Any]]: if db is None: return [] @@ -63,25 +61,7 @@ async def hydrate_connected_target_rows( stmt = ( select(Document, DocumentChunk, DocumentSection, JobResult) - .join( - DocumentChunk, - ( - (DocumentChunk.document_id == Document.document_id) - if revision_pins is None - else and_( - DocumentChunk.document_id == Document.document_id, - or_( - *[ - and_( - DocumentChunk.document_id == document_id, - DocumentChunk.job_result_id == job_result_id, - ) - for document_id, job_result_id in target_ids_by_revision - ] - ), - ) - ), - ) + .join(DocumentChunk, DocumentChunk.document_id == Document.document_id) .outerjoin(DocumentSection, DocumentSection.section_id == DocumentChunk.section_id) .join(JobResult, JobResult.id == DocumentChunk.job_result_id) .where(or_(*revision_filters)) diff --git a/packages/shared-python/shared/services/retrieval/hydration/reference.py b/packages/shared-python/shared/services/retrieval/hydration/reference.py index cd240654..46208df5 100644 --- a/packages/shared-python/shared/services/retrieval/hydration/reference.py +++ b/packages/shared-python/shared/services/retrieval/hydration/reference.py @@ -1,9 +1,8 @@ from __future__ import annotations -from collections.abc import Mapping from typing import Any -from sqlalchemy import and_, or_, select +from sqlalchemy import select from sqlalchemy.ext.asyncio import AsyncSession from shared.models.database.document import Document, DocumentChunk, DocumentSection @@ -21,7 +20,6 @@ async def hydrate_referenced_chunk_rows( namespace: str, refs: list[dict[str, Any]], score_by_chunk_id: dict[str, float] | None = None, - revision_pins: Mapping[str, str] | None = None, ) -> list[dict[str, Any]]: if db is None or not refs: return [] @@ -41,48 +39,22 @@ async def hydrate_referenced_chunk_rows( document_ids = sorted({document_id for document_id, _, _, _ in ref_keys}) chunk_ids = sorted({chunk_id for _, chunk_id, _, _ in ref_keys}) - pinned_document_ids = [ - document_id for document_id in document_ids if revision_pins and document_id in revision_pins - ] - if revision_pins is not None and not pinned_document_ids: - return [] - - if revision_pins is None: - chunk_join = ( - (DocumentChunk.document_id == Document.document_id) - & (DocumentChunk.job_result_id == Document.current_job_result_id) - ) - else: - chunk_join = and_( - DocumentChunk.document_id == Document.document_id, - or_( - *[ - and_( - DocumentChunk.document_id == document_id, - DocumentChunk.job_result_id == str(revision_pins[document_id]), - ) - for document_id in pinned_document_ids - ] - ), - ) - stmt = ( select(Document, DocumentChunk, DocumentSection, JobResult) - .join(DocumentChunk, chunk_join) + .join( + DocumentChunk, + (DocumentChunk.document_id == Document.document_id) + & (DocumentChunk.job_result_id == Document.current_job_result_id), + ) .outerjoin(DocumentSection, DocumentSection.section_id == DocumentChunk.section_id) .join(JobResult, JobResult.id == DocumentChunk.job_result_id) .where(Document.user_id == user_id) .where(Document.namespace == namespace) - .where( - Document.document_id.in_( - document_ids if revision_pins is None else pinned_document_ids - ) - ) + .where(Document.status == 'active') + .where(Document.document_id.in_(document_ids)) .where(DocumentChunk.chunk_id.in_(chunk_ids)) .order_by(DocumentChunk.sort_order) ) - if revision_pins is None: - stmt = stmt.where(Document.status == 'active') result = await db.execute(stmt) rows_by_key: dict[ReferenceLookupKey, dict[str, Any]] = {} diff --git a/packages/shared-python/shared/services/retrieval/hydration/result_assembly.py b/packages/shared-python/shared/services/retrieval/hydration/result_assembly.py index 863bfadb..ed448704 100644 --- a/packages/shared-python/shared/services/retrieval/hydration/result_assembly.py +++ b/packages/shared-python/shared/services/retrieval/hydration/result_assembly.py @@ -1,6 +1,5 @@ from __future__ import annotations -from collections.abc import Mapping from typing import Any from sqlalchemy.ext.asyncio import AsyncSession @@ -22,7 +21,6 @@ async def assemble_retrieval_results( exclude_document_ids: list[str], exclude_sections: list[dict[str, str]], allowed_chunk_types: set[str] | None = None, - revision_pins: Mapping[str, str] | None = None, ) -> list[dict[str, Any]]: filtered_rows = filter_excluded_rows( rows, @@ -39,7 +37,6 @@ async def assemble_retrieval_results( rows=filtered_rows, exclude_document_ids=exclude_document_ids, exclude_sections=exclude_sections, - revision_pins=revision_pins, ) rows_by_chunk_id = { str(row.get('chunk_id') or ''): row diff --git a/packages/shared-python/shared/services/retrieval/manifest_cache.py b/packages/shared-python/shared/services/retrieval/manifest_cache.py deleted file mode 100644 index 3861903e..00000000 --- a/packages/shared-python/shared/services/retrieval/manifest_cache.py +++ /dev/null @@ -1,47 +0,0 @@ -"""Request-scoped decoded serving-manifest reuse.""" - -from __future__ import annotations - -from collections.abc import Mapping -from typing import Any - - -_CACHE_INFO_KEY = "retrieval_serving_manifest_payloads" - - -def manifest_revision_key(revisions: Mapping[str, str]) -> tuple[tuple[str, str], ...]: - """Return a stable key for one immutable revision pin set.""" - return tuple(sorted((str(document_id), str(job_result_id)) for document_id, job_result_id in revisions.items())) - - -def get_cached_manifest_payloads( - db: object, - *, - revisions: Mapping[str, str], -) -> dict[tuple[str, str], dict[str, Any]] | None: - """Return decoded manifests cached on this request's SQLAlchemy session.""" - info = getattr(db, "info", None) - if not isinstance(info, dict): - return None - batches = info.get(_CACHE_INFO_KEY) - if not isinstance(batches, dict): - return None - payloads = batches.get(manifest_revision_key(revisions)) - if not isinstance(payloads, dict): - return None - return payloads - - -def cache_manifest_payloads( - db: object, - *, - revisions: Mapping[str, str], - payloads: dict[tuple[str, str], dict[str, Any]], -) -> None: - """Store one complete decoded manifest batch for this request only.""" - info = getattr(db, "info", None) - if not isinstance(info, dict): - return - batches = info.setdefault(_CACHE_INFO_KEY, {}) - if isinstance(batches, dict): - batches[manifest_revision_key(revisions)] = payloads diff --git a/packages/shared-python/shared/services/retrieval/nav/nav_hierarchy.py b/packages/shared-python/shared/services/retrieval/nav/nav_hierarchy.py index e98a9d4c..d9e42b07 100644 --- a/packages/shared-python/shared/services/retrieval/nav/nav_hierarchy.py +++ b/packages/shared-python/shared/services/retrieval/nav/nav_hierarchy.py @@ -145,22 +145,14 @@ def get_structure(self, section_id: str) -> dict: def _children_for_section_path( self, section_id: str, doc_id: str, limit: Optional[int] = None ) -> List[dict]: + del doc_id # section_id is globally addressable in this provider model. child_ids = [str(c) for c in self._provider.children(section_id)] if limit is not None: child_ids = child_ids[: max(0, int(limit))] - # ``node_meta`` may materialize a lazy subtree to calculate chunk - # counts. Tree traversal needs only the child id/title; avoid an N+1 - # payload load while building the scoring tree. - out: List[dict] = [] - for cid in child_ids: - path = self.path_titles(cid, doc_id) - out.append( - { - "section_id": cid, - "preview": path.rsplit(" / ", 1)[-1] if path else "", - } - ) - return out + return [ + {"section_id": cid, "preview": self._provider.node_meta(cid).title} + for cid in child_ids + ] def section_relation_ids( self, section_id: str, doc_id: str diff --git a/packages/shared-python/shared/services/retrieval/nav/nav_knowhere.py b/packages/shared-python/shared/services/retrieval/nav/nav_knowhere.py index e25234ea..cf1de779 100644 --- a/packages/shared-python/shared/services/retrieval/nav/nav_knowhere.py +++ b/packages/shared-python/shared/services/retrieval/nav/nav_knowhere.py @@ -21,8 +21,6 @@ from __future__ import annotations import os -import logging -import time from hashlib import sha256 from dataclasses import dataclass, field from typing import ( @@ -47,20 +45,12 @@ PersistedScoreUnit, tokenize_query_for_ranker, ) -from shared.services.retrieval.serving_manifest import decode_serving_manifest _ASSET_TYPES = ("table", "image") # Knowhere sentinel path for the virtual document container (not a collectable leaf). ROOT_SECTION_PATH = "Root" _DEFAULT_DSN = "postgresql://root:root123@127.0.0.1:5433/Knowhere" _MAP_UNIT_INDEX_FORMAT_VERSION = 1 -_MAP_SCORE_CHANNELS: Tuple[str, str] = ("path", "content") -# PostgreSQL's trigram planner evaluates one LIKE branch per pattern. Once a -# query has more than a handful of patterns, fetching the bounded pinned unit -# set once and evaluating literal substring hits in Python is faster and keeps -# SQL planning time from growing with the planner's subgoal length. -_TERM_SCORE_FULL_SCAN_PATTERN_THRESHOLD = 8 -_logger = logging.getLogger(__name__) @dataclass(frozen=True) @@ -171,13 +161,6 @@ def knowhere_database_url() -> str: class ChunkStore(Protocol): - def load_chunk_reference_metadata( - self, - document_id: str, - chunk_id: str, - ) -> Optional[Mapping[str, Any]]: - raise NotImplementedError - def load_persisted_score_corpus( self, document_ids: Sequence[str], @@ -226,19 +209,6 @@ def __init__( self._revisions = dict(revisions) self._excluded_sections = set(excluded_sections or ()) self._conn: Optional[_SyncConnection] = None - self._score_manifest_cache: Optional[ - tuple[tuple[tuple[str, str], ...], list[Sequence[object]]] - ] = None - self._score_unit_rows_cache: dict[ - tuple[tuple[str, str], ...], list[Sequence[object]] - ] = {} - self._score_frequency_cache: dict[ - tuple[tuple[str, str], ...], - dict[tuple[str, str], dict[str, int]], - ] = {} - self._score_term_cache: dict[ - tuple[tuple[str, str], ...], dict[str, Tuple[float, ...]] - ] = {} def _connection(self) -> "_SyncConnection": if self._conn is None: @@ -292,30 +262,6 @@ def load_section_units( finally: cur.close() - def load_chunk_reference_metadata( - self, - document_id: str, - chunk_id: str, - ) -> Optional[Mapping[str, Any]]: - """Resolve deferred reference fields for one selected chunk.""" - doc_id = str(document_id).strip() - cid = str(chunk_id).strip() - job_result_id = self._revisions.get(doc_id) - if not doc_id or not cid or not job_result_id: - return None - cur = self._connection().cursor() - try: - cur.execute( - "SELECT file_path FROM document_chunks " - "WHERE document_id = %s AND job_result_id = %s AND chunk_id = %s " - "ORDER BY sort_order DESC, id DESC LIMIT 1", - (doc_id, job_result_id, cid), - ) - row = cur.fetchone() - return {"file_path": str(row[0] or "") or None} if row else None - finally: - cur.close() - def load_persisted_score_corpus( self, document_ids: Sequence[str], @@ -336,77 +282,48 @@ def load_persisted_score_corpus( ] cur = self._connection().cursor() try: - revision_key = tuple(revisions) - cached_manifests = self._score_manifest_cache - if cached_manifests is not None and cached_manifests[0] == revision_key: - manifests = cached_manifests[1] - _logger.info( - "retrieval map-index load stage=manifests cache_hit rows=%d", - len(manifests), - ) - else: - stage_started = time.perf_counter() - try: - cur.execute( - "SELECT indexes.document_id, indexes.job_result_id, " - "indexes.format_version, indexes.unit_count, indexes.token_count, " - "manifests.payload_zlib, manifests.checksum, manifests.format_version, " - "statistics.payload_zlib, statistics.checksum, statistics.format_version " - "FROM document_map_unit_indexes AS indexes " - f"JOIN (VALUES {values_sql}) AS revisions(document_id, job_result_id) " - "ON indexes.document_id = revisions.document_id " - "AND indexes.job_result_id = revisions.job_result_id " - "JOIN retrieval_serving_revision_manifests AS manifests " - "ON manifests.document_id = indexes.document_id " - "AND manifests.job_result_id = indexes.job_result_id " - "JOIN retrieval_serving_revision_stats AS statistics " - "ON statistics.document_id = indexes.document_id " - "AND statistics.job_result_id = indexes.job_result_id", - revision_params, - ) - except Exception as exc: - if getattr(exc, "pgcode", None) == "42P01": - return None - raise - manifests = list(cur.fetchall()) - self._score_manifest_cache = (revision_key, manifests) - _logger.info( - "retrieval map-index load stage=manifests seconds=%.3f rows=%d", - time.perf_counter() - stage_started if cached_manifests is None else 0.0, - len(manifests), + cur.execute( + "SELECT indexes.document_id, indexes.job_result_id, " + "indexes.format_version, indexes.unit_count, indexes.token_count " + "FROM document_map_unit_indexes AS indexes " + f"JOIN (VALUES {values_sql}) AS revisions(document_id, job_result_id) " + "ON indexes.document_id = revisions.document_id " + "AND indexes.job_result_id = revisions.job_result_id", + revision_params, ) + manifests = list(cur.fetchall()) if len(manifests) != len(revisions) or any( - len(row) < 11 - or int(row[2]) != _MAP_UNIT_INDEX_FORMAT_VERSION - or not row[5] - or not row[6] - or not row[8] - or not row[9] - for row in manifests + int(row[2]) != _MAP_UNIT_INDEX_FORMAT_VERSION for row in manifests ): return None - decoded_manifests: dict[tuple[str, str], dict[str, Any]] = {} - decoded_statistics: dict[tuple[str, str], dict[str, Any]] = {} - try: - for row in manifests: - decoded_manifests[(str(row[0]), str(row[1]))] = decode_serving_manifest( - bytes(row[5]), - checksum=str(row[6]), - format_version=int(row[7]), - ) - decoded_statistics[(str(row[0]), str(row[1]))] = decode_serving_manifest( - bytes(row[8]), - checksum=str(row[9]), - format_version=int(row[10]), - ) - except ValueError: - return None - # The marker is written last in the same transaction that inserts - # all units and token rows. A committed marker therefore denotes - # one complete revision snapshot; avoid recounting millions of - # token rows on every request. Any rebuild deletes the marker - # first, so readers fall back to legacy scoring until completion. + cur.execute( + "SELECT COUNT(*), COUNT(DISTINCT (units.document_id, units.unit_id)) " + "FROM document_map_units AS units " + f"JOIN (VALUES {values_sql}) AS revisions(document_id, job_result_id) " + "ON units.document_id = revisions.document_id " + "AND units.job_result_id = revisions.job_result_id", + revision_params, + ) + unit_count_row = cur.fetchone() + indexed_unit_count = int(unit_count_row[0]) if unit_count_row else 0 + distinct_unit_count = int(unit_count_row[1]) if unit_count_row else 0 + expected_count = sum(int(row[3]) for row in manifests) + if indexed_unit_count != expected_count or distinct_unit_count != indexed_unit_count: + return None + cur.execute( + "SELECT COUNT(*) FROM document_map_unit_tokens AS tokens " + "JOIN document_map_units AS units ON units.id = tokens.map_unit_id " + f"JOIN (VALUES {values_sql}) AS revisions(document_id, job_result_id) " + "ON units.document_id = revisions.document_id " + "AND units.job_result_id = revisions.job_result_id", + revision_params, + ) + token_row = cur.fetchone() + indexed_token_count = int(token_row[0]) if token_row else 0 + expected_token_count = sum(int(row[4]) for row in manifests) + if indexed_token_count != expected_token_count: + return None # The public chunk id is content-derived and may repeat within a # revision. The persisted scorer keys scores by that id, so use # the legacy payload path whenever ambiguity would change results. @@ -419,64 +336,25 @@ def load_persisted_score_corpus( for document_id, section_ids in allowed_by_document.items() for section_id in section_ids ] - unit_cache_key = revision_key - all_unit_rows = self._score_unit_rows_cache.get(unit_cache_key, []) - units_cache_hit = unit_cache_key in self._score_unit_rows_cache unit_rows: list[Sequence[object]] = [] if allowed_pairs: - if not units_cache_hit: - stage_started = time.perf_counter() - manifest_rows: list[Sequence[object]] = [] - for document_id, job_result_id in revisions: - payload = decoded_manifests.get((document_id, job_result_id), {}) - raw_units = payload.get("map_units") - if not isinstance(raw_units, list): - manifest_rows = [] - break - for raw_unit in raw_units: - if not isinstance(raw_unit, dict): - manifest_rows = [] - break - row_id = str(raw_unit.get("row_id") or "").strip() - unit_id = str(raw_unit.get("unit_id") or "").strip() - if not row_id or not unit_id: - manifest_rows = [] - break - manifest_rows.append( - ( - row_id, - document_id, - unit_id, - str(raw_unit.get("section_id") or ""), - int(raw_unit.get("path_token_count") or 0), - int(raw_unit.get("content_token_count") or 0), - ) - ) - if not manifest_rows and raw_units: - break - expected_unit_count = sum(int(row[3]) for row in manifests) - all_unit_rows = ( - manifest_rows - if len(manifest_rows) == expected_unit_count - else [] - ) - if expected_unit_count and not all_unit_rows: - return None - self._score_unit_rows_cache[unit_cache_key] = all_unit_rows - else: - stage_started = time.perf_counter() - allowed_pairs_set = set(allowed_pairs) - unit_rows = [ - row - for row in all_unit_rows - if (str(row[1]), str(row[3])) in allowed_pairs_set - ] - _logger.info( - "retrieval map-index load stage=units seconds=%.3f rows=%d cache_hit=%s", - time.perf_counter() - stage_started if not units_cache_hit else 0.0, - len(unit_rows), - units_cache_hit, + allowed_document_ids = [pair[0] for pair in allowed_pairs] + allowed_section_ids = [pair[1] for pair in allowed_pairs] + cur.execute( + "SELECT units.id, units.document_id, units.unit_id, units.section_id, " + "units.path_token_count, units.content_token_count " + "FROM document_map_units AS units " + f"JOIN (VALUES {values_sql}) AS revisions(document_id, job_result_id) " + "ON units.document_id = revisions.document_id " + "AND units.job_result_id = revisions.job_result_id " + "JOIN UNNEST(%s::text[], %s::text[]) " + "AS allowed(document_id, section_id) " + "ON units.document_id = allowed.document_id " + "AND units.section_id = allowed.section_id " + "ORDER BY units.document_id, units.sort_order, units.unit_id", + [*revision_params, allowed_document_ids, allowed_section_ids], ) + unit_rows = list(cur.fetchall()) map_unit_ids = [str(row[0]) for row in unit_rows] unique_queries = list(dict.fromkeys(str(query) for query in queries)) query_tokens_by_query = { @@ -491,88 +369,26 @@ def load_persisted_score_corpus( ) frequencies: Dict[Tuple[str, str], Dict[str, int]] = {} if map_unit_ids and query_tokens: - full_frequency_map = self._score_frequency_cache.get(revision_key) - if full_frequency_map is None: - full_frequency_map = {} - for document_id, job_result_id in revisions: - payload = decoded_statistics.get((document_id, job_result_id), {}) - unit_frequencies = payload.get("unit_frequencies", {}) - if not isinstance(unit_frequencies, dict): - continue - for map_unit_id, by_channel in unit_frequencies.items(): - if not isinstance(by_channel, dict): - continue - for channel in _MAP_SCORE_CHANNELS: - values = by_channel.get(channel, {}) - if isinstance(values, dict): - full_frequency_map[(str(map_unit_id), channel)] = { - str(token): int(value) - for token, value in values.items() - } - self._score_frequency_cache[revision_key] = full_frequency_map - frequencies = { - key: { - token: value - for token, value in values.items() - if token in query_tokens - } - for key, values in full_frequency_map.items() - if key[0] in map_unit_ids - } - expected_units_by_revision = { - (str(row[0]), str(row[1])): int(row[3]) for row in manifests - } - statistics_complete = all( - int( - decoded_statistics.get((document_id, job_result_id), {}).get( - "unit_count", -1 - ) - ) - == expected_units_by_revision.get((document_id, job_result_id), -1) - for document_id, job_result_id in revisions - ) - manifest_frequency_complete = bool(map_unit_ids) and statistics_complete - if map_unit_ids and query_tokens and not manifest_frequency_complete: query_token_hashes = [ sha256(token.encode("utf-8")).hexdigest() for token in query_tokens ] - stage_started = time.perf_counter() cur.execute( "SELECT map_unit_id, channel, token, frequency " "FROM document_map_unit_tokens " "WHERE map_unit_id = ANY(%s) AND token_hash = ANY(%s) " - "AND token = ANY(%s) AND channel = ANY(%s)", - ( - map_unit_ids, - query_token_hashes, - query_tokens, - list(_MAP_SCORE_CHANNELS), - ), + "AND token = ANY(%s)", + (map_unit_ids, query_token_hashes, query_tokens), ) for map_unit_id, channel, token, frequency in cur.fetchall(): frequencies.setdefault((str(map_unit_id), str(channel)), {})[ str(token) ] = int(frequency) - _logger.info( - "retrieval map-index load stage=frequencies seconds=%.3f units=%d", - time.perf_counter() - stage_started, - len(map_unit_ids), - ) - term_cache_key = (revision_key, tuple(unique_queries)) - term_scores = self._score_term_cache.get(term_cache_key) - if term_scores is None: - term_scores = self._load_term_scores( - cur, - map_unit_ids=map_unit_ids, - queries=unique_queries, - query_tokens_by_query=query_tokens_by_query, - ) - self._score_term_cache[term_cache_key] = term_scores - _logger.info( - "retrieval map-index load stage=complete units=%d queries=%d", - len(unit_rows), - len(unique_queries), + term_scores = self._load_term_scores( + cur, + map_unit_ids=map_unit_ids, + queries=unique_queries, + query_tokens_by_query=query_tokens_by_query, ) path_stats = self._load_persisted_bm25_stats( cur, @@ -624,81 +440,34 @@ def _load_term_scores( ) -> Dict[str, Tuple[float, ...]]: if not map_unit_ids or not queries: return {} - # Keep candidate selection in PostgreSQL so the trigram index can - # discard non-matching units, but compute the exact score once per - # returned row in Python. The previous query generated one POSITION - # expression and one OR predicate for every query token. Long planner - # subgoals therefore produced very large SQL statements and repeated - # substring evaluation for the same row. LIKE ANY keeps the SQL shape - # constant while preserving the existing literal substring semantics in - # the final Python scoring pass (LIKE may over-select wildcard matches, - # which are rejected by the literal checks below). - candidate_patterns: list[str] = [] + expressions: List[str] = [] + params: List[object] = [] for query in queries: query_lower = query.lower().strip() if not query_lower: + expressions.append("0.0") continue - candidate_patterns.append(f"%{query_lower}%") - candidate_patterns.extend( - f"%{token}%" for token in query_tokens_by_query[query] if token - ) - candidate_patterns = list(dict.fromkeys(candidate_patterns)) - if not candidate_patterns: - return {} - stage_started = time.perf_counter() - if len(candidate_patterns) > _TERM_SCORE_FULL_SCAN_PATTERN_THRESHOLD: - # Long subgoals make LIKE ANY increasingly expensive even with the - # trigram index. The map-unit id list is already bounded by the - # pinned serving projection, so one text fetch plus literal Python - # checks avoids a query whose shape grows with token count. - cur.execute( - "SELECT id, term_search_text_lower " - "FROM document_map_units WHERE id = ANY(%s)", - (list(map_unit_ids),), - ) - candidate_mode = "full_scan" - else: - cur.execute( - "SELECT id, term_search_text_lower " - "FROM document_map_units " - "WHERE id = ANY(%s) AND term_search_text_lower LIKE ANY(%s)", - (list(map_unit_ids), candidate_patterns), + token_expressions = [ + "CASE WHEN POSITION(%s IN term_search_text_lower) > 0 THEN 1 ELSE 0 END" + for _token in query_tokens_by_query[query] + ] + token_sum = " + ".join(token_expressions) or "0" + expressions.append( + "CASE WHEN POSITION(%s IN term_search_text_lower) > 0 " + f"THEN 100.0 ELSE ({token_sum})::double precision END" ) - candidate_mode = "trigram" - rows = cur.fetchall() - _logger.info( - "retrieval map-index load stage=term_scores units=%d queries=%d mode=%s seconds=%.3f", - len(map_unit_ids), - len(queries), - candidate_mode, - time.perf_counter() - stage_started, + params.append(query_lower) + params.extend(query_tokens_by_query[query]) + params.append(list(map_unit_ids)) + cur.execute( + "SELECT id, " + ", ".join(expressions) + " " + "FROM document_map_units WHERE id = ANY(%s)", + params, ) - scores_by_unit: Dict[str, Tuple[float, ...]] = {} - for row in rows: - if len(row) < 2: - continue - unit_id = str(row[0]) - haystack = str(row[1] or "").lower() - scores: list[float] = [] - for query in queries: - query_lower = query.lower().strip() - if not query_lower: - scores.append(0.0) - elif query_lower in haystack: - scores.append(100.0) - else: - scores.append( - float( - sum( - 1 - for token in query_tokens_by_query[query] - if token in haystack - ) - ) - ) - if any(score > 0.0 for score in scores): - scores_by_unit[unit_id] = tuple(scores) - return scores_by_unit + return { + str(row[0]): tuple(float(value) for value in row[1:]) + for row in cur.fetchall() + } def _load_persisted_bm25_stats( self, @@ -728,7 +497,6 @@ def _load_persisted_bm25_stats( ) average_idf = 0.0 if needs_average_idf and map_unit_ids and document_count: - stage_started = time.perf_counter() cur.execute( "SELECT COALESCE(AVG(LN((%s - frequencies.document_frequency + 0.5) " "/ (frequencies.document_frequency + 0.5))), 0.0) " @@ -740,11 +508,6 @@ def _load_persisted_bm25_stats( ) row = cur.fetchone() average_idf = float(row[0]) if row else 0.0 - _logger.info( - "retrieval map-index load stage=average_idf channel=%s seconds=%.3f", - channel, - time.perf_counter() - stage_started, - ) return PersistedBm25Stats( document_count=document_count, total_length=sum(lengths), diff --git a/packages/shared-python/shared/services/retrieval/nav/nav_map_scores.py b/packages/shared-python/shared/services/retrieval/nav/nav_map_scores.py index 18d793e9..6318f490 100644 --- a/packages/shared-python/shared/services/retrieval/nav/nav_map_scores.py +++ b/packages/shared-python/shared/services/retrieval/nav/nav_map_scores.py @@ -1,8 +1,6 @@ from __future__ import annotations from collections.abc import Iterator -import logging -import time from typing import Any, Dict, List, Optional, Sequence, Set, Tuple from .knowhere_hybrid import ( @@ -18,7 +16,6 @@ # Keep one bulk read bounded when a namespace contains a very large document, # while still replacing the per-document N+1 access pattern. _CORPUS_PREFETCH_GROUP_SIZE = 8 -_logger = logging.getLogger(__name__) def _children_ids(ts: Any, section_id: str, doc_id: str) -> List[str]: @@ -415,17 +412,10 @@ def compute_corpus_map_and_unit_scores_many( str, Tuple[Dict[str, List[str]], Set[str], Dict[str, str]], ] = {} - tree_started = time.perf_counter() for doc_id in valid_doc_ids: root_ids = list(ts.sections_for_doc(doc_id)) children_map, leaves, titles = _walk_tree(ts, doc_id, root_ids) tree_by_doc[doc_id] = (children_map, leaves, titles) - _logger.info( - "retrieval mapnav phase=tree_build seconds=%.3f documents=%d sections=%d", - time.perf_counter() - tree_started, - len(valid_doc_ids), - sum(len(value[0]) for value in tree_by_doc.values()), - ) def unit_factory() -> Iterator[ScoreUnitRow]: prefetch_batch = getattr(ts, "prefetch_document_units_batch", None) @@ -474,32 +464,17 @@ def unit_factory() -> Iterator[ScoreUnitRow]: release(document_id) persisted_loader = getattr(ts, "load_persisted_score_corpus", None) - loader_started = time.perf_counter() persisted_corpus = ( persisted_loader(valid_doc_ids, unique_queries) if callable(persisted_loader) else None ) - _logger.info( - "retrieval mapnav phase=index_load seconds=%.3f persisted=%s", - time.perf_counter() - loader_started, - persisted_corpus is not None, - ) - score_started = time.perf_counter() unit_scores_by_query = ( score_persisted_corpus_many(persisted_corpus, unique_queries) if persisted_corpus is not None else score_unit_stream_hybrid_many(unit_factory, unique_queries) ) - _logger.info( - "retrieval mapnav phase=unit_scoring persisted=%s seconds=%.3f units=%d queries=%d", - persisted_corpus is not None, - time.perf_counter() - score_started, - sum(len(scores) for scores in unit_scores_by_query.values()), - len(unique_queries), - ) results: Dict[str, Tuple[Dict[str, float], Dict[str, float]]] = {} - pooling_started = time.perf_counter() for query in unique_queries: unit_scores = unit_scores_by_query.get(query, {}) map_scores: Dict[str, float] = {} @@ -515,12 +490,6 @@ def unit_factory() -> Iterator[ScoreUnitRow]: ) map_scores[doc_id] = doc_max results[query] = (map_scores, unit_scores) - _logger.info( - "retrieval mapnav phase=map_pooling seconds=%.3f documents=%d sections=%d", - time.perf_counter() - pooling_started, - len(valid_doc_ids), - sum(len(value[0]) for value in tree_by_doc.values()), - ) return results diff --git a/packages/shared-python/shared/services/retrieval/nav/nav_orchestrate.py b/packages/shared-python/shared/services/retrieval/nav/nav_orchestrate.py index 52cf8867..2e82be08 100644 --- a/packages/shared-python/shared/services/retrieval/nav/nav_orchestrate.py +++ b/packages/shared-python/shared/services/retrieval/nav/nav_orchestrate.py @@ -203,8 +203,6 @@ def _relit_map( """ relit = prepared q = (query or "").strip() - if relit is None and q: - relit = state.relit_map_cache.get(q) if relit is None and q: try: from .nav_map_scores import relight_map_for_query @@ -217,7 +215,6 @@ def _relit_map( ) if scores: relit = (scores, units, highlights) - state.relit_map_cache[q] = relit except Exception: relit = None if relit is None: diff --git a/packages/shared-python/shared/services/retrieval/nav/nav_types.py b/packages/shared-python/shared/services/retrieval/nav/nav_types.py index 3c6b5241..8f3366c0 100644 --- a/packages/shared-python/shared/services/retrieval/nav/nav_types.py +++ b/packages/shared-python/shared/services/retrieval/nav/nav_types.py @@ -305,12 +305,6 @@ class NavState: # PLAN's rewritten retrieval_query per subgoal after widen. Overrides the # planned query for the next harvest, and re-scores the shared map with it. subgoal_refined_queries: Dict[str, str] = field(default_factory=dict) - # Episode-local map scores keyed by retrieval query. Checklist waves may - # revisit the same subgoal query; reuse the exact score snapshot instead of - # rebuilding the persisted index and rescoring the corpus. - relit_map_cache: Dict[ - str, Tuple[Dict[str, float], Dict[str, float], List[str]] - ] = field(default_factory=dict) # Per-subgoal "seen but not selected" section ids — hidden from later map # views for that subgoal so widen surfaces siblings instead of dead ends. subgoal_dismissed_section_ids: Dict[str, set[str]] = field(default_factory=dict) diff --git a/packages/shared-python/shared/services/retrieval/nav_snapshot.py b/packages/shared-python/shared/services/retrieval/nav_snapshot.py index c68d6d8c..1edfd374 100644 --- a/packages/shared-python/shared/services/retrieval/nav_snapshot.py +++ b/packages/shared-python/shared/services/retrieval/nav_snapshot.py @@ -8,31 +8,13 @@ from __future__ import annotations import json -import logging -import time from dataclasses import dataclass -from collections.abc import Callable, Iterator, Mapping from typing import Any, Protocol -from sqlalchemy import ( - ARRAY, - Executable, - String, - bindparam, - cast, - func, - literal, - select, - tuple_, -) +from sqlalchemy import Executable, literal, select, tuple_ from sqlalchemy.engine import Result -from shared.models.database.document import ( - Document, - DocumentChunk, - DocumentSection, - RetrievalServingRevisionManifest, -) +from shared.models.database.document import Document, DocumentChunk, DocumentSection from shared.models.database.job_result import JobResult from shared.services.retrieval.nav.nav_knowhere import ( LazyKnowhereProvider, @@ -44,18 +26,14 @@ knowhere_database_url, ) from shared.services.retrieval.search.section_filters import is_excluded_section -from shared.services.retrieval.serving_manifest import decode_serving_manifest -from shared.services.retrieval.manifest_cache import get_cached_manifest_payloads # Keep each payload query bounded under the API's 30-second statement timeout. -# Ten-thousand-row keyset pages avoid OFFSET scans while keeping the reference -# payload bounded under the API's 30-second asyncpg command timeout. +# Ten-thousand-row keyset pages avoid OFFSET scans while keeping each payload +# statement bounded. The contract benchmark verifies this page size against +# the full 2 KiB content and metadata payload. _CHUNK_BATCH_SIZE = 10_000 -# Keep revision predicates bounded while reducing round trips for large -# namespaces. Keyset paging still caps each payload query at 10,000 rows. -_REVISION_GROUP_SIZE = 64 -_logger = logging.getLogger(__name__) +_REVISION_GROUP_SIZE = 32 class SnapshotSession(Protocol): @@ -64,19 +42,15 @@ class SnapshotSession(Protocol): async def execute(self, statement: Executable) -> Result[tuple[object, ...]]: raise NotImplementedError - async def rollback(self) -> None: - raise NotImplementedError - @dataclass(frozen=True) class NavSnapshot: """In-memory corpus for one map-nav episode.""" provider: NamespaceKnowhereProvider - chunk_ref_index: Mapping[str, dict[str, Any]] + chunk_ref_index: dict[str, dict[str, Any]] document_ids: list[str] document_titles: dict[str, str] - document_revisions: Mapping[str, str] | None def close(self) -> None: close = getattr(self.provider, "close", None) @@ -90,12 +64,9 @@ def build_nav_snapshot( sections_by_doc: dict[str, list[SectionRow]], units_by_doc: dict[str, list[UnitRow]], chunk_ref_index: dict[str, dict[str, Any]], - document_revisions: Mapping[str, str] | None = None, ) -> NavSnapshot: """Assemble provider + index from already-fetched rows (also used by tests).""" - doc_ids = [ - did for did in document_titles if did in sections_by_doc or did in units_by_doc - ] + doc_ids = [did for did in document_titles if did in sections_by_doc or did in units_by_doc] if not doc_ids: raise ValueError("nav snapshot requires at least one active document") @@ -113,18 +84,7 @@ def build_nav_snapshot( provider=provider, chunk_ref_index=dict(chunk_ref_index), document_ids=list(provider.document_ids()), - document_titles={ - did: document_titles.get(did, did) for did in provider.document_ids() - }, - document_revisions=( - { - did: str(document_revisions.get(did, "")) - for did in provider.document_ids() - if document_revisions.get(did) - } - if document_revisions is not None - else None - ), + document_titles={did: document_titles.get(did, did) for did in provider.document_ids()}, ) @@ -136,38 +96,19 @@ async def load_nav_snapshot( exclude_document_ids: list[str] | None = None, exclude_sections: list[dict[str, str]] | None = None, lazy: bool = False, - revision_pins: Mapping[str, str] | None = None, ) -> NavSnapshot: """Preload namespace current revision into a sync map-nav snapshot.""" - excluded_docs = [ - str(x).strip() for x in (exclude_document_ids or ()) if str(x).strip() - ] + excluded_docs = [str(x).strip() for x in (exclude_document_ids or ()) if str(x).strip()] excluded_secs = list(exclude_sections or ()) - if revision_pins is None: - doc_stmt = ( - select( - Document.document_id, - Document.source_file_name, - Document.current_job_result_id, - ) - .where(Document.user_id == user_id) - .where(Document.namespace == namespace) - .where(Document.status == "active") - .where(Document.current_job_result_id.is_not(None)) - .order_by(Document.document_id) - ) - else: - pinned_document_ids = [str(document_id) for document_id in revision_pins] - if not pinned_document_ids: - raise ValueError("revision pins must include at least one document") - doc_stmt = ( - select(Document.document_id, Document.source_file_name) - .where(Document.user_id == user_id) - .where(Document.namespace == namespace) - .where(Document.document_id.in_(pinned_document_ids)) - .order_by(Document.document_id) - ) + doc_stmt = ( + select(Document.document_id, Document.source_file_name, Document.current_job_result_id) + .where(Document.user_id == user_id) + .where(Document.namespace == namespace) + .where(Document.status == "active") + .where(Document.current_job_result_id.is_not(None)) + .order_by(Document.document_id) + ) if excluded_docs: doc_stmt = doc_stmt.where(Document.document_id.notin_(excluded_docs)) doc_rows = list((await db.execute(doc_stmt)).all()) @@ -180,18 +121,12 @@ async def load_nav_snapshot( document_titles: dict[str, str] = {} current_job_result_ids: set[str] = set() document_revisions: list[tuple[str, str]] = [] - for doc_row in doc_rows: - document_id = doc_row[0] - source_file_name = doc_row[1] + for document_id, source_file_name, current_job_result_id in doc_rows: did = str(document_id) title = str(source_file_name or "").strip() or did document_titles[did] = title - job_result_id = ( - str(revision_pins.get(did, "")) - if revision_pins is not None - else str(doc_row[2] or "") - ) - if job_result_id: + if current_job_result_id: + job_result_id = str(current_job_result_id) current_job_result_ids.add(job_result_id) document_revisions.append((did, job_result_id)) @@ -206,47 +141,28 @@ async def load_nav_snapshot( if job_result_id and job_id } - manifest_sections = await _load_manifest_sections( + sections_by_doc, section_path_by_id = await _load_sections( db, document_revisions=document_revisions, exclude_sections=excluded_secs, - job_id_by_result_id=job_id_by_result_id, ) - if manifest_sections is None: - sections_by_doc, section_path_by_id = await _load_sections( + if lazy: + chunk_ids_by_doc, chunk_ref_index, remounted_assets = await _load_chunk_index( db, document_revisions=document_revisions, exclude_sections=excluded_secs, + section_path_by_id=section_path_by_id, + job_id_by_result_id=job_id_by_result_id, ) - manifest_chunk_index = None - else: - sections_by_doc, section_path_by_id, manifest_chunk_index = manifest_sections - if lazy: - if manifest_chunk_index is None: - chunk_ids_by_doc, chunk_ref_index, remounted_assets = await _load_chunk_index( - db, - document_revisions=document_revisions, - exclude_sections=excluded_secs, - section_path_by_id=section_path_by_id, - job_id_by_result_id=job_id_by_result_id, - ) - else: - chunk_ids_by_doc, chunk_ref_index, remounted_assets = manifest_chunk_index kept_titles = { - did: title - for did, title in document_titles.items() - if sections_by_doc.get(did) + did: title for did, title in document_titles.items() if sections_by_doc.get(did) } if not kept_titles: raise ValueError( f"nav snapshot empty after excludes for " f"user_id={user_id!r} namespace={namespace!r}" ) - revisions = { - did: result_id - for did, result_id in document_revisions - if did in kept_titles - } + revisions = {did: result_id for did, result_id in document_revisions if did in kept_titles} store = ReadOnlyChunkStore( dsn=knowhere_database_url(), revisions=revisions, @@ -270,9 +186,7 @@ async def load_nav_snapshot( chunk_store=store, known_chunk_ids=chunk_ids_by_doc.get(did, ()), root_asset_ids=remounted_assets.get(did, {}).get("root", ()), - remounted_assets_by_section=remounted_assets.get(did, {}).get( - "owners", {} - ), + remounted_assets_by_section=remounted_assets.get(did, {}).get("owners", {}), ) for did in kept_titles ] @@ -280,9 +194,11 @@ async def load_nav_snapshot( providers, titles=kept_titles, chunk_owner_by_id={ - chunk_id: document_id - for document_id, chunk_ids in chunk_ids_by_doc.items() - for chunk_id in chunk_ids + chunk_id: meta["document_id"] + for chunk_id, meta in chunk_ref_index.items() + if ":" not in chunk_id + and isinstance(meta, dict) + and meta.get("document_id") }, ) except Exception: @@ -290,15 +206,9 @@ async def load_nav_snapshot( raise return NavSnapshot( provider=provider, - chunk_ref_index=LazyChunkRefIndex( - chunk_ref_index, - resolver=store.load_chunk_reference_metadata, - ), + chunk_ref_index=dict(chunk_ref_index), document_ids=list(provider.document_ids()), - document_titles={ - did: kept_titles.get(did, did) for did in provider.document_ids() - }, - document_revisions=dict(revisions), + document_titles={did: kept_titles.get(did, did) for did in provider.document_ids()}, ) units_by_doc, chunk_ref_index = await _load_chunks( @@ -311,7 +221,9 @@ async def load_nav_snapshot( # Keep only documents that still have sections after exclude filters. kept_titles = { - did: title for did, title in document_titles.items() if sections_by_doc.get(did) + did: title + for did, title in document_titles.items() + if sections_by_doc.get(did) } if not kept_titles: raise ValueError( @@ -324,172 +236,7 @@ async def load_nav_snapshot( sections_by_doc={did: sections_by_doc.get(did, []) for did in kept_titles}, units_by_doc={did: units_by_doc.get(did, []) for did in kept_titles}, chunk_ref_index=chunk_ref_index, - document_revisions={ - document_id: job_result_id - for document_id, job_result_id in document_revisions - if document_id in kept_titles - }, - ) - - -async def _load_manifest_sections( - db: SnapshotSession, - *, - document_revisions: list[tuple[str, str]], - exclude_sections: list[dict[str, str]], - job_id_by_result_id: dict[str, str], -) -> tuple[ - dict[str, list[SectionRow]], - dict[str, str], - tuple[dict[str, list[str]], dict[str, dict[str, Any]], dict[str, dict[str, Any]]], -] | None: - """Load section metadata from complete serving manifests when available.""" - cached_payloads = get_cached_manifest_payloads( - db, - revisions=dict(document_revisions), ) - if cached_payloads is not None and len(cached_payloads) == len(document_revisions): - manifest_entries: list[tuple[object, ...]] = [ - (document_id, job_result_id, payload, None, None) - for document_id, job_result_id in document_revisions - for payload in [cached_payloads.get((document_id, job_result_id))] - if payload is not None - ] - if len(manifest_entries) != len(document_revisions): - return None - else: - statement = select( - RetrievalServingRevisionManifest.document_id, - RetrievalServingRevisionManifest.job_result_id, - RetrievalServingRevisionManifest.payload_zlib, - RetrievalServingRevisionManifest.checksum, - RetrievalServingRevisionManifest.format_version, - ).where( - tuple_( - RetrievalServingRevisionManifest.document_id, - RetrievalServingRevisionManifest.job_result_id, - ).in_(document_revisions) - ) - try: - rows = (await db.execute(statement)).all() - except Exception: - await db.rollback() - return None - if len(rows) != len(document_revisions): - return None - manifest_entries = [tuple(row) for row in rows] - by_doc: dict[str, list[SectionRow]] = {} - path_by_id: dict[str, str] = {} - ids_by_doc: dict[str, list[str]] = {} - ref_index: dict[str, dict[str, Any]] = {} - root_assets_by_doc: dict[str, set[str]] = {} - text_connections_by_doc: dict[str, list[tuple[str, str]]] = {} - try: - for document_id, _job_result_id, payload_zlib, checksum, format_version in manifest_entries: - if isinstance(payload_zlib, dict): - payload = payload_zlib - else: - if not isinstance(payload_zlib, (bytes, bytearray, memoryview)): - return None - if checksum is None or format_version is None: - return None - payload = decode_serving_manifest( - bytes(payload_zlib), - checksum=str(checksum), - format_version=int(str(format_version)), - ) - raw_sections = payload.get("sections") - if not isinstance(raw_sections, list): - return None - for raw_section in raw_sections: - if not isinstance(raw_section, dict): - return None - section_path = str(raw_section.get("section_path") or "") - if is_excluded_section( - document_id=str(document_id), - section_path=section_path, - exclude_sections=exclude_sections, - ): - continue - section_id = str(raw_section.get("section_id") or "") - if not section_id or not section_path: - return None - section = SectionRow( - section_id=section_id, - parent_section_id=( - str(raw_section["parent_section_id"]) - if raw_section.get("parent_section_id") - else None - ), - section_path=section_path, - section_title=str(raw_section.get("section_title") or "").strip(), - section_level=int(raw_section.get("section_level") or 0), - summary=str(raw_section.get("summary") or "").strip(), - sort_order=int(raw_section.get("sort_order") or 0), - ) - by_doc.setdefault(str(document_id), []).append(section) - path_by_id[section_id] = section_path - raw_chunks = payload.get("chunks") - if not isinstance(raw_chunks, list): - return None - for raw_chunk in raw_chunks: - if not isinstance(raw_chunk, dict): - return None - chunk_id = str(raw_chunk.get("chunk_id") or "").strip() - if not chunk_id: - return None - section_id = ( - str(raw_chunk["section_id"]) - if raw_chunk.get("section_id") - else None - ) - section_path = path_by_id.get(section_id) if section_id else None - if is_excluded_section( - document_id=str(document_id), - section_path=section_path, - exclude_sections=exclude_sections, - ) or (section_id and section_id not in path_by_id): - continue - chunk_type = str(raw_chunk.get("chunk_type") or "text") - meta = { - "document_id": str(document_id), - "section_path": section_path, - "chunk_type": chunk_type, - "file_path": None, - "job_id": job_id_by_result_id.get(str(_job_result_id)), - } - document_key = str(document_id) - ids_by_doc.setdefault(document_key, []).append(chunk_id) - ref_index[chunk_id] = meta - ref_index[f"{document_key}:{chunk_id}"] = meta - if ( - chunk_type in {"image", "table"} - and section_id - and section_path == "Root" - ): - root_assets_by_doc.setdefault(document_key, set()).add(chunk_id) - connections = raw_chunk.get("connect_to") - if chunk_type == "text" and isinstance(connections, list): - for connection in connections: - target = ( - str(connection.get("target") or "").strip() - if isinstance(connection, dict) - else str(connection or "").strip() - ) - if target: - text_connections_by_doc.setdefault(document_key, []).append( - (section_id or "", target) - ) - except (TypeError, ValueError, KeyError): - return None - remounted: dict[str, dict[str, Any]] = {} - for document_id, asset_ids in root_assets_by_doc.items(): - owners: dict[str, list[str]] = {} - for section_id, target in text_connections_by_doc.get(document_id, ()): - if target in asset_ids: - owners.setdefault(section_id, []).append(target) - remounted[document_id] = {"root": sorted(asset_ids), "owners": owners} - return by_doc, path_by_id, (ids_by_doc, ref_index, remounted) async def _load_chunk_index( @@ -500,14 +247,11 @@ async def _load_chunk_index( section_path_by_id: dict[str, str], job_id_by_result_id: dict[str, str], ) -> tuple[dict[str, list[str]], dict[str, dict[str, Any]], dict[str, dict[str, Any]]]: - """Load only IDs/reference metadata; content and asset paths remain lazy.""" + """Load only IDs/reference metadata; content remains lazy.""" ids_by_doc: dict[str, list[str]] = {} ref_index: dict[str, dict[str, Any]] = {} root_assets_by_doc: dict[str, set[str]] = {} text_connections_by_doc: dict[str, list[tuple[str, str]]] = {} - query_seconds = 0.0 - assembly_seconds = 0.0 - query_count = 0 for group_start in range(0, len(document_revisions), _REVISION_GROUP_SIZE): revision_group = document_revisions[group_start : group_start + _REVISION_GROUP_SIZE] last_key: tuple[str, str, int, str, str] | None = None @@ -519,8 +263,7 @@ async def _load_chunk_index( DocumentChunk.chunk_id, DocumentChunk.section_id, DocumentChunk.chunk_type, - # Asset paths are only needed for selected references and - # are resolved by ``LazyChunkRefIndex`` at bridge time. + DocumentChunk.file_path, DocumentChunk.chunk_metadata["connect_to"].label("connect_to"), DocumentChunk.sort_order, DocumentChunk.id, @@ -546,13 +289,9 @@ async def _load_chunk_index( ) > tuple_(*[literal(value) for value in last_key]) ) - query_started = time.perf_counter() rows = (await db.execute(stmt)).all() - query_seconds += time.perf_counter() - query_started - query_count += 1 if not rows: break - assembly_started = time.perf_counter() for row in rows: document_id = str(row[0]) job_result_id = str(row[1]) @@ -572,7 +311,7 @@ async def _load_chunk_index( "document_id": document_id, "section_path": section_path, "chunk_type": chunk_type, - "file_path": None, + "file_path": str(row[5] or "") or None, "job_id": job_id_by_result_id.get(job_result_id), } ids_by_doc.setdefault(document_id, []).append(chunk_id) @@ -584,7 +323,7 @@ async def _load_chunk_index( and section_path == "Root" ): root_assets_by_doc.setdefault(document_id, set()).add(chunk_id) - connections = row[5] + connections = row[6] if isinstance(connections, str) and connections.strip(): try: connections = json.loads(connections) @@ -600,14 +339,13 @@ async def _load_chunk_index( text_connections_by_doc.setdefault(document_id, []).append( (section_id or "", target) ) - assembly_seconds += time.perf_counter() - assembly_started last = rows[-1] last_key = ( str(last[0]), str(last[1]), - int(last[6] or 0), + int(last[7] or 0), str(last[2] or ""), - str(last[7]), + str(last[8]), ) if len(rows) < _CHUNK_BATCH_SIZE: break @@ -618,58 +356,9 @@ async def _load_chunk_index( if target in asset_ids: owners.setdefault(section_id, []).append(target) remounted[document_id] = {"root": sorted(asset_ids), "owners": owners} - _logger.info( - "retrieval snapshot phase=chunk_index rows=%d queries=%d query_seconds=%.3f assembly_seconds=%.3f", - sum(len(chunk_ids) for chunk_ids in ids_by_doc.values()), - query_count, - query_seconds, - assembly_seconds, - ) return ids_by_doc, ref_index, remounted -class LazyChunkRefIndex(Mapping[str, dict[str, Any]]): - """Reference metadata map that resolves selected asset paths on demand. - - Snapshot construction still records every chunk identity, section path, - type, and job id so navigation ownership is unchanged. ``file_path`` is - fetched only when the exit bridge asks for a selected reference. - """ - - def __init__( - self, - base: Mapping[str, dict[str, Any]], - *, - resolver: Callable[[str, str], Mapping[str, Any] | None], - ) -> None: - self._base = {str(key): dict(value) for key, value in base.items()} - self._resolver = resolver - - def __getitem__(self, key: str) -> dict[str, Any]: - value = self.get(key) - if value is None: - raise KeyError(key) - return value - - def __iter__(self) -> Iterator[str]: - return iter(self._base) - - def __len__(self) -> int: - return len(self._base) - - def get(self, key: str, default: Any = None) -> dict[str, Any] | Any: - value = self._base.get(str(key)) - if value is None: - return default - if value.get("file_path") is None: - document_id = str(value.get("document_id") or "").strip() - chunk_id = str(key).split(":", 1)[-1].strip() - resolved = self._resolver(document_id, chunk_id) - if resolved is not None: - value.update({"file_path": resolved.get("file_path") or None}) - return value - - async def _load_sections( db: SnapshotSession, *, @@ -677,110 +366,53 @@ async def _load_sections( exclude_sections: list[dict[str, str]], ) -> tuple[dict[str, list[SectionRow]], dict[str, str]]: # Captured pairs replace DocumentSection.job_result_id == Document.current_job_result_id. - by_doc: dict[str, list[SectionRow]] = {} - path_by_id: dict[str, str] = {} - document_ids = [document_id for document_id, _ in document_revisions] - job_result_ids = [job_result_id for _, job_result_id in document_revisions] - revision_rows = ( - func.unnest( - cast( - bindparam("section_document_ids", value=document_ids), ARRAY(String()) - ), - cast( - bindparam("section_job_result_ids", value=job_result_ids), - ARRAY(String()), - ), + stmt = ( + select( + DocumentSection.document_id, + DocumentSection.section_id, + DocumentSection.parent_section_id, + DocumentSection.section_path, + DocumentSection.section_title, + DocumentSection.section_level, + DocumentSection.summary, + DocumentSection.sort_order, ) - .table_valued("document_id", "job_result_id") - .render_derived(name="revisions") - ) - last_key: tuple[str, str, int, str] | None = None - query_seconds = 0.0 - assembly_seconds = 0.0 - query_count = 0 - row_count = 0 - while True: - stmt = ( - select( - DocumentSection.document_id, - DocumentSection.section_id, - DocumentSection.parent_section_id, - DocumentSection.section_path, - DocumentSection.section_title, - DocumentSection.section_level, - DocumentSection.summary, - DocumentSection.sort_order, - DocumentSection.job_result_id, - ) - .join( - revision_rows, - (DocumentSection.document_id == revision_rows.c.document_id) - & (DocumentSection.job_result_id == revision_rows.c.job_result_id), - ) - .order_by( + .where( + tuple_( DocumentSection.document_id, DocumentSection.job_result_id, - DocumentSection.sort_order, - DocumentSection.section_id, - ) - .limit(_CHUNK_BATCH_SIZE) + ).in_(document_revisions) ) - if last_key is not None: - stmt = stmt.where( - tuple_( - DocumentSection.document_id, - DocumentSection.job_result_id, - DocumentSection.sort_order, - DocumentSection.section_id, - ) - > tuple_(*[literal(value) for value in last_key]) - ) - query_started = time.perf_counter() - rows = (await db.execute(stmt)).all() - query_seconds += time.perf_counter() - query_started - query_count += 1 - if not rows: - break - row_count += len(rows) - assembly_started = time.perf_counter() - for row in rows: - document_id = str(row[0]) - section_path = str(row[3] or "") - if is_excluded_section( - document_id=document_id, - section_path=section_path, - exclude_sections=exclude_sections, - ): - continue - section_id = str(row[1]) - section = SectionRow( - section_id=section_id, - parent_section_id=str(row[2]) if row[2] else None, - section_path=section_path, - section_title=str(row[4] or "").strip(), - section_level=int(row[5] or 0), - summary=str(row[6] or "").strip(), - sort_order=int(row[7] or 0), - ) - by_doc.setdefault(document_id, []).append(section) - path_by_id[section_id] = section_path - assembly_seconds += time.perf_counter() - assembly_started - last = rows[-1] - last_key = ( - str(last[0]), - str(last[8]), - int(last[7] or 0), - str(last[1]), + .order_by( + DocumentSection.document_id, + DocumentSection.sort_order, + DocumentSection.section_id, ) - if len(rows) < _CHUNK_BATCH_SIZE: - break - _logger.info( - "retrieval snapshot phase=sections rows=%d queries=%d query_seconds=%.3f assembly_seconds=%.3f", - row_count, - query_count, - query_seconds, - assembly_seconds, ) + + by_doc: dict[str, list[SectionRow]] = {} + path_by_id: dict[str, str] = {} + for row in (await db.execute(stmt)).all(): + document_id = str(row[0]) + section_path = str(row[3] or "") + if is_excluded_section( + document_id=document_id, + section_path=section_path, + exclude_sections=exclude_sections, + ): + continue + section_id = str(row[1]) + section = SectionRow( + section_id=section_id, + parent_section_id=str(row[2]) if row[2] else None, + section_path=section_path, + section_title=str(row[4] or "").strip(), + section_level=int(row[5] or 0), + summary=str(row[6] or "").strip(), + sort_order=int(row[7] or 0), + ) + by_doc.setdefault(document_id, []).append(section) + path_by_id[section_id] = section_path return by_doc, path_by_id diff --git a/packages/shared-python/shared/services/retrieval/publication_content.py b/packages/shared-python/shared/services/retrieval/publication_content.py index 34092436..19a7f207 100644 --- a/packages/shared-python/shared/services/retrieval/publication_content.py +++ b/packages/shared-python/shared/services/retrieval/publication_content.py @@ -14,7 +14,6 @@ ) from shared.services.retrieval.map_unit_index import replace_document_map_units from shared.services.retrieval.publication_models import DocumentPublicationScope -from shared.services.retrieval.serving_manifest import persist_revision_serving_state from shared.services.retrieval.search.lexical_text import ( build_content_lexical_text, build_content_search_text, @@ -92,8 +91,6 @@ def replace_document_revision_content( ) db.flush() replace_document_map_units(db, scope=scope) - db.flush() - persist_revision_serving_state(db, scope=scope) class DocumentSectionPublisher: diff --git a/packages/shared-python/shared/services/retrieval/publication_service.py b/packages/shared-python/shared/services/retrieval/publication_service.py index d7c879de..e201916e 100644 --- a/packages/shared-python/shared/services/retrieval/publication_service.py +++ b/packages/shared-python/shared/services/retrieval/publication_service.py @@ -31,13 +31,6 @@ ExistingDocumentScope, PublishedDocumentState, ) -from shared.services.retrieval.serving_generation import ( - advance_namespace_generation, - lock_namespace_generation, -) -from shared.services.retrieval.serving_manifest import ( - rebuild_namespace_serving_statistics, -) def utc_now_naive() -> datetime: @@ -129,23 +122,6 @@ def _publish_document_state_for_job( skipped_all_duplicate=True, ) - existing_namespace = None - if document_id: - existing_namespace = db.execute( - select(Document.namespace).where( - Document.document_id == str(document_id) - ) - ).scalar_one_or_none() - namespaces_to_lock = {namespace} - if existing_namespace: - namespaces_to_lock.add(str(existing_namespace)) - for namespace_to_lock in sorted(namespaces_to_lock): - lock_namespace_generation( - db, - user_id=str(job.user_id), - namespace=namespace_to_lock, - ) - document = self._upsert_document_revision( db, job=job, @@ -180,22 +156,6 @@ def _publish_document_state_for_job( ) db.flush() - rebuild_namespace_serving_statistics( - db, - user_id=scope.user_id, - namespace=scope.namespace, - ) - if existing_namespace and str(existing_namespace) != scope.namespace: - rebuild_namespace_serving_statistics( - db, - user_id=scope.user_id, - namespace=str(existing_namespace), - ) - advance_namespace_generation( - db, - user_id=scope.user_id, - namespace=scope.namespace, - ) return PublishedDocumentState( user_id=str(job.user_id), namespace=namespace, @@ -249,7 +209,6 @@ def _upsert_document_revision( ) return None document.status = "active" - document.namespace = namespace document.archived_at = None document.current_job_result_id = job_result_id document.source_file_name = source_file_name or document.source_file_name diff --git a/packages/shared-python/shared/services/retrieval/search/channels.py b/packages/shared-python/shared/services/retrieval/search/channels.py index cdf67326..1ddea1be 100644 --- a/packages/shared-python/shared/services/retrieval/search/channels.py +++ b/packages/shared-python/shared/services/retrieval/search/channels.py @@ -8,7 +8,6 @@ from __future__ import annotations import time -from collections.abc import Mapping from typing import Any from loguru import logger @@ -62,7 +61,7 @@ FROM document_chunks dc JOIN documents d ON d.document_id = dc.document_id - {revision_join} + AND d.current_job_result_id = dc.job_result_id LEFT JOIN document_sections ds ON ds.section_id = dc.section_id JOIN job_results jr @@ -70,42 +69,12 @@ WHERE d.user_id = :user_id AND d.namespace = :namespace AND d.status = 'active' - {revision_clause} {exclude_clause} {extra_filters} ) """ -def _build_revision_scope( - revision_pins: Mapping[str, str] | None, -) -> tuple[str, str, dict[str, Any]]: - if revision_pins is None: - return ( - "AND d.current_job_result_id = dc.job_result_id", - "", - {}, - ) - - pairs = [ - (str(document_id).strip(), str(job_result_id).strip()) - for document_id, job_result_id in revision_pins.items() - if str(document_id).strip() and str(job_result_id).strip() - ] - if not pairs: - return "", "AND FALSE", {} - - params: dict[str, Any] = {} - placeholders: list[str] = [] - for index, (document_id, job_result_id) in enumerate(pairs): - document_key = f"_pin_document_{index}" - revision_key = f"_pin_revision_{index}" - placeholders.append(f"(:{document_key}, :{revision_key})") - params[document_key] = document_id - params[revision_key] = job_result_id - return "", f"AND (dc.document_id, dc.job_result_id) IN ({', '.join(placeholders)})", params - - def _build_exclude_clause(exclude_document_ids: list[str]) -> str: if not exclude_document_ids: return "" @@ -259,7 +228,6 @@ async def path_channel( allowed_chunk_types: set[str] | None = None, signal_paths: list[str] | None = None, filter_mode: str = "delete", - revision_pins: Mapping[str, str] | None = None, ) -> list[dict[str, Any]]: """Path channel: BM25 over pre-tokenized path search text. @@ -278,7 +246,6 @@ async def path_channel( signal_paths=signal_paths, filter_mode=filter_mode, search_field="path_search_text", - revision_pins=revision_pins, ) @@ -294,7 +261,6 @@ async def content_channel( allowed_chunk_types: set[str] | None = None, signal_paths: list[str] | None = None, filter_mode: str = "delete", - revision_pins: Mapping[str, str] | None = None, ) -> list[dict[str, Any]]: """Content channel: BM25 over pre-tokenized content search text.""" return await _bm25_channel( @@ -309,7 +275,6 @@ async def content_channel( signal_paths=signal_paths, filter_mode=filter_mode, search_field="content_search_text", - revision_pins=revision_pins, ) @@ -326,7 +291,6 @@ async def _bm25_channel( signal_paths: list[str] | None, filter_mode: str, search_field: str, - revision_pins: Mapping[str, str] | None, ) -> list[dict[str, Any]]: if search_field not in {"content_search_text", "path_search_text"}: raise ValueError(f"Unsupported search_field: {search_field}") @@ -353,14 +317,7 @@ async def _bm25_channel( params.update(extra_params) params.update(section_params) - revision_join, revision_clause, revision_params = _build_revision_scope( - revision_pins - ) - params.update(revision_params) - corpus_cte = _SCOPED_CORPUS_CTE.format( - revision_join=revision_join, - revision_clause=revision_clause, exclude_clause=exclude_clause, extra_filters=extra_sql, ) @@ -452,7 +409,6 @@ async def term_channel( allowed_chunk_types: set[str] | None = None, signal_paths: list[str] | None = None, filter_mode: str = "delete", - revision_pins: Mapping[str, str] | None = None, ) -> list[dict[str, Any]]: """Term/grep channel: substring matching on term_search_text. @@ -479,11 +435,6 @@ async def term_channel( ) params.update(extra_params) - revision_join, revision_clause, revision_params = _build_revision_scope( - revision_pins - ) - params.update(revision_params) - ilike_conditions = [] for i, unit in enumerate(query_tokens): param_key = f"unit_{i}" @@ -497,8 +448,6 @@ async def term_channel( where_clause = " OR ".join(ilike_conditions) sql = ( _SCOPED_CORPUS_CTE.format( - revision_join=revision_join, - revision_clause=revision_clause, exclude_clause=exclude_clause, extra_filters=extra_sql ) + f""" diff --git a/packages/shared-python/shared/services/retrieval/search/discovery.py b/packages/shared-python/shared/services/retrieval/search/discovery.py index d71ccdd9..2e3fdd3c 100644 --- a/packages/shared-python/shared/services/retrieval/search/discovery.py +++ b/packages/shared-python/shared/services/retrieval/search/discovery.py @@ -18,9 +18,7 @@ from __future__ import annotations -import asyncio import time -from collections.abc import Awaitable, Callable, Mapping from dataclasses import dataclass, field from typing import Any @@ -70,13 +68,11 @@ async def bottom_discovery( channels: list[str] | None = None, channel_weights: dict[str, float] | None = None, internal_recall_k: int | None = None, - revision_pins: Mapping[str, str] | None = None, **_kwargs: Any, ) -> DiscoveryResult: """Run 3-channel BM25 discovery plus RRF fusion.""" t0 = time.monotonic() try: - del db allowed_chunk_types = chunk_types effective_recall_k = ( internal_recall_k @@ -85,10 +81,13 @@ async def bottom_discovery( ) active_channels = set(channels) if channels else {"path", "content", "term"} - path_rows, content_rows, term_rows = await asyncio.gather( - _run_channel( - path_channel, - enabled="path" in active_channels, + path_rows: list[dict[str, Any]] = [] + content_rows: list[dict[str, Any]] = [] + term_rows: list[dict[str, Any]] = [] + + if "path" in active_channels: + path_rows = await path_channel( + db, user_id=user_id, namespace=namespace, query=query, @@ -98,11 +97,11 @@ async def bottom_discovery( allowed_chunk_types=allowed_chunk_types, signal_paths=signal_paths, filter_mode=filter_mode, - revision_pins=revision_pins, - ), - _run_channel( - content_channel, - enabled="content" in active_channels, + ) + + if "content" in active_channels: + content_rows = await content_channel( + db, user_id=user_id, namespace=namespace, query=query, @@ -112,11 +111,11 @@ async def bottom_discovery( allowed_chunk_types=allowed_chunk_types, signal_paths=signal_paths, filter_mode=filter_mode, - revision_pins=revision_pins, - ), - _run_channel( - term_channel, - enabled="term" in active_channels, + ) + + if "term" in active_channels: + term_rows = await term_channel( + db, user_id=user_id, namespace=namespace, query=query, @@ -126,9 +125,7 @@ async def bottom_discovery( allowed_chunk_types=allowed_chunk_types, signal_paths=signal_paths, filter_mode=filter_mode, - revision_pins=revision_pins, - ), - ) + ) default_weights = { "path": CHANNEL_WEIGHT_PATH, @@ -197,20 +194,3 @@ async def bottom_discovery( latency = int((time.monotonic() - t0) * 1000) logger.error(f" search.bottom_discovery failed: {exc}") return DiscoveryResult(status="error", error=str(exc), latency_ms=latency) - - -async def _run_channel( - channel: Callable[..., Awaitable[list[dict[str, Any]]]], - *, - enabled: bool, - **kwargs: Any, -) -> list[dict[str, Any]]: - if not enabled: - return [] - - # Import lazily so discovery remains usable by lightweight contract tests - # without creating a database context until a channel is actually enabled. - from shared.core.database import get_db_context - - async with get_db_context() as channel_db: - return await channel(channel_db, **kwargs) diff --git a/packages/shared-python/shared/services/retrieval/search/ranking.py b/packages/shared-python/shared/services/retrieval/search/ranking.py index c4906e46..392aa1f1 100644 --- a/packages/shared-python/shared/services/retrieval/search/ranking.py +++ b/packages/shared-python/shared/services/retrieval/search/ranking.py @@ -1,7 +1,6 @@ from __future__ import annotations import math -from collections.abc import Mapping from typing import Any from loguru import logger @@ -16,9 +15,9 @@ def get_candidate_key(row: dict[str, Any]) -> str: path = get_row_path(row) if path: - return f"path:{path}" - chunk_id = str(row.get("chunk_id") or "").strip() - return f"chunk:{chunk_id}" if chunk_id else "" + return f'path:{path}' + chunk_id = str(row.get('chunk_id') or '').strip() + return f'chunk:{chunk_id}' if chunk_id else '' async def load_chunk_importance_scores( @@ -27,11 +26,12 @@ async def load_chunk_importance_scores( user_id: str, namespace: str, rows: list[dict[str, Any]], - revision_pins: Mapping[str, str] | None = None, ) -> dict[str, float]: - chunk_ids = sorted( - {str(row.get("chunk_id") or "").strip() for row in rows if row.get("chunk_id")} - ) + chunk_ids = sorted({ + str(row.get('chunk_id') or '').strip() + for row in rows + if row.get('chunk_id') + }) if not chunk_ids: return {} stmt = ( @@ -43,31 +43,22 @@ async def load_chunk_importance_scores( ) .where(RetrievalHitStat.user_id == user_id) .where(RetrievalHitStat.namespace == namespace) - .where(RetrievalHitStat.hit_kind == "chunk") + .where(RetrievalHitStat.hit_kind == 'chunk') .where(RetrievalHitStat.chunk_id.in_(chunk_ids)) ) - if revision_pins is not None: - pinned_document_ids = {str(document_id) for document_id in revision_pins} - stmt = stmt.where(RetrievalHitStat.document_id.in_(pinned_document_ids)) result = await db.execute(stmt) importance_scores: dict[str, float] = {} for chunk_id, hit_count, last_hit_at, created_at in result.all(): if not chunk_id: continue - score = compute_importance_score(hit_count, last_hit_at, created_at) - importance_scores[str(chunk_id)] = score - for row in rows: - if str(row.get("chunk_id") or "") == str(chunk_id): - document_id = str(row.get("document_id") or "") - if document_id: - importance_scores[f"{document_id}:{chunk_id}"] = score + importance_scores[str(chunk_id)] = compute_importance_score(hit_count, last_hit_at, created_at) return importance_scores def apply_importance_multiplier( rows: list[dict[str, Any]], *, - raw_field: str = "importance_raw_score", + raw_field: str = 'importance_raw_score', low: float = 0.1, high: float = 2.0, ) -> None: @@ -76,11 +67,7 @@ def apply_importance_multiplier( values = sorted(float(row.get(raw_field, 0.0) or 0.0) for row in rows) item_count = len(values) - median = ( - values[item_count // 2] - if item_count % 2 - else (values[item_count // 2 - 1] + values[item_count // 2]) / 2 - ) + median = values[item_count // 2] if item_count % 2 else (values[item_count // 2 - 1] + values[item_count // 2]) / 2 q1 = values[item_count // 4] if item_count >= 4 else values[0] q3 = values[3 * item_count // 4] if item_count >= 4 else values[-1] iqr = q3 - q1 @@ -93,13 +80,13 @@ def apply_importance_multiplier( z_score = (raw_score - median) / iqr sigmoid_score = 1.0 / (1.0 + math.exp(-z_score)) multiplier = low + (high - low) * sigmoid_score - row["importance_multiplier"] = round(multiplier, 4) - row["agent_score"] = round( - float(row.get("agent_score", 0.0) or 0.0) * multiplier, + row['importance_multiplier'] = round(multiplier, 4) + row['agent_score'] = round( + float(row.get('agent_score', 0.0) or 0.0) * multiplier, 6, ) - row["discovery_score"] = round( - float(row.get("discovery_score", 0.0) or 0.0) * multiplier, + row['discovery_score'] = round( + float(row.get('discovery_score', 0.0) or 0.0) * multiplier, 6, ) @@ -120,9 +107,9 @@ def rank_candidates_by_path( if not key: continue candidate = dict(row) - candidate["discovery_score"] = float(row.get("discovery_score", 0.0) or 0.0) - candidate["agent_score"] = 0.0 - candidate.setdefault("hydrate_mode", "chunks") + candidate['discovery_score'] = float(row.get('discovery_score', 0.0) or 0.0) + candidate['agent_score'] = 0.0 + candidate.setdefault('hydrate_mode', 'chunks') merged[key] = candidate insertion_order[key] = counter counter += 1 @@ -131,33 +118,25 @@ def rank_candidates_by_path( key = get_candidate_key(row) if not key: continue - routed_agent_score = float(row.get("agent_score", 0.0) or 0.0) + routed_agent_score = float(row.get('agent_score', 0.0) or 0.0) if key not in merged: candidate = dict(row) - candidate["discovery_score"] = float(row.get("discovery_score", 0.0) or 0.0) - candidate["agent_score"] = routed_agent_score + candidate['discovery_score'] = float(row.get('discovery_score', 0.0) or 0.0) + candidate['agent_score'] = routed_agent_score merged[key] = candidate insertion_order[key] = counter counter += 1 continue candidate = merged[key] - candidate["agent_score"] = max( - float(candidate.get("agent_score", 0.0) or 0.0), routed_agent_score - ) - if not candidate.get("source_chunk_path") and row.get("source_chunk_path"): - candidate["source_chunk_path"] = row.get("source_chunk_path") - if not candidate.get("section_path") and row.get("section_path"): - candidate["section_path"] = row.get("section_path") + candidate['agent_score'] = max(float(candidate.get('agent_score', 0.0) or 0.0), routed_agent_score) + if not candidate.get('source_chunk_path') and row.get('source_chunk_path'): + candidate['source_chunk_path'] = row.get('source_chunk_path') + if not candidate.get('section_path') and row.get('section_path'): + candidate['section_path'] = row.get('section_path') for row in merged.values(): - document_id = str(row.get("document_id") or "") - chunk_id = str(row.get("chunk_id") or "") - row["importance_raw_score"] = float( - (importance_scores or {}).get( - f"{document_id}:{chunk_id}", - (importance_scores or {}).get(chunk_id, 0.0), - ) - or 0.0 + row['importance_raw_score'] = float( + (importance_scores or {}).get(str(row.get('chunk_id') or ''), 0.0) or 0.0 ) apply_importance_multiplier(list(merged.values())) @@ -166,13 +145,11 @@ def rank_candidates_by_path( fallback_rows: list[dict[str, Any]] = [] for key, row in merged.items(): - agent_score = float(row.get("agent_score", 0.0) or 0.0) - discovery_score = float(row.get("discovery_score", 0.0) or 0.0) - row["evidence_score"] = round( - agent_score if has_agent_results else max(discovery_score, agent_score), 6 - ) - row["score"] = row["evidence_score"] - row["_candidate_order"] = insertion_order[key] + agent_score = float(row.get('agent_score', 0.0) or 0.0) + discovery_score = float(row.get('discovery_score', 0.0) or 0.0) + row['evidence_score'] = round(agent_score if has_agent_results else max(discovery_score, agent_score), 6) + row['score'] = row['evidence_score'] + row['_candidate_order'] = insertion_order[key] if has_agent_results and agent_score <= 0.0: fallback_rows.append(row) @@ -181,9 +158,9 @@ def rank_candidates_by_path( def get_sort_key(row: dict[str, Any]) -> tuple[float, float, int]: return ( - float(row.get("agent_score", 0.0) or 0.0), - float(row.get("discovery_score", 0.0) or 0.0), - -int(row.get("_candidate_order", 0) or 0), + float(row.get('agent_score', 0.0) or 0.0), + float(row.get('discovery_score', 0.0) or 0.0), + -int(row.get('_candidate_order', 0) or 0), ) primary_rows.sort(key=get_sort_key, reverse=True) @@ -191,10 +168,10 @@ def get_sort_key(row: dict[str, Any]) -> tuple[float, float, int]: if len(ranked_rows) < top_k and fallback_rows: fallback_rows.sort(key=get_sort_key, reverse=True) - ranked_rows.extend(fallback_rows[: top_k - len(ranked_rows)]) + ranked_rows.extend(fallback_rows[:top_k - len(ranked_rows)]) for row in ranked_rows: - row.pop("_candidate_order", None) + row.pop('_candidate_order', None) return ranked_rows @@ -206,7 +183,6 @@ async def rank_retrieval_candidates( discovery_rows: list[dict[str, Any]], routed_rows: list[dict[str, Any]], top_k: int, - revision_pins: Mapping[str, str] | None = None, ) -> list[dict[str, Any]]: try: importance_scores = await load_chunk_importance_scores( @@ -214,12 +190,9 @@ async def rank_retrieval_candidates( user_id=user_id, namespace=namespace, rows=[*discovery_rows, *routed_rows], - revision_pins=revision_pins, ) except Exception as exc: - logger.warning( - f"Failed to load chunk importance scores, continuing without importance: {exc}" - ) + logger.warning(f'Failed to load chunk importance scores, continuing without importance: {exc}') importance_scores = {} return rank_candidates_by_path( discovery_rows, diff --git a/packages/shared-python/shared/services/retrieval/search/scoped_corpus.py b/packages/shared-python/shared/services/retrieval/search/scoped_corpus.py index ed6f2719..4a60e65b 100644 --- a/packages/shared-python/shared/services/retrieval/search/scoped_corpus.py +++ b/packages/shared-python/shared/services/retrieval/search/scoped_corpus.py @@ -1,72 +1,13 @@ from __future__ import annotations -from collections.abc import Mapping from typing import Any -from sqlalchemy import and_, func, select, tuple_ +from sqlalchemy import func, select from sqlalchemy.ext.asyncio import AsyncSession -from sqlalchemy.exc import SQLAlchemyError -from shared.models.database.document import ( - Document, - DocumentChunk, - DocumentSection, - RetrievalServingRevisionManifest, -) +from shared.models.database.document import Document, DocumentChunk, DocumentSection from shared.models.database.job_result import JobResult from shared.services.retrieval.search.section_filters import is_excluded_section -from shared.services.retrieval.serving_manifest import decode_serving_manifest -from shared.services.retrieval.manifest_cache import cache_manifest_payloads - - -async def count_manifest_chunks( - db: AsyncSession, - *, - revision_pins: Mapping[str, str], -) -> int | None: - """Count chunks exactly from complete pinned manifests when available.""" - if not revision_pins: - return None - statement = select( - RetrievalServingRevisionManifest.document_id, - RetrievalServingRevisionManifest.job_result_id, - RetrievalServingRevisionManifest.payload_zlib, - RetrievalServingRevisionManifest.checksum, - RetrievalServingRevisionManifest.format_version, - ).where( - tuple_( - RetrievalServingRevisionManifest.document_id, - RetrievalServingRevisionManifest.job_result_id, - ).in_(list(revision_pins.items())) - ) - try: - rows = (await db.execute(statement)).all() - if len(rows) != len(revision_pins): - return None - total = 0 - decoded_payloads: dict[tuple[str, str], dict[str, Any]] = {} - for document_id, job_result_id, payload_zlib, checksum, format_version in rows: - payload = decode_serving_manifest( - bytes(payload_zlib), - checksum=str(checksum), - format_version=int(format_version), - ) - chunks = payload.get("chunks") - if not isinstance(chunks, list): - return None - total += len(chunks) - decoded_payloads[(str(document_id), str(job_result_id))] = payload - cache_manifest_payloads( - db, - revisions=revision_pins, - payloads=decoded_payloads, - ) - return total - except SQLAlchemyError: - await db.rollback() - return None - except (TypeError, ValueError, KeyError): - return None async def count_scoped_chunks( @@ -76,32 +17,18 @@ async def count_scoped_chunks( namespace: str, exclude_document_ids: list[str], allowed_chunk_types: set[str] | None, - revision_pins: Mapping[str, str] | None = None, ) -> int: - if revision_pins is None: - stmt = ( - select(func.count(DocumentChunk.id)) - .join( - Document, - (Document.document_id == DocumentChunk.document_id) - & (Document.current_job_result_id == DocumentChunk.job_result_id), - ) - .where(Document.user_id == user_id) - .where(Document.namespace == namespace) - .where(Document.status == 'active') - ) - else: - stmt = ( - select(func.count(DocumentChunk.id)) - .join(Document, Document.document_id == DocumentChunk.document_id) - .where(Document.user_id == user_id) - .where(Document.namespace == namespace) - .where( - tuple_(DocumentChunk.document_id, DocumentChunk.job_result_id).in_( - list(revision_pins.items()) - ) - ) + stmt = ( + select(func.count(DocumentChunk.id)) + .join( + Document, + (Document.document_id == DocumentChunk.document_id) + & (Document.current_job_result_id == DocumentChunk.job_result_id), ) + .where(Document.user_id == user_id) + .where(Document.namespace == namespace) + .where(Document.status == 'active') + ) if exclude_document_ids: stmt = stmt.where(Document.document_id.notin_(list(exclude_document_ids))) if allowed_chunk_types is not None: @@ -120,31 +47,21 @@ async def load_all_scoped_chunks( allowed_chunk_types: set[str] | None, signal_paths: list[str], filter_mode: str, - revision_pins: Mapping[str, str] | None = None, ) -> list[dict[str, Any]]: - if revision_pins is None: - chunk_join = ( - (DocumentChunk.document_id == Document.document_id) - & (DocumentChunk.job_result_id == Document.current_job_result_id) - ) - else: - chunk_join = and_( - DocumentChunk.document_id == Document.document_id, - tuple_(DocumentChunk.document_id, DocumentChunk.job_result_id).in_( - list(revision_pins.items()) - ), - ) stmt = ( select(Document, DocumentChunk, DocumentSection, JobResult) - .join(DocumentChunk, chunk_join) + .join( + DocumentChunk, + (DocumentChunk.document_id == Document.document_id) + & (DocumentChunk.job_result_id == Document.current_job_result_id), + ) .outerjoin(DocumentSection, DocumentSection.section_id == DocumentChunk.section_id) .join(JobResult, JobResult.id == DocumentChunk.job_result_id) .where(Document.user_id == user_id) .where(Document.namespace == namespace) + .where(Document.status == 'active') .order_by(DocumentChunk.sort_order) ) - if revision_pins is None: - stmt = stmt.where(Document.status == 'active') if exclude_document_ids: stmt = stmt.where(Document.document_id.notin_(list(exclude_document_ids))) if allowed_chunk_types is not None: diff --git a/packages/shared-python/shared/services/retrieval/search/scoring.py b/packages/shared-python/shared/services/retrieval/search/scoring.py index 2746de34..848a3ada 100644 --- a/packages/shared-python/shared/services/retrieval/search/scoring.py +++ b/packages/shared-python/shared/services/retrieval/search/scoring.py @@ -51,18 +51,14 @@ def merge_channels_rrf( for channel_idx, channel_rows in enumerate(channels): weight = weights[channel_idx] if channel_idx < len(weights) else 1.0 - seen_chunk_ids: set[str] = set() - unique_rank = 0 - for row in channel_rows: + for rank, row in enumerate(channel_rows): chunk_id = str(row.get('chunk_id') or '') - if not chunk_id or chunk_id in seen_chunk_ids: + if not chunk_id: continue - seen_chunk_ids.add(chunk_id) - rrf_score = weight / (k + unique_rank + 1) + rrf_score = weight / (k + rank + 1) score_dict[chunk_id] = score_dict.get(chunk_id, 0.0) + rrf_score if chunk_id not in row_by_chunk_id: row_by_chunk_id[chunk_id] = row - unique_rank += 1 ranked = sorted(score_dict.items(), key=lambda x: x[1], reverse=True) results: list[dict[str, Any]] = [] diff --git a/packages/shared-python/shared/services/retrieval/serving_generation.py b/packages/shared-python/shared/services/retrieval/serving_generation.py deleted file mode 100644 index 5dc7baf5..00000000 --- a/packages/shared-python/shared/services/retrieval/serving_generation.py +++ /dev/null @@ -1,60 +0,0 @@ -"""Namespace generation locking for serving-state lifecycle updates.""" - -from __future__ import annotations - -from hashlib import sha256 - -from sqlalchemy import select -from sqlalchemy.dialects.postgresql import insert -from sqlalchemy.orm import Session - -from shared.models.database.document import RetrievalNamespaceGeneration - - -def lock_namespace_generation( - db: Session, - *, - user_id: str, - namespace: str, -) -> RetrievalNamespaceGeneration: - """Create if needed, then lock and return one namespace generation row.""" - generation_id = f"rng_{sha256(f'{user_id}:{namespace}'.encode()).hexdigest()}" - db.execute( - insert(RetrievalNamespaceGeneration) - .values( - id=generation_id, - user_id=user_id, - namespace=namespace, - generation=0, - ) - .on_conflict_do_nothing( - index_elements=[ - RetrievalNamespaceGeneration.user_id, - RetrievalNamespaceGeneration.namespace, - ] - ) - ) - generation = db.execute( - select(RetrievalNamespaceGeneration) - .where(RetrievalNamespaceGeneration.user_id == user_id) - .where(RetrievalNamespaceGeneration.namespace == namespace) - .with_for_update() - ).scalar_one() - return generation - - -def advance_namespace_generation( - db: Session, - *, - user_id: str, - namespace: str, -) -> int: - """Increment a locked namespace generation and return its new value.""" - generation = lock_namespace_generation( - db, - user_id=user_id, - namespace=namespace, - ) - generation.generation += 1 - db.flush() - return generation.generation diff --git a/packages/shared-python/shared/services/retrieval/serving_manifest.py b/packages/shared-python/shared/services/retrieval/serving_manifest.py deleted file mode 100644 index b2dd40c5..00000000 --- a/packages/shared-python/shared/services/retrieval/serving_manifest.py +++ /dev/null @@ -1,405 +0,0 @@ -"""Versioned compression and integrity checks for serving manifests.""" - -from __future__ import annotations - -import hashlib -import json -import zlib -from typing import Any - -from sqlalchemy import delete, select -from sqlalchemy.orm import Session - -from shared.models.database.document import ( - Document, - DocumentChunk, - DocumentMapUnit, - DocumentMapUnitToken, - DocumentSection, - RetrievalNamespaceGeneration, - RetrievalNamespaceStat, - RetrievalNamespaceTokenStat, - RetrievalServingRevisionManifest, - RetrievalServingRevisionStat, -) -from shared.models.database.job_result import JobResult -from shared.services.retrieval.publication_models import DocumentPublicationScope - -SERVING_MANIFEST_FORMAT_VERSION = 1 - - -def build_revision_serving_payload( - db: Session, - *, - scope: DocumentPublicationScope, -) -> dict[str, Any]: - """Build ordered metadata for one published document revision.""" - document = db.execute( - select(Document).where(Document.document_id == scope.document_id) - ).scalar_one() - job_result = db.execute( - select(JobResult).where(JobResult.id == scope.job_result_id) - ).scalar_one() - sections = list( - db.scalars( - select(DocumentSection) - .where(DocumentSection.document_id == scope.document_id) - .where(DocumentSection.job_result_id == scope.job_result_id) - .order_by(DocumentSection.sort_order, DocumentSection.section_id) - ) - ) - chunks = list( - db.scalars( - select(DocumentChunk) - .where(DocumentChunk.document_id == scope.document_id) - .where(DocumentChunk.job_result_id == scope.job_result_id) - .order_by( - DocumentChunk.sort_order, DocumentChunk.chunk_id, DocumentChunk.id - ) - ) - ) - map_units = list( - db.scalars( - select(DocumentMapUnit) - .where(DocumentMapUnit.document_id == scope.document_id) - .where(DocumentMapUnit.job_result_id == scope.job_result_id) - .order_by(DocumentMapUnit.sort_order, DocumentMapUnit.unit_id) - ) - ) - section_path_by_id = { - section.section_id: section.section_path for section in sections - } - root_asset_ids = { - chunk.chunk_id - for chunk in chunks - if chunk.chunk_type in {"image", "table"} - and chunk.section_id is not None - and section_path_by_id.get(chunk.section_id) == "Root" - } - remounted_assets: dict[str, list[str]] = {} - for chunk in chunks: - if chunk.chunk_type != "text" or not isinstance(chunk.chunk_metadata, dict): - continue - connections = chunk.chunk_metadata.get("connect_to") - if not isinstance(connections, list): - continue - targets = [ - str(connection.get("target") or "").strip() - for connection in connections - if isinstance(connection, dict) - and str(connection.get("target") or "").strip() in root_asset_ids - ] - if targets: - remounted_assets[chunk.section_id or ""] = targets - - return { - "document_id": scope.document_id, - "job_result_id": scope.job_result_id, - "job_id": str(job_result.job_id), - "source_file_name": str( - document.source_file_name or scope.source_file_name or "" - ), - "sections": [ - { - "section_id": section.section_id, - "parent_section_id": section.parent_section_id, - "section_path": section.section_path, - "section_title": section.section_title, - "section_level": section.section_level, - "summary": section.summary, - "sort_order": section.sort_order, - } - for section in sections - ], - "chunks": [ - { - "chunk_id": chunk.chunk_id, - "section_id": chunk.section_id, - "chunk_type": chunk.chunk_type, - "sort_order": chunk.sort_order, - "connect_to": _connection_target_ids(chunk.chunk_metadata), - } - for chunk in chunks - ], - "map_units": [ - { - "row_id": unit.id, - "unit_id": unit.unit_id, - "section_id": unit.section_id, - "unit_kind": unit.unit_kind, - "path_token_count": unit.path_token_count, - "content_token_count": unit.content_token_count, - "sort_order": unit.sort_order, - } - for unit in map_units - ], - "root_asset_ids": sorted(root_asset_ids), - "remounted_assets_by_section": remounted_assets, - } - - -def build_revision_statistics_payload( - db: Session, - *, - scope: DocumentPublicationScope, -) -> dict[str, Any]: - """Build compressed scoring contributions for one revision.""" - units = list( - db.scalars( - select(DocumentMapUnit) - .where(DocumentMapUnit.document_id == scope.document_id) - .where(DocumentMapUnit.job_result_id == scope.job_result_id) - ) - ) - unit_ids = [unit.id for unit in units] - frequencies: dict[str, dict[str, int]] = {"path": {}, "content": {}} - unit_frequencies: dict[str, dict[str, dict[str, int]]] = {} - if unit_ids: - for map_unit_id, channel, token, frequency in db.execute( - select( - DocumentMapUnitToken.map_unit_id, - DocumentMapUnitToken.channel, - DocumentMapUnitToken.token, - DocumentMapUnitToken.frequency, - ).where(DocumentMapUnitToken.map_unit_id.in_(unit_ids)) - ).all(): - channel_key = str(channel) - if channel_key in frequencies: - token_key = str(token) - frequency_value = int(frequency) - frequencies[channel_key][token_key] = ( - frequencies[channel_key].get(token_key, 0) + frequency_value - ) - unit_frequencies.setdefault(str(map_unit_id), {}).setdefault( - channel_key, {} - )[token_key] = frequency_value - return { - "document_id": scope.document_id, - "job_result_id": scope.job_result_id, - "unit_count": len(units), - "path_token_count": sum(int(unit.path_token_count or 0) for unit in units), - "content_token_count": sum( - int(unit.content_token_count or 0) for unit in units - ), - "token_frequencies": frequencies, - "unit_frequencies": unit_frequencies, - } - - -def persist_revision_serving_state( - db: Session, - *, - scope: DocumentPublicationScope, -) -> None: - """Replace manifest and statistics rows for one revision atomically.""" - manifest_payload = build_revision_serving_payload(db, scope=scope) - statistics_payload = build_revision_statistics_payload(db, scope=scope) - manifest_bytes, manifest_checksum, manifest_version = encode_serving_manifest( - manifest_payload - ) - statistics_bytes, statistics_checksum, statistics_version = encode_serving_manifest( - statistics_payload - ) - db.execute( - delete(RetrievalServingRevisionManifest) - .where(RetrievalServingRevisionManifest.document_id == scope.document_id) - .where(RetrievalServingRevisionManifest.job_result_id == scope.job_result_id) - ) - db.execute( - delete(RetrievalServingRevisionStat) - .where(RetrievalServingRevisionStat.document_id == scope.document_id) - .where(RetrievalServingRevisionStat.job_result_id == scope.job_result_id) - ) - db.add( - RetrievalServingRevisionManifest( - user_id=scope.user_id, - namespace=scope.namespace, - document_id=scope.document_id, - job_result_id=scope.job_result_id, - format_version=manifest_version, - payload_zlib=manifest_bytes, - checksum=manifest_checksum, - ) - ) - db.add( - RetrievalServingRevisionStat( - user_id=scope.user_id, - namespace=scope.namespace, - document_id=scope.document_id, - job_result_id=scope.job_result_id, - format_version=statistics_version, - payload_zlib=statistics_bytes, - checksum=statistics_checksum, - ) - ) - - -def rebuild_namespace_serving_statistics( - db: Session, - *, - user_id: str, - namespace: str, -) -> int: - """Recompute namespace aggregates from active current revisions. - - Callers hold the namespace generation lock. The aggregate is prepared for - the generation that the caller will publish next. - """ - generation = db.execute( - select(RetrievalNamespaceGeneration) - .where(RetrievalNamespaceGeneration.user_id == user_id) - .where(RetrievalNamespaceGeneration.namespace == namespace) - .with_for_update() - ).scalar_one() - target_generation = int(generation.generation) + 1 - revisions = { - (str(document_id), str(job_result_id)) - for document_id, job_result_id in db.execute( - select(Document.document_id, Document.current_job_result_id) - .where(Document.user_id == user_id) - .where(Document.namespace == namespace) - .where(Document.status == "active") - .where(Document.current_job_result_id.is_not(None)) - ).all() - if document_id and job_result_id - } - aggregate: dict[str, Any] = { - "document_count": 0, - "unit_count": 0, - "path_token_count": 0, - "content_token_count": 0, - "token_frequencies": {"path": {}, "content": {}}, - } - document_frequencies: dict[tuple[str, str], int] = {} - for row in db.scalars( - select(RetrievalServingRevisionStat) - .where(RetrievalServingRevisionStat.user_id == user_id) - .where(RetrievalServingRevisionStat.namespace == namespace) - ): - if (row.document_id, row.job_result_id) not in revisions: - continue - payload = decode_serving_manifest( - row.payload_zlib, - checksum=row.checksum, - format_version=row.format_version, - ) - aggregate["document_count"] += 1 - aggregate["unit_count"] += int(payload.get("unit_count", 0)) - aggregate["path_token_count"] += int(payload.get("path_token_count", 0)) - aggregate["content_token_count"] += int(payload.get("content_token_count", 0)) - token_frequencies = payload.get("token_frequencies", {}) - if not isinstance(token_frequencies, dict): - continue - for channel, values in token_frequencies.items(): - if channel not in aggregate["token_frequencies"] or not isinstance( - values, dict - ): - continue - for token, value in values.items(): - token_key = str(token) - aggregate["token_frequencies"][channel][token_key] = aggregate[ - "token_frequencies" - ][channel].get(token_key, 0) + int(value) - if int(value) > 0: - key = (str(channel), token_key) - document_frequencies[key] = document_frequencies.get(key, 0) + 1 - - encoded, checksum, _version = encode_serving_manifest(aggregate) - namespace_stat = db.execute( - select(RetrievalNamespaceStat) - .where(RetrievalNamespaceStat.user_id == user_id) - .where(RetrievalNamespaceStat.namespace == namespace) - ).scalar_one_or_none() - if namespace_stat is None: - db.add( - RetrievalNamespaceStat( - user_id=user_id, - namespace=namespace, - generation=target_generation, - payload_zlib=encoded, - checksum=checksum, - ) - ) - else: - namespace_stat.generation = target_generation - namespace_stat.payload_zlib = encoded - namespace_stat.checksum = checksum - db.execute( - delete(RetrievalNamespaceTokenStat) - .where(RetrievalNamespaceTokenStat.user_id == user_id) - .where(RetrievalNamespaceTokenStat.namespace == namespace) - ) - db.add_all( - [ - RetrievalNamespaceTokenStat( - user_id=user_id, - namespace=namespace, - generation=target_generation, - channel=channel, - token_hash=hashlib.sha256(token.encode("utf-8")).hexdigest(), - document_frequency=frequency, - ) - for (channel, token), frequency in document_frequencies.items() - ] - ) - db.flush() - return target_generation - - -def _connection_target_ids(metadata: Any) -> list[str]: - if not isinstance(metadata, dict): - return [] - connections = metadata.get("connect_to") - if not isinstance(connections, list): - return [] - return [ - target - for connection in connections - if isinstance(connection, dict) - for target in [str(connection.get("target") or "").strip()] - if target - ] - - -def encode_serving_manifest(payload: dict[str, Any]) -> tuple[bytes, str, int]: - """Return compressed canonical JSON, checksum, and format version.""" - canonical_payload = json.dumps( - payload, - ensure_ascii=False, - separators=(",", ":"), - sort_keys=True, - ).encode("utf-8") - checksum = hashlib.sha256(canonical_payload).hexdigest() - return ( - zlib.compress(canonical_payload), - checksum, - SERVING_MANIFEST_FORMAT_VERSION, - ) - - -def decode_serving_manifest( - payload_zlib: bytes, - *, - checksum: str, - format_version: int, -) -> dict[str, Any]: - """Validate and decode one persisted serving manifest.""" - if format_version != SERVING_MANIFEST_FORMAT_VERSION: - raise ValueError(f"unsupported serving manifest version: {format_version}") - - try: - canonical_payload = zlib.decompress(payload_zlib) - except zlib.error as exc: - raise ValueError("invalid serving manifest compression") from exc - - actual_checksum = hashlib.sha256(canonical_payload).hexdigest() - if actual_checksum != checksum: - raise ValueError("serving manifest checksum mismatch") - - try: - decoded = json.loads(canonical_payload.decode("utf-8")) - except (UnicodeDecodeError, json.JSONDecodeError) as exc: - raise ValueError("invalid serving manifest JSON") from exc - if not isinstance(decoded, dict): - raise ValueError("serving manifest payload must be an object") - return decoded From 319ccf7e5e7e0ff9e41bf65596d7531eb659871a Mon Sep 17 00:00:00 2001 From: suguanYang Date: Mon, 31 Aug 2026 11:06:53 +0800 Subject: [PATCH 04/10] Revert "Merge pull request #350 from Ontos-AI/perf/wangbinqi/mapnav-timing-instrumentation" This reverts commit 9b355e5b7d0c66d68e11a49de118e3cd9d063589, reversing changes made to d796ff475161890323cee42b260d27a076a74a05. --- .../services/retrieval/nav/nav_agent.py | 27 +------ .../shared/services/retrieval/nav/nav_llm.py | 74 +++++++------------ .../services/retrieval/nav/nav_orchestrate.py | 37 ++-------- 3 files changed, 34 insertions(+), 104 deletions(-) diff --git a/packages/shared-python/shared/services/retrieval/nav/nav_agent.py b/packages/shared-python/shared/services/retrieval/nav/nav_agent.py index 197af5af..f4d0d12d 100644 --- a/packages/shared-python/shared/services/retrieval/nav/nav_agent.py +++ b/packages/shared-python/shared/services/retrieval/nav/nav_agent.py @@ -1,8 +1,7 @@ from __future__ import annotations -import logging -import os import time +import os from typing import Any, List, Optional, Sequence, Tuple from ._compat import AgentStep, EpisodeResult @@ -38,7 +37,6 @@ # Back-compat aliases for tests / callers. _evidence_owner_section_id = evidence_owner_section_id _unit_score_for_evidence_chunk = unit_score_for_evidence_chunk -_logger = logging.getLogger(__name__) def _chunks_to_retrieved_nodes(chunks: List[Chunk]) -> List[str]: @@ -400,7 +398,6 @@ def _run_nav_episode_body( state = NavState(doc_id=episode_doc, query=query, task_type=task_type) steps: List[AgentStep] = [] - map_started = time.perf_counter() if namespace_mode: section_ids = list(ts.sections_for_doc("")) state.map_scores, state.unit_scores = compute_corpus_map_and_unit_scores( @@ -411,12 +408,6 @@ def _run_nav_episode_body( state.map_scores, state.unit_scores = compute_map_and_unit_scores( ts, doc_id=episode_doc, query=query, root_ids=section_ids ) - _logger.info( - "retrieval mapnav phase=map_scoring seconds=%.3f documents=%d sections=%d", - time.perf_counter() - map_started, - len(corpus_ids), - len(section_ids), - ) state.highlight_ids = select_map_highlights( state.unit_scores, k=int(cfg.collect_top_k) ) @@ -443,11 +434,6 @@ def _run_nav_episode_body( }, t0=plan_t0), ) ) - _logger.info( - "retrieval mapnav phase=planner seconds=%.3f subgoals=%d", - time.perf_counter() - plan_t0, - len(retrieval_plan.subgoals), - ) # Checklist: wave orchestration; navigate mode: classic single navigate. if cfg.is_checklist and state.retrieval_plan is not None: @@ -466,11 +452,6 @@ def _run_nav_episode_body( detail=stamp_step_detail(orch_detail, t0=orch_t0), ) ) - _logger.info( - "retrieval mapnav phase=orchestration seconds=%.3f waves=%d", - time.perf_counter() - orch_t0, - len(orch_detail.get("waves", [])), - ) else: navigate( ts, @@ -483,7 +464,6 @@ def _run_nav_episode_body( steps_out=steps, ) - evidence_started = time.perf_counter() fill = pack_nav_evidence( _dedupe_scored(list(state.collected)), ts, @@ -491,11 +471,6 @@ def _run_nav_episode_body( cfg, budget_chars=budget_chars, ) - _logger.info( - "retrieval mapnav phase=evidence_pack seconds=%.3f chunks=%d", - time.perf_counter() - evidence_started, - len(fill.kept_chunks), - ) scored_chunks = list(fill.scored_chunks) retrieval_seconds = time.perf_counter() - retrieval_t0 composed = "" diff --git a/packages/shared-python/shared/services/retrieval/nav/nav_llm.py b/packages/shared-python/shared/services/retrieval/nav/nav_llm.py index 767c1675..f8ba3187 100644 --- a/packages/shared-python/shared/services/retrieval/nav/nav_llm.py +++ b/packages/shared-python/shared/services/retrieval/nav/nav_llm.py @@ -20,9 +20,7 @@ from __future__ import annotations -import logging import os -import time from contextlib import contextmanager from contextvars import ContextVar from typing import Any, Callable, Dict, Iterator, Optional, Sequence @@ -32,7 +30,6 @@ NavChatBackend = Callable[..., Dict[str, Any]] _backend: Optional[NavChatBackend] = None -_logger = logging.getLogger(__name__) _DS_DEFAULT_MODEL = "deepseek-v4-flash" _DEFAULT_PLANNER_THINK_MAX = 16384 @@ -193,34 +190,25 @@ def nav_chat( raise NavTokenLimit(used=nav_tokens_used(), limit=nav_token_limit()) merged_extra = _merge_thinking_extra(extra, role=thinking_role, model=model) - call_started = time.perf_counter() if _backend is not None: - try: - result = _backend( - purpose=purpose, - messages=list(messages), - model=model, - temperature=temperature, - max_tokens=max_tokens, - response_format=response_format, - extra=merged_extra, - thinking_role=thinking_role, - context=context, - api_key_env=api_key_env, - base_url_env=base_url_env, - timeout=timeout, - usage_tag=usage_tag, - ) - record_episode_tokens((result or {}).get("usage")) - return result - finally: - _logger.info( - "retrieval mapnav llm_call purpose=%s role=%s seconds=%.3f", - purpose, - thinking_role, - time.perf_counter() - call_started, - ) + result = _backend( + purpose=purpose, + messages=list(messages), + model=model, + temperature=temperature, + max_tokens=max_tokens, + response_format=response_format, + extra=merged_extra, + thinking_role=thinking_role, + context=context, + api_key_env=api_key_env, + base_url_env=base_url_env, + timeout=timeout, + usage_tag=usage_tag, + ) + record_episode_tokens((result or {}).get("usage")) + return result from ._compat import cached_chat_completion # type: ignore from ._compat import ( # type: ignore @@ -242,24 +230,16 @@ def nav_chat( f"(model={model!r}; need DS_KEY for deepseek-* or OPENAI_API_KEY)." ) client = make_openai_client(api_key=key, base_url=base_url, timeout=timeout) - try: - cached = cached_chat_completion( - client, - purpose=purpose, - model=model, - messages=list(messages), - temperature=temperature, - max_tokens=max_tokens, - response_format=response_format, - extra=merged_extra, - ) - finally: - _logger.info( - "retrieval mapnav llm_call purpose=%s role=%s seconds=%.3f", - purpose, - thinking_role, - time.perf_counter() - call_started, - ) + cached = cached_chat_completion( + client, + purpose=purpose, + model=model, + messages=list(messages), + temperature=temperature, + max_tokens=max_tokens, + response_format=response_format, + extra=merged_extra, + ) if usage_tag: record_usage(usage_tag, cached.get("usage")) record_episode_tokens(cached.get("usage")) diff --git a/packages/shared-python/shared/services/retrieval/nav/nav_orchestrate.py b/packages/shared-python/shared/services/retrieval/nav/nav_orchestrate.py index 2e82be08..91f45e1c 100644 --- a/packages/shared-python/shared/services/retrieval/nav/nav_orchestrate.py +++ b/packages/shared-python/shared/services/retrieval/nav/nav_orchestrate.py @@ -8,14 +8,14 @@ from __future__ import annotations -import logging +from .nav_token_budget import stamp_step_detail + import re import time from contextlib import contextmanager from dataclasses import asdict from typing import Any, Dict, Iterator, List, Optional, Sequence, Set, Tuple -from .nav_token_budget import stamp_step_detail from .nav_navigate import navigate from .nav_plan import ( RetrievalPlan, @@ -29,7 +29,6 @@ from .nav_verify import apply_bindings_from_result, build_subgoal_result _SLOT_STRIP_RE = re.compile(r"\{\{\s*[^}]+\s*\}\}") -_logger = logging.getLogger(__name__) def ready_subgoal_ids( @@ -484,18 +483,8 @@ def _run_one(sid: str, working_state: NavState, out_steps: Optional[List[Any]]) ) # Serial wave execution (parallel fan-out retired with ThreadPoolExecutor). - wave_started = time.perf_counter() for sid in ready: - subgoal_started = time.perf_counter() - try: - outputs.append(_run_one(sid, state, steps_out)) - finally: - _logger.info( - "retrieval mapnav harvest subgoal=%s wave=%d seconds=%.3f", - sid, - wave_idx, - time.perf_counter() - subgoal_started, - ) + outputs.append(_run_one(sid, state, steps_out)) # Bookkeeping shared by both decision paths. for item in outputs: @@ -519,17 +508,9 @@ def _run_one(sid: str, working_state: NavState, out_steps: Optional[List[Any]]) state.subgoal_attempt_counts.get(sid, 0) ) + 1 - control_started = time.perf_counter() - try: - control_detail = _apply_plan_control( - ts, state, config, plan=plan, outputs=outputs, by_id=by_id, steps_out=steps_out - ) - finally: - _logger.info( - "retrieval mapnav plan_control wave=%d seconds=%.3f", - wave_idx, - time.perf_counter() - control_started, - ) + control_detail = _apply_plan_control( + ts, state, config, plan=plan, outputs=outputs, by_id=by_id, steps_out=steps_out + ) wave_detail["plan_control"] = control_detail replan_requested = bool(control_detail.get("replan")) if control_detail.get("done"): @@ -544,12 +525,6 @@ def _run_one(sid: str, working_state: NavState, out_steps: Optional[List[Any]]) ) ) summary["waves"].append(wave_detail) - _logger.info( - "retrieval mapnav wave=%d ready=%d seconds=%.3f", - wave_idx, - len(ready), - time.perf_counter() - wave_started, - ) if replan_requested: cap = int(getattr(config, "max_replans", 0) or 0) From ece74d6a37b6f9eda1928400ccb4bceabea9f266 Mon Sep 17 00:00:00 2001 From: suguanYang Date: Mon, 31 Aug 2026 11:07:14 +0800 Subject: [PATCH 05/10] Revert "Merge pull request #344 from Ontos-AI/fix/wangbinqi/normalize-native-dsn" This reverts commit b051c0813819ddd2f1f6078771dc453ab0a3ea42, reversing changes made to a87134fdec545df5bd7d8252d88e8dfaa70aa70d. --- .../1d2e3f4a5b6c_add_document_map_units.py | 122 ----- apps/api/scripts/backfill_map_unit_indexes.py | 113 ---- ...etrieval_lazy_snapshot_quality_contract.py | 158 +----- .../test_retrieval_map_unit_index_contract.py | 423 --------------- deploy/ecs/README.md | 28 - .../shared/models/database/__init__.py | 6 - .../shared/models/database/document.py | 157 +----- .../services/retrieval/map_unit_index.py | 163 ------ .../services/retrieval/nav/knowhere_hybrid.py | 287 +++-------- .../services/retrieval/nav/nav_hierarchy.py | 49 +- .../services/retrieval/nav/nav_knowhere.py | 482 +----------------- .../services/retrieval/nav/nav_map_scores.py | 178 ++----- .../services/retrieval/nav/nav_orchestrate.py | 67 +-- .../services/retrieval/publication_content.py | 24 +- 14 files changed, 138 insertions(+), 2119 deletions(-) delete mode 100644 apps/api/alembic/versions/1d2e3f4a5b6c_add_document_map_units.py delete mode 100644 apps/api/scripts/backfill_map_unit_indexes.py delete mode 100644 apps/api/tests/contract/test_retrieval_map_unit_index_contract.py delete mode 100644 packages/shared-python/shared/services/retrieval/map_unit_index.py diff --git a/apps/api/alembic/versions/1d2e3f4a5b6c_add_document_map_units.py b/apps/api/alembic/versions/1d2e3f4a5b6c_add_document_map_units.py deleted file mode 100644 index 79aa7583..00000000 --- a/apps/api/alembic/versions/1d2e3f4a5b6c_add_document_map_units.py +++ /dev/null @@ -1,122 +0,0 @@ -"""Add revision-pinned map-nav score units.""" - -from __future__ import annotations - -import sqlalchemy as sa -from alembic import op - -revision = "1d2e3f4a5b6c" -down_revision = "0c1d2e3f4a5b" -branch_labels = None -depends_on = None - - -def upgrade() -> None: - bind = op.get_bind() - inspector = sa.inspect(bind) - if not inspector.has_table("document_map_unit_indexes"): - op.create_table( - "document_map_unit_indexes", - sa.Column("id", sa.String(length=100), nullable=False), - sa.Column("document_id", sa.String(length=36), nullable=False), - sa.Column("job_result_id", sa.String(length=36), nullable=False), - sa.Column("format_version", sa.Integer(), nullable=False), - sa.Column("unit_count", sa.Integer(), nullable=False), - sa.Column("token_count", sa.Integer(), nullable=False), - sa.Column("created_at", sa.DateTime(), nullable=False), - sa.ForeignKeyConstraint( - ["document_id"], ["documents.document_id"], ondelete="CASCADE" - ), - sa.ForeignKeyConstraint( - ["job_result_id"], ["job_results.id"], ondelete="CASCADE" - ), - sa.PrimaryKeyConstraint("id"), - sa.UniqueConstraint( - "document_id", - "job_result_id", - name="uq_document_map_unit_indexes_revision", - ), - ) - if not inspector.has_table("document_map_units"): - op.create_table( - "document_map_units", - sa.Column("id", sa.String(length=160), nullable=False), - sa.Column("document_id", sa.String(length=36), nullable=False), - sa.Column("job_result_id", sa.String(length=36), nullable=False), - sa.Column("unit_id", sa.String(length=128), nullable=False), - sa.Column("section_id", sa.String(length=36), nullable=False), - sa.Column("unit_kind", sa.String(length=32), nullable=False), - sa.Column("path_token_count", sa.Integer(), nullable=False), - sa.Column("content_token_count", sa.Integer(), nullable=False), - sa.Column("term_search_text_lower", sa.Text(), nullable=False), - sa.Column("sort_order", sa.Integer(), nullable=False), - sa.Column("created_at", sa.DateTime(), nullable=False), - sa.ForeignKeyConstraint( - ["document_id"], ["documents.document_id"], ondelete="CASCADE" - ), - sa.ForeignKeyConstraint( - ["job_result_id"], ["job_results.id"], ondelete="CASCADE" - ), - sa.PrimaryKeyConstraint("id"), - ) - if not inspector.has_table("document_map_unit_tokens"): - op.create_table( - "document_map_unit_tokens", - sa.Column("id", sa.String(length=36), nullable=False), - sa.Column("map_unit_id", sa.String(length=160), nullable=False), - sa.Column("channel", sa.String(length=16), nullable=False), - sa.Column("token", sa.Text(), nullable=False), - sa.Column("token_hash", sa.String(length=64), nullable=False), - sa.Column("frequency", sa.Integer(), nullable=False), - sa.ForeignKeyConstraint( - ["map_unit_id"], ["document_map_units.id"], ondelete="CASCADE" - ), - sa.PrimaryKeyConstraint("id"), - ) - inspector = sa.inspect(bind) - indexes = { - item["name"] for item in inspector.get_indexes("document_map_unit_indexes") - } - if "idx_document_map_unit_indexes_revision" not in indexes: - op.create_index( - "idx_document_map_unit_indexes_revision", - "document_map_unit_indexes", - ["document_id", "job_result_id"], - ) - indexes = {item["name"] for item in inspector.get_indexes("document_map_units")} - if "idx_document_map_units_revision_order" not in indexes: - op.create_index( - "idx_document_map_units_revision_order", - "document_map_units", - ["document_id", "job_result_id", "sort_order", "unit_id"], - ) - if "idx_document_map_units_section" not in indexes: - op.create_index( - "idx_document_map_units_section", "document_map_units", ["section_id"] - ) - indexes = { - item["name"] for item in inspector.get_indexes("document_map_unit_tokens") - } - if "idx_document_map_unit_tokens_lookup" not in indexes: - op.create_index( - "idx_document_map_unit_tokens_lookup", - "document_map_unit_tokens", - ["channel", "token_hash", "map_unit_id"], - ) - if "idx_document_map_unit_tokens_unit" not in indexes: - op.create_index( - "idx_document_map_unit_tokens_unit", - "document_map_unit_tokens", - ["map_unit_id", "channel"], - ) - - -def downgrade() -> None: - existing_tables = set(sa.inspect(op.get_bind()).get_table_names()) - for table_name in ( - "document_map_unit_tokens", - "document_map_units", - "document_map_unit_indexes", - ): - if table_name in existing_tables: - op.drop_table(table_name) diff --git a/apps/api/scripts/backfill_map_unit_indexes.py b/apps/api/scripts/backfill_map_unit_indexes.py deleted file mode 100644 index c0301d38..00000000 --- a/apps/api/scripts/backfill_map_unit_indexes.py +++ /dev/null @@ -1,113 +0,0 @@ -"""Backfill persisted MAP-NAV lexical indexes for existing revisions. - -The migration creates empty derived tables intentionally. Run this command -after deployment with ``--apply`` so each revision is rebuilt and committed -independently; without ``--apply`` it is a read-only inventory. -""" - -# ruff: noqa: E402 - -from __future__ import annotations - -import argparse -import os -import sys -from pathlib import Path - - -def _resolve_shared_root(api_root: Path) -> Path: - """Resolve the shared package in source checkouts and runtime images.""" - runtime_shared_root = api_root / "packages" / "shared-python" - if runtime_shared_root.is_dir(): - return runtime_shared_root - - repository_shared_root = api_root.parents[1] / "packages" / "shared-python" - if repository_shared_root.is_dir(): - return repository_shared_root - - raise RuntimeError(f"Could not locate shared-python package from {api_root}") - - -def _bootstrap_python_path() -> None: - api_root = Path(__file__).resolve().parents[1] - shared_root = _resolve_shared_root(api_root) - for path in (api_root, shared_root): - value = os.fspath(path) - if value not in sys.path: - sys.path.insert(0, value) - - -_bootstrap_python_path() - -from sqlalchemy import select - -from shared.core.database_sync import get_sync_session_factory -from shared.models.database.document import Document -from shared.services.retrieval.map_unit_index import replace_document_map_units -from shared.services.retrieval.publication_models import DocumentPublicationScope - - -def _build_parser() -> argparse.ArgumentParser: - parser = argparse.ArgumentParser( - description="Backfill MAP-NAV indexes for current document revisions." - ) - parser.add_argument( - "--apply", - action="store_true", - help="Build and commit each current revision index.", - ) - parser.add_argument( - "--document-id", default="", help="Limit the backfill to one document." - ) - return parser - - -def _load_documents(document_id: str) -> list[Document]: - session_factory = get_sync_session_factory() - with session_factory() as db: - statement = select(Document).where(Document.current_job_result_id.is_not(None)) - normalized_document_id = document_id.strip() - if normalized_document_id: - statement = statement.where(Document.document_id == normalized_document_id) - return list(db.scalars(statement).all()) - - -def backfill_map_unit_indexes(*, apply: bool, document_id: str = "") -> int: - documents = _load_documents(document_id) - if not apply: - for document in documents: - print( - f"would backfill document={document.document_id} revision={document.current_job_result_id}" - ) - return len(documents) - - session_factory = get_sync_session_factory() - for document in documents: - job_result_id = document.current_job_result_id - if not job_result_id: - continue - scope = DocumentPublicationScope( - user_id=document.user_id, - namespace=document.namespace, - document_id=document.document_id, - job_result_id=job_result_id, - source_file_name=str(document.source_file_name or ""), - ) - with session_factory() as db: - replace_document_map_units(db, scope=scope) - db.commit() - print(f"backfilled document={document.document_id} revision={job_result_id}") - return len(documents) - - -def main() -> None: - arguments = _build_parser().parse_args() - count = backfill_map_unit_indexes( - apply=bool(arguments.apply), document_id=str(arguments.document_id) - ) - action = "backfilled" if arguments.apply else "found" - print(f"{action} revisions={count}") - - -if __name__ == "__main__": - main() diff --git a/apps/api/tests/contract/test_retrieval_lazy_snapshot_quality_contract.py b/apps/api/tests/contract/test_retrieval_lazy_snapshot_quality_contract.py index ae6d78ed..cfb82c9e 100644 --- a/apps/api/tests/contract/test_retrieval_lazy_snapshot_quality_contract.py +++ b/apps/api/tests/contract/test_retrieval_lazy_snapshot_quality_contract.py @@ -1,7 +1,7 @@ from __future__ import annotations from dataclasses import dataclass -from collections.abc import Mapping, Sequence +from collections.abc import Sequence from typing import Any from shared.services.retrieval.nav.nav_hierarchy import ProviderToolSpace @@ -16,13 +16,11 @@ from shared.services.retrieval.nav.nav_map_scores import ( build_score_units, compute_corpus_map_and_unit_scores, - compute_corpus_map_and_unit_scores_many, ) from shared.services.retrieval.nav.knowhere_hybrid import ( ScoreUnitRow, score_rows_hybrid_all, score_unit_stream_hybrid_all, - score_unit_stream_hybrid_many, ) @@ -30,21 +28,6 @@ class _FakeChunkStore: units_by_section: dict[str, list[UnitRow]] document_loads: int = 0 - batch_loads: int = 0 - - def load_documents_units( - self, - section_ids_by_document: Mapping[str, Sequence[str]], - ) -> dict[str, list[UnitRow]]: - self.batch_loads += 1 - return { - document_id: [ - unit - for section_id in section_ids - for unit in self.units_by_section.get(section_id, ()) - ] - for document_id, section_ids in section_ids_by_document.items() - } def load_document_units( self, @@ -135,63 +118,6 @@ def _providers() -> tuple[ProviderToolSpace, ProviderToolSpace, _FakeChunkStore] return ProviderToolSpace(eager), ProviderToolSpace(lazy), store -def _multi_document_providers() -> tuple[ - ProviderToolSpace, - ProviderToolSpace, - _FakeChunkStore, -]: - first_sections = [ - SectionRow("root-a", None, "Root A", "Root A", 0, "", 0), - SectionRow("leaf-a", "root-a", "Root A / Leaf A", "Leaf A", 1, "", 1), - ] - second_sections = [ - SectionRow("root-b", None, "Root B", "Root B", 0, "", 0), - SectionRow("leaf-b", "root-b", "Root B / Leaf B", "Leaf B", 1, "", 1), - ] - first_unit = UnitRow("chunk-a", "leaf-a", "text", "alpha evidence", 1) - second_unit = UnitRow("chunk-b", "leaf-b", "text", "beta evidence", 1) - eager = NamespaceKnowhereProvider( - [ - KnowhereProvider( - doc_id="doc-a", - sections=first_sections, - units=[first_unit], - ), - KnowhereProvider( - doc_id="doc-b", - sections=second_sections, - units=[second_unit], - ), - ], - titles={"doc-a": "Document A", "doc-b": "Document B"}, - ) - store = _FakeChunkStore( - { - "leaf-a": [first_unit], - "leaf-b": [second_unit], - } - ) - lazy = NamespaceKnowhereProvider( - [ - LazyKnowhereProvider( - doc_id="doc-a", - sections=first_sections, - chunk_store=store, - known_chunk_ids=[first_unit.chunk_id], - ), - LazyKnowhereProvider( - doc_id="doc-b", - sections=second_sections, - chunk_store=store, - known_chunk_ids=[second_unit.chunk_id], - ), - ], - titles={"doc-a": "Document A", "doc-b": "Document B"}, - chunk_owner_by_id={"chunk-a": "doc-a", "chunk-b": "doc-b"}, - ) - return ProviderToolSpace(eager), ProviderToolSpace(lazy), store - - def test_lazy_provider_preserves_score_units_and_scores() -> None: eager, lazy, store = _providers() @@ -281,88 +207,6 @@ def test_streaming_scorer_preserves_duplicate_id_eager_semantics() -> None: assert score_unit_stream_hybrid_all(lambda: rows, "alpha beta") == eager_scores -def test_streaming_scorer_scores_multiple_queries_with_one_corpus_read() -> None: - rows: list[ScoreUnitRow] = [ - { - "chunk_id": "unit-a", - "path_search_text": "root alpha", - "content_search_text": "alpha alpha evidence", - "term_search_text": "alpha alpha evidence root", - }, - { - "chunk_id": "unit-b", - "path_search_text": "root beta", - "content_search_text": "beta evidence", - "term_search_text": "beta evidence root", - }, - { - "chunk_id": "unit-c", - "path_search_text": "root common", - "content_search_text": "common evidence", - "term_search_text": "common evidence root", - }, - ] - queries: list[str] = ["alpha evidence", "beta evidence"] - expected = { - query: score_unit_stream_hybrid_all(lambda: rows, query) - for query in queries - } - read_count: int = 0 - - def unit_factory() -> Sequence[ScoreUnitRow]: - nonlocal read_count - read_count += 1 - return rows - - assert score_unit_stream_hybrid_many(unit_factory, queries) == expected - assert read_count == 1 - - -def test_corpus_map_scores_multiple_queries_with_one_lazy_load() -> None: - eager, lazy, store = _providers() - queries: list[str] = ["alpha retrieval", "supporting image"] - expected = { - query: compute_corpus_map_and_unit_scores( - eager, - doc_ids=["doc"], - query=query, - ) - for query in queries - } - - store.document_loads = 0 - store.batch_loads = 0 - actual = compute_corpus_map_and_unit_scores_many( - lazy, - doc_ids=["doc"], - queries=queries, - ) - - assert actual == expected - assert store.batch_loads == 1 - assert store.document_loads == 0 - - -def test_corpus_map_batches_multiple_documents_without_score_drift() -> None: - eager, lazy, store = _multi_document_providers() - queries: list[str] = ["alpha evidence", "beta evidence"] - expected = compute_corpus_map_and_unit_scores_many( - eager, - doc_ids=["doc-a", "doc-b"], - queries=queries, - ) - - actual = compute_corpus_map_and_unit_scores_many( - lazy, - doc_ids=["doc-a", "doc-b"], - queries=queries, - ) - - assert actual == expected - assert store.batch_loads == 1 - assert store.document_loads == 0 - - def test_native_chunk_store_strips_async_driver_from_database_url( monkeypatch: Any, ) -> None: diff --git a/apps/api/tests/contract/test_retrieval_map_unit_index_contract.py b/apps/api/tests/contract/test_retrieval_map_unit_index_contract.py deleted file mode 100644 index 2f2027b0..00000000 --- a/apps/api/tests/contract/test_retrieval_map_unit_index_contract.py +++ /dev/null @@ -1,423 +0,0 @@ -from __future__ import annotations - -from collections.abc import Callable, Mapping, Sequence -from contextlib import AbstractAsyncContextManager -from uuid import uuid4 - -from httpx import AsyncClient -from sqlalchemy import delete, select - -from shared.models.database.document import ( - DocumentMapUnit, - DocumentMapUnitIndex, - DocumentMapUnitToken, -) -from shared.services.retrieval.nav.nav_hierarchy import ProviderToolSpace -from shared.services.retrieval.nav.nav_map_scores import ( - build_score_units, - compute_corpus_map_and_unit_scores, -) -from shared.services.retrieval.nav.nav_knowhere import ( - KnowhereProvider, - LazyKnowhereProvider, - NamespaceKnowhereProvider, - ReadOnlyChunkStore, - SectionRow, - UnitRow, -) -from shared.services.retrieval.nav_snapshot import load_nav_snapshot -from shared.services.retrieval.publication_content import ( - replace_document_revision_content, -) -from shared.services.retrieval.publication_models import DocumentPublicationScope -from tests.support.contract_database import ContractDatabase -from tests.support.retrieval_snapshot_support import contract_db_session - -_USER_ID = "local-dev-user" - - -class _IncompleteIndexStore: - """Minimal lazy store whose incomplete index forces legacy scoring.""" - - def __init__(self, units_by_section: Mapping[str, Sequence[UnitRow]]) -> None: - self.units_by_section = { - str(section_id): list(units) - for section_id, units in units_by_section.items() - } - self.persisted_loads = 0 - self.batch_loads = 0 - - def load_persisted_score_corpus( - self, - document_ids: Sequence[str], - allowed_section_ids_by_document: Mapping[str, Sequence[str]], - queries: Sequence[str], - ) -> None: - del document_ids, allowed_section_ids_by_document, queries - self.persisted_loads += 1 - return None - - def load_documents_units( - self, - section_ids_by_document: Mapping[str, Sequence[str]], - ) -> dict[str, list[UnitRow]]: - self.batch_loads += 1 - return { - str(document_id): [ - unit - for section_id in section_ids - for unit in self.units_by_section.get(str(section_id), ()) - ] - for document_id, section_ids in section_ids_by_document.items() - } - - def load_document_units( - self, - document_id: str, - section_ids: Sequence[str], - extra_chunk_ids_by_section: Mapping[str, Sequence[str]] | None = None, - ) -> list[UnitRow]: - del document_id, extra_chunk_ids_by_section - return [ - unit - for section_id in section_ids - for unit in self.units_by_section.get(str(section_id), ()) - ] - - def load_section_units( - self, - document_id: str, - section_id: str, - extra_chunk_ids: Sequence[str] = (), - ) -> list[UnitRow]: - del document_id, extra_chunk_ids - return list(self.units_by_section.get(str(section_id), ())) - - def close(self) -> None: - return None - - -async def test_published_map_units_preserve_scores_without_chunk_payload_reads( - developer_api_client_factory: Callable[ - [], AbstractAsyncContextManager[AsyncClient] - ], - monkeypatch, -) -> None: - identifier = uuid4().hex[:8] - namespace = f"map-unit-index-{identifier}" - document_id = f"doc_map_{identifier}" - job_id = f"job_map_{identifier}" - job_result_id = f"result_map_{identifier}" - async with developer_api_client_factory(): - await _seed_revision( - namespace=namespace, - document_id=document_id, - job_id=job_id, - job_result_id=job_result_id, - ) - scope = DocumentPublicationScope( - user_id=_USER_ID, - namespace=namespace, - document_id=document_id, - job_result_id=job_result_id, - source_file_name="indexed.pdf", - ) - chunks = [ - { - "chunk_id": "parent-a", - "type": "text", - "content": "common alpha parent evidence", - "path": "indexed.pdf/Root/Parent/intro-a", - "order": 1, - "metadata": {}, - }, - { - "chunk_id": "parent-b", - "type": "text", - "content": "common beta parent evidence", - "path": "indexed.pdf/Root/Parent/intro-b", - "order": 2, - "metadata": {}, - }, - { - "chunk_id": "leaf-a", - "type": "text", - "content": "common alpha leaf evidence", - "path": "indexed.pdf/Root/Parent/Leaf A/body", - "order": 3, - "metadata": {}, - }, - { - "chunk_id": "leaf-b", - "type": "text", - "content": "common beta leaf evidence", - "path": "indexed.pdf/Root/Parent/Leaf B/body", - "order": 4, - "metadata": {}, - }, - ] - async with contract_db_session() as db: - await db.run_sync( - lambda sync_db: replace_document_revision_content( - sync_db, - scope=scope, - chunks=chunks, - ) - ) - await db.commit() - - async with contract_db_session() as db: - eager_snapshot = await load_nav_snapshot( - db, - user_id=_USER_ID, - namespace=namespace, - ) - eager_toolspace = ProviderToolSpace(eager_snapshot.provider) - expected_units = build_score_units(eager_toolspace, document_id) - expected_scores = compute_corpus_map_and_unit_scores( - eager_toolspace, - doc_ids=[document_id], - query="common alpha", - ) - - async with contract_db_session() as db: - index = ( - await db.execute( - select(DocumentMapUnitIndex).where( - DocumentMapUnitIndex.document_id == document_id - ) - ) - ).scalar_one() - persisted_units = list( - ( - await db.execute( - select(DocumentMapUnit) - .where(DocumentMapUnit.document_id == document_id) - .order_by(DocumentMapUnit.sort_order) - ) - ).scalars() - ) - persisted_tokens = list( - ( - await db.execute( - select(DocumentMapUnitToken).where( - DocumentMapUnitToken.map_unit_id.in_( - [unit.id for unit in persisted_units] - ) - ) - ) - ).scalars() - ) - lazy_snapshot = await load_nav_snapshot( - db, - user_id=_USER_ID, - namespace=namespace, - lazy=True, - ) - - assert index.unit_count == len(expected_units) - assert [unit.unit_id for unit in persisted_units] == [ - str(unit["chunk_id"]) for unit in expected_units - ] - assert persisted_tokens - - def reject_payload_read( - _store: ReadOnlyChunkStore, - _section_ids_by_document: Mapping[str, Sequence[str]], - ) -> dict[str, list[UnitRow]]: - raise AssertionError("persisted map scoring loaded full chunk payloads") - - original_payload_loader = ReadOnlyChunkStore.load_documents_units - monkeypatch.setattr( - ReadOnlyChunkStore, - "load_documents_units", - reject_payload_read, - ) - actual_scores = compute_corpus_map_and_unit_scores( - ProviderToolSpace(lazy_snapshot.provider), - doc_ids=[document_id], - query="common alpha", - ) - monkeypatch.setattr( - ReadOnlyChunkStore, - "load_documents_units", - original_payload_loader, - ) - async with contract_db_session() as db: - token_id = ( - select(DocumentMapUnitToken.id) - .join( - DocumentMapUnit, - DocumentMapUnit.id == DocumentMapUnitToken.map_unit_id, - ) - .where(DocumentMapUnit.document_id == document_id) - .limit(1) - .scalar_subquery() - ) - await db.execute( - delete(DocumentMapUnitToken).where(DocumentMapUnitToken.id == token_id) - ) - await db.commit() - incomplete_snapshot = await load_nav_snapshot( - db, - user_id=_USER_ID, - namespace=namespace, - lazy=True, - ) - fallback_scores = compute_corpus_map_and_unit_scores( - ProviderToolSpace(incomplete_snapshot.provider), - doc_ids=[document_id], - query="common alpha", - ) - incomplete_snapshot.close() - lazy_snapshot.close() - eager_snapshot.close() - - assert actual_scores == expected_scores - assert fallback_scores == expected_scores - - -def test_incomplete_index_falls_back_for_duplicate_unit_ids() -> None: - first_sections = [ - SectionRow("root-a", None, "Root A", "Root A", 0, "", 0), - SectionRow("leaf-a", "root-a", "Root A / Leaf A", "Leaf A", 1, "", 1), - ] - second_sections = [ - SectionRow("root-b", None, "Root B", "Root B", 0, "", 0), - SectionRow("leaf-b", "root-b", "Root B / Leaf B", "Leaf B", 1, "", 1), - ] - first_unit = UnitRow("same-chunk", "leaf-a", "text", "alpha evidence", 1) - second_unit = UnitRow("same-chunk", "leaf-b", "text", "beta evidence", 1) - - eager = ProviderToolSpace( - NamespaceKnowhereProvider( - [ - KnowhereProvider( - doc_id="doc-a", sections=first_sections, units=[first_unit] - ), - KnowhereProvider( - doc_id="doc-b", sections=second_sections, units=[second_unit] - ), - ], - titles={"doc-a": "Document A", "doc-b": "Document B"}, - ) - ) - store = _IncompleteIndexStore( - {"leaf-a": [first_unit], "leaf-b": [second_unit]} - ) - lazy = ProviderToolSpace( - NamespaceKnowhereProvider( - [ - LazyKnowhereProvider( - doc_id="doc-a", - sections=first_sections, - chunk_store=store, - known_chunk_ids=[first_unit.chunk_id], - ), - LazyKnowhereProvider( - doc_id="doc-b", - sections=second_sections, - chunk_store=store, - known_chunk_ids=[second_unit.chunk_id], - ), - ], - titles={"doc-a": "Document A", "doc-b": "Document B"}, - chunk_owner_by_id={"same-chunk": "doc-a"}, - ) - ) - - expected = compute_corpus_map_and_unit_scores( - eager, doc_ids=["doc-a", "doc-b"], query="alpha beta" - ) - actual = compute_corpus_map_and_unit_scores( - lazy, doc_ids=["doc-a", "doc-b"], query="alpha beta" - ) - - assert actual == expected - assert store.persisted_loads == 1 - assert store.batch_loads == 1 - - -def test_titleless_leaf_has_identical_eager_and_lazy_path_scoring() -> None: - sections = [ - SectionRow("root", None, "Root", "Root", 0, "", 0), - SectionRow("leaf", "root", "Root / Leaf", "", 1, "", 1), - ] - unit = UnitRow("titleless-chunk", "leaf", "text", "alpha evidence", 1) - eager = ProviderToolSpace( - KnowhereProvider(doc_id="doc", sections=sections, units=[unit]) - ) - store = _IncompleteIndexStore({"leaf": [unit]}) - lazy = ProviderToolSpace( - LazyKnowhereProvider( - doc_id="doc", - sections=sections, - chunk_store=store, - known_chunk_ids=[unit.chunk_id], - ) - ) - - assert build_score_units(eager, "doc") == build_score_units(lazy, "doc") - assert compute_corpus_map_and_unit_scores( - eager, doc_ids=["doc"], query="alpha" - ) == compute_corpus_map_and_unit_scores( - lazy, doc_ids=["doc"], query="alpha" - ) - - -async def _seed_revision( - *, - namespace: str, - document_id: str, - job_id: str, - job_result_id: str, -) -> None: - await ContractDatabase.execute( - """ - INSERT INTO jobs ( - job_id, user_id, job_type, status, source_type, version, - webhook_enabled, created_at, updated_at, credits_charged, billing_status - ) VALUES ( - :job_id, :user_id, 'document_ingestion', 'done', 'file', 0, - false, NOW(), NOW(), 0, 'skipped' - ) - """, - {"job_id": job_id, "user_id": _USER_ID}, - ) - await ContractDatabase.execute( - """ - INSERT INTO documents ( - document_id, user_id, namespace, status, source_file_name, - parse_track, created_at, updated_at - ) VALUES ( - :document_id, :user_id, :namespace, 'active', - 'indexed.pdf', 'chunk', NOW(), NOW() - ) - """, - { - "document_id": document_id, - "user_id": _USER_ID, - "namespace": namespace, - }, - ) - await ContractDatabase.execute( - """ - INSERT INTO job_results ( - id, job_id, document_id, delivery_mode, created_at, updated_at - ) VALUES ( - :job_result_id, :job_id, :document_id, 'inline', NOW(), NOW() - ) - """, - { - "job_result_id": job_result_id, - "job_id": job_id, - "document_id": document_id, - }, - ) - await ContractDatabase.execute( - """ - UPDATE documents SET current_job_result_id = :job_result_id - WHERE document_id = :document_id - """, - {"job_result_id": job_result_id, "document_id": document_id}, - ) diff --git a/deploy/ecs/README.md b/deploy/ecs/README.md index dd0742b4..e19200af 100644 --- a/deploy/ecs/README.md +++ b/deploy/ecs/README.md @@ -69,34 +69,6 @@ workflow validates these resources, registers immutable image-digest task definitions, runs the production migration first, and then updates the ECS services. It does not create or delete AWS resources. -## Required post-deploy backfill - -The map-nav lexical-index migration creates the derived index tables, but it does -not rebuild indexes for revisions that already exist. Until those revisions are -backfilled, retrieval remains quality-preserving but uses the legacy scoring -path. Every release containing the map-nav index change must include the -following DevOps action in its release notification. - -Run the commands as a one-off container using the newly deployed API image and -the production database secret. Do not run them inside the long-lived API task. - -```bash -# Read-only inventory -python /app/scripts/backfill_map_unit_indexes.py - -# Optional canary: apply one affected document first -python /app/scripts/backfill_map_unit_indexes.py \ - --document-id \ - --apply - -# Apply to all current document revisions -python /app/scripts/backfill_map_unit_indexes.py --apply -``` - -The script commits each document revision independently and is safe to rerun. -Verify the canary retrieval before starting the full apply. New or republished -documents build their index automatically during publication. - ## Manual staging availability `.github/workflows/manage-staging.yml` exposes three manually dispatched diff --git a/packages/shared-python/shared/models/database/__init__.py b/packages/shared-python/shared/models/database/__init__.py index 5535cd9b..2d0e7d8a 100644 --- a/packages/shared-python/shared/models/database/__init__.py +++ b/packages/shared-python/shared/models/database/__init__.py @@ -13,9 +13,6 @@ from .document import ( Document, DocumentChunk, - DocumentMapUnit, - DocumentMapUnitIndex, - DocumentMapUnitToken, DocumentSection, GraphEdge, GraphNode, @@ -59,9 +56,6 @@ "Document", "DocumentSection", "DocumentChunk", - "DocumentMapUnit", - "DocumentMapUnitIndex", - "DocumentMapUnitToken", "DocumentPagePlan", "DemoMaterialization", "GraphNode", diff --git a/packages/shared-python/shared/models/database/document.py b/packages/shared-python/shared/models/database/document.py index 33dda2f9..b2381996 100644 --- a/packages/shared-python/shared/models/database/document.py +++ b/packages/shared-python/shared/models/database/document.py @@ -49,9 +49,7 @@ class Document(Base): document_metadata: Mapped[Optional[Dict[str, Any]]] = mapped_column( JSON, nullable=True ) - parse_track: Mapped[str] = mapped_column( - String(32), nullable=False, default="chunk" - ) + parse_track: Mapped[str] = mapped_column(String(32), nullable=False, default="chunk") created_at: Mapped[datetime] = mapped_column( DateTime, default=utc_now_naive, nullable=False ) @@ -141,7 +139,6 @@ class DocumentChunk(Base): id: Mapped[str] = mapped_column( String(36), primary_key=True, default=lambda: f"dchk_{uuid4().hex[:12]}" ) - chunk_id: Mapped[str] = mapped_column(String(64), nullable=False) user_id: Mapped[str] = mapped_column(Text, nullable=False) namespace: Mapped[str] = mapped_column( @@ -236,110 +233,6 @@ class DocumentChunk(Base): ) -class DocumentMapUnit(Base): - """Persisted lexical map unit for one document revision. - - These rows are a derived index of the exact leaf and interstitial units - used by map-nav. Full chunk payloads remain in ``document_chunks`` and are - loaded separately for evidence hydration. - """ - - __tablename__ = "document_map_units" - - id: Mapped[str] = mapped_column(String(160), primary_key=True) - document_id: Mapped[str] = mapped_column( - String(36), - ForeignKey("documents.document_id", ondelete="CASCADE"), - nullable=False, - ) - job_result_id: Mapped[str] = mapped_column( - String(36), ForeignKey("job_results.id", ondelete="CASCADE"), nullable=False - ) - unit_id: Mapped[str] = mapped_column(String(128), nullable=False) - section_id: Mapped[str] = mapped_column(String(36), nullable=False) - unit_kind: Mapped[str] = mapped_column(String(32), nullable=False) - path_token_count: Mapped[int] = mapped_column(Integer, nullable=False) - content_token_count: Mapped[int] = mapped_column(Integer, nullable=False) - term_search_text_lower: Mapped[str] = mapped_column(Text, nullable=False) - sort_order: Mapped[int] = mapped_column(Integer, nullable=False, default=0) - created_at: Mapped[datetime] = mapped_column( - DateTime, default=utc_now_naive, nullable=False - ) - - __table_args__ = ( - Index( - "idx_document_map_units_revision_order", - "document_id", - "job_result_id", - "sort_order", - "unit_id", - ), - Index("idx_document_map_units_section", "section_id"), - ) - - -class DocumentMapUnitToken(Base): - """One exact token frequency in a persisted map unit channel.""" - - __tablename__ = "document_map_unit_tokens" - - id: Mapped[str] = mapped_column(String(36), primary_key=True) - map_unit_id: Mapped[str] = mapped_column( - String(160), - ForeignKey("document_map_units.id", ondelete="CASCADE"), - nullable=False, - ) - channel: Mapped[str] = mapped_column(String(16), nullable=False) - token: Mapped[str] = mapped_column(Text, nullable=False) - token_hash: Mapped[str] = mapped_column(String(64), nullable=False) - frequency: Mapped[int] = mapped_column(Integer, nullable=False) - - __table_args__ = ( - Index( - "idx_document_map_unit_tokens_lookup", - "channel", - "token_hash", - "map_unit_id", - ), - Index("idx_document_map_unit_tokens_unit", "map_unit_id", "channel"), - ) - - -class DocumentMapUnitIndex(Base): - """Completeness marker for a revision's materialized map-unit index.""" - - __tablename__ = "document_map_unit_indexes" - - id: Mapped[str] = mapped_column(String(100), primary_key=True) - document_id: Mapped[str] = mapped_column( - String(36), - ForeignKey("documents.document_id", ondelete="CASCADE"), - nullable=False, - ) - job_result_id: Mapped[str] = mapped_column( - String(36), ForeignKey("job_results.id", ondelete="CASCADE"), nullable=False - ) - format_version: Mapped[int] = mapped_column(Integer, nullable=False) - unit_count: Mapped[int] = mapped_column(Integer, nullable=False) - token_count: Mapped[int] = mapped_column(Integer, nullable=False) - created_at: Mapped[datetime] = mapped_column( - DateTime, default=utc_now_naive, nullable=False - ) - - __table_args__ = ( - UniqueConstraint( - "document_id", - "job_result_id", - name="uq_document_map_unit_indexes_revision", - ), - Index( - "idx_document_map_unit_indexes_revision", - "document_id", - "job_result_id", - ), - ) - - class GraphNode(Base): """Persisted derived graph node used for routing and expansion.""" @@ -492,58 +385,44 @@ class RetrievalHitStat(Base): class RetrievalRun(Base): """One row per agentic retrieval query. Append-only analytics.""" - __tablename__ = "retrieval_runs" + __tablename__ = 'retrieval_runs' - run_id: Mapped[str] = mapped_column( - String(36), primary_key=True, default=lambda: f"aret_{uuid4().hex[:12]}" - ) + run_id: Mapped[str] = mapped_column(String(36), primary_key=True, default=lambda: f'aret_{uuid4().hex[:12]}') user_id: Mapped[str] = mapped_column(Text, nullable=False) - namespace: Mapped[str] = mapped_column( - String(255), nullable=False, default="default" - ) + namespace: Mapped[str] = mapped_column(String(255), nullable=False, default='default') query: Mapped[str] = mapped_column(Text, nullable=False) - query_hash: Mapped[str] = mapped_column(String(32), nullable=False, default="") + query_hash: Mapped[str] = mapped_column(String(32), nullable=False, default='') top_k: Mapped[int] = mapped_column(Integer, nullable=False, default=10) chunk_types: Mapped[Optional[List[str]]] = mapped_column(JSON, nullable=True) filters: Mapped[Optional[Dict[str, Any]]] = mapped_column(JSON, nullable=True) - policy_name: Mapped[str] = mapped_column( - String(64), nullable=False, default="rule_based_v1" - ) + policy_name: Mapped[str] = mapped_column(String(64), nullable=False, default='rule_based_v1') agentic_enabled: Mapped[bool] = mapped_column(nullable=False, default=True) cache_hit: Mapped[bool] = mapped_column(nullable=False, default=False) result_count: Mapped[int] = mapped_column(Integer, nullable=False, default=0) final_doc_ids: Mapped[Optional[List[str]]] = mapped_column(JSON, nullable=True) - result_provenance: Mapped[Optional[Dict[str, Any]]] = mapped_column( - JSON, nullable=True - ) - parent_run_id: Mapped[Optional[str]] = mapped_column( - String(36), nullable=True, index=True - ) + result_provenance: Mapped[Optional[Dict[str, Any]]] = mapped_column(JSON, nullable=True) + parent_run_id: Mapped[Optional[str]] = mapped_column(String(36), nullable=True, index=True) workflow_step_id: Mapped[Optional[str]] = mapped_column(String(64), nullable=True) workflow_plan: Mapped[Optional[Dict[str, Any]]] = mapped_column(JSON, nullable=True) latency_ms: Mapped[int] = mapped_column(Integer, nullable=False, default=0) token_count: Mapped[Optional[int]] = mapped_column(Integer, nullable=True) error: Mapped[Optional[str]] = mapped_column(Text, nullable=True) - created_at: Mapped[datetime] = mapped_column( - DateTime, default=datetime.utcnow, nullable=False - ) + created_at: Mapped[datetime] = mapped_column(DateTime, default=datetime.utcnow, nullable=False) completed_at: Mapped[Optional[datetime]] = mapped_column(DateTime, nullable=True) __table_args__ = ( - Index("idx_retrieval_runs_user_namespace", "user_id", "namespace"), - Index("idx_retrieval_runs_created", "created_at"), - Index("idx_retrieval_runs_query_hash", "query_hash"), + Index('idx_retrieval_runs_user_namespace', 'user_id', 'namespace'), + Index('idx_retrieval_runs_created', 'created_at'), + Index('idx_retrieval_runs_query_hash', 'query_hash'), ) class RetrievalStep(Base): """One row per agent step within a retrieval run. Append-only analytics.""" - __tablename__ = "retrieval_steps" + __tablename__ = 'retrieval_steps' - step_id: Mapped[str] = mapped_column( - String(36), primary_key=True, default=lambda: f"arst_{uuid4().hex[:12]}" - ) + step_id: Mapped[str] = mapped_column(String(36), primary_key=True, default=lambda: f'arst_{uuid4().hex[:12]}') run_id: Mapped[str] = mapped_column(String(36), nullable=False, index=True) step_index: Mapped[int] = mapped_column(Integer, nullable=False, default=0) action_type: Mapped[str] = mapped_column(String(64), nullable=False) @@ -555,11 +434,9 @@ class RetrievalStep(Base): token_count: Mapped[Optional[int]] = mapped_column(Integer, nullable=True) model_name: Mapped[Optional[str]] = mapped_column(String(128), nullable=True) error: Mapped[Optional[str]] = mapped_column(Text, nullable=True) - created_at: Mapped[datetime] = mapped_column( - DateTime, default=datetime.utcnow, nullable=False - ) + created_at: Mapped[datetime] = mapped_column(DateTime, default=datetime.utcnow, nullable=False) __table_args__ = ( - Index("idx_retrieval_steps_run", "run_id", "step_index"), - Index("idx_retrieval_steps_created", "created_at"), + Index('idx_retrieval_steps_run', 'run_id', 'step_index'), + Index('idx_retrieval_steps_created', 'created_at'), ) diff --git a/packages/shared-python/shared/services/retrieval/map_unit_index.py b/packages/shared-python/shared/services/retrieval/map_unit_index.py deleted file mode 100644 index 15454dc9..00000000 --- a/packages/shared-python/shared/services/retrieval/map_unit_index.py +++ /dev/null @@ -1,163 +0,0 @@ -"""Publication-time materialization of exact map-nav lexical units.""" - -from __future__ import annotations - -from collections import Counter -from hashlib import sha256 -from uuid import uuid4 - -from sqlalchemy import delete, select -from sqlalchemy.orm import Session - -from shared.models.database.document import ( - DocumentChunk, - DocumentMapUnit, - DocumentMapUnitIndex, - DocumentMapUnitToken, - DocumentSection, -) -from shared.services.retrieval.nav.nav_hierarchy import ProviderToolSpace -from shared.services.retrieval.nav.nav_knowhere import ( - KnowhereProvider, - SectionRow, - UnitRow, -) -from shared.services.retrieval.nav.nav_map_scores import build_score_units -from shared.services.retrieval.publication_models import DocumentPublicationScope - - -MAP_UNIT_INDEX_FORMAT_VERSION = 1 - - -def replace_document_map_units( - db: Session, - *, - scope: DocumentPublicationScope, -) -> None: - """Build the derived index through the authoritative map-unit constructor.""" - db.execute( - delete(DocumentMapUnitToken).where( - DocumentMapUnitToken.map_unit_id.in_( - select(DocumentMapUnit.id) - .where(DocumentMapUnit.document_id == scope.document_id) - .where(DocumentMapUnit.job_result_id == scope.job_result_id) - ) - ) - ) - db.execute( - delete(DocumentMapUnit) - .where(DocumentMapUnit.document_id == scope.document_id) - .where(DocumentMapUnit.job_result_id == scope.job_result_id) - ) - db.execute( - delete(DocumentMapUnitIndex) - .where(DocumentMapUnitIndex.document_id == scope.document_id) - .where(DocumentMapUnitIndex.job_result_id == scope.job_result_id) - ) - section_models = list( - db.scalars( - select(DocumentSection) - .where(DocumentSection.document_id == scope.document_id) - .where(DocumentSection.job_result_id == scope.job_result_id) - .order_by(DocumentSection.sort_order, DocumentSection.section_id) - ) - ) - chunk_models = list( - db.scalars( - select(DocumentChunk) - .where(DocumentChunk.document_id == scope.document_id) - .where(DocumentChunk.job_result_id == scope.job_result_id) - .order_by( - DocumentChunk.sort_order, - DocumentChunk.chunk_id, - DocumentChunk.id, - ) - ) - ) - provider = KnowhereProvider( - doc_id=scope.document_id, - sections=[_to_section_row(section) for section in section_models], - units=[_to_unit_row(chunk) for chunk in chunk_models], - ) - score_units = build_score_units( - ProviderToolSpace(provider), - scope.document_id, - ) - persisted_count = 0 - token_count = 0 - for sort_order, unit in enumerate(score_units): - unit_id = str(unit.get("chunk_id") or "").strip() - section_id = str(unit.get("section_id") or "").strip() - if not unit_id or not section_id: - continue - map_unit_id = f"dmu_{uuid4().hex}" - path_tokens = str(unit.get("path_search_text") or "").split() - content_tokens = str(unit.get("content_search_text") or "").split() - db.add( - DocumentMapUnit( - id=map_unit_id, - document_id=scope.document_id, - job_result_id=scope.job_result_id, - unit_id=unit_id, - section_id=section_id, - unit_kind=str(unit.get("kind") or "leaf"), - path_token_count=len(path_tokens), - content_token_count=len(content_tokens), - term_search_text_lower=str(unit.get("term_search_text") or "").lower(), - sort_order=sort_order, - ) - ) - for channel, frequencies in ( - ("path", Counter(path_tokens)), - ("content", Counter(content_tokens)), - ): - for token, frequency in frequencies.items(): - db.add( - DocumentMapUnitToken( - id=f"dmut_{uuid4().hex[:31]}", - map_unit_id=map_unit_id, - channel=channel, - token=token, - token_hash=sha256(token.encode("utf-8")).hexdigest(), - frequency=frequency, - ) - ) - token_count += len(frequencies) - persisted_count += 1 - db.add( - DocumentMapUnitIndex( - id=f"dmui_{uuid4().hex}", - document_id=scope.document_id, - job_result_id=scope.job_result_id, - format_version=MAP_UNIT_INDEX_FORMAT_VERSION, - unit_count=persisted_count, - token_count=token_count, - ) - ) - - -def _to_section_row(section: DocumentSection) -> SectionRow: - return SectionRow( - section_id=section.section_id, - parent_section_id=section.parent_section_id, - section_path=section.section_path, - section_title=str(section.section_title or ""), - section_level=section.section_level, - summary=str(section.summary or ""), - sort_order=section.sort_order, - ) - - -def _to_unit_row(chunk: DocumentChunk) -> UnitRow: - raw_metadata = chunk.chunk_metadata - metadata = dict(raw_metadata) if isinstance(raw_metadata, dict) else {} - return UnitRow( - chunk_id=chunk.chunk_id, - section_id=chunk.section_id, - chunk_type=chunk.chunk_type, - content=str(chunk.content or ""), - sort_order=chunk.sort_order, - source_chunk_path=str(chunk.source_chunk_path or ""), - file_path=str(chunk.file_path or ""), - metadata=metadata, - ) diff --git a/packages/shared-python/shared/services/retrieval/nav/knowhere_hybrid.py b/packages/shared-python/shared/services/retrieval/nav/knowhere_hybrid.py index 369da962..7d8b195b 100644 --- a/packages/shared-python/shared/services/retrieval/nav/knowhere_hybrid.py +++ b/packages/shared-python/shared/services/retrieval/nav/knowhere_hybrid.py @@ -5,7 +5,6 @@ Reference: https://github.com/Ontos-AI/knowhere """ - from __future__ import annotations import os @@ -57,9 +56,7 @@ def _space_join_tokens(text: str) -> str: return " ".join(tokenize_for_retrieval(text, dedupe=False)) -def build_content_search_text( - content: str, *, section_summary: Optional[str] = None -) -> str: +def build_content_search_text(content: str, *, section_summary: Optional[str] = None) -> str: parts = [str(content or "").strip()] if section_summary and str(section_summary).strip(): parts.append(str(section_summary).strip()) @@ -119,9 +116,7 @@ def rank_rows_by_bm25( try: from rank_bm25 import BM25Okapi except ImportError: - return _rank_rows_by_token_overlap( - rows, query_tokens, search_field=search_field - ) + return _rank_rows_by_token_overlap(rows, query_tokens, search_field=search_field) corpus: List[List[str]] = [] ranked_rows: List[dict[str, Any]] = [] @@ -145,9 +140,7 @@ def rank_rows_by_bm25( return ranked_rows -def rank_rows_by_term_channel( - rows: List[dict[str, Any]], query: str -) -> List[dict[str, Any]]: +def rank_rows_by_term_channel(rows: List[dict[str, Any]], query: str) -> List[dict[str, Any]]: query_lower = query.lower().strip() query_tokens = tokenize_query_for_ranker(query) if not query_lower or not query_tokens: @@ -224,24 +217,12 @@ def normalize_row_scores( def _channel_weights() -> Tuple[float, float, float]: - path_w = float( - os.environ.get( - "NAV_DISCOVERY_CHANNEL_WEIGHT_PATH", str(CHANNEL_WEIGHT_PATH) - ).strip() - or CHANNEL_WEIGHT_PATH - ) + path_w = float(os.environ.get("NAV_DISCOVERY_CHANNEL_WEIGHT_PATH", str(CHANNEL_WEIGHT_PATH)).strip() or CHANNEL_WEIGHT_PATH) content_w = float( - os.environ.get( - "NAV_DISCOVERY_CHANNEL_WEIGHT_CONTENT", str(CHANNEL_WEIGHT_CONTENT) - ).strip() + os.environ.get("NAV_DISCOVERY_CHANNEL_WEIGHT_CONTENT", str(CHANNEL_WEIGHT_CONTENT)).strip() or CHANNEL_WEIGHT_CONTENT ) - term_w = float( - os.environ.get( - "NAV_DISCOVERY_CHANNEL_WEIGHT_TERM", str(CHANNEL_WEIGHT_TERM) - ).strip() - or CHANNEL_WEIGHT_TERM - ) + term_w = float(os.environ.get("NAV_DISCOVERY_CHANNEL_WEIGHT_TERM", str(CHANNEL_WEIGHT_TERM)).strip() or CHANNEL_WEIGHT_TERM) return path_w, content_w, term_w @@ -261,23 +242,14 @@ def hybrid_search_rows( recall_k = internal_recall_k if recall_k is None: - mult = int( - os.environ.get( - "NAV_DISCOVERY_RECALL_MULT", str(INTERNAL_RECALL_K_MULTIPLIER) - ).strip() - or INTERNAL_RECALL_K_MULTIPLIER - ) + mult = int(os.environ.get("NAV_DISCOVERY_RECALL_MULT", str(INTERNAL_RECALL_K_MULTIPLIER)).strip() or INTERNAL_RECALL_K_MULTIPLIER) recall_k = max(top_k, top_k * max(1, mult)) rrf_k = int(os.environ.get("NAV_DISCOVERY_RRF_K", str(RRF_K)).strip() or RRF_K) path_w, content_w, term_w = _channel_weights() - path_rows = rank_rows_by_bm25( - list(rows), query_tokens, search_field="path_search_text" - )[:recall_k] - content_rows = rank_rows_by_bm25( - list(rows), query_tokens, search_field="content_search_text" - )[:recall_k] + path_rows = rank_rows_by_bm25(list(rows), query_tokens, search_field="path_search_text")[:recall_k] + content_rows = rank_rows_by_bm25(list(rows), query_tokens, search_field="content_search_text")[:recall_k] term_rows = rank_rows_by_term_channel(list(rows), query)[:recall_k] fused = merge_channels_rrf( @@ -290,32 +262,27 @@ def hybrid_search_rows( return fused + def map_channel_weights() -> Tuple[float, float, float]: """Channel weights for map scoring (prefer NAV_MAP_* env, fall back to legacy names).""" path_w = float( os.environ.get( "NAV_MAP_CHANNEL_WEIGHT_PATH", - os.environ.get( - "NAV_DISCOVERY_CHANNEL_WEIGHT_PATH", str(CHANNEL_WEIGHT_PATH) - ), + os.environ.get("NAV_DISCOVERY_CHANNEL_WEIGHT_PATH", str(CHANNEL_WEIGHT_PATH)), ).strip() or CHANNEL_WEIGHT_PATH ) content_w = float( os.environ.get( "NAV_MAP_CHANNEL_WEIGHT_CONTENT", - os.environ.get( - "NAV_DISCOVERY_CHANNEL_WEIGHT_CONTENT", str(CHANNEL_WEIGHT_CONTENT) - ), + os.environ.get("NAV_DISCOVERY_CHANNEL_WEIGHT_CONTENT", str(CHANNEL_WEIGHT_CONTENT)), ).strip() or CHANNEL_WEIGHT_CONTENT ) term_w = float( os.environ.get( "NAV_MAP_CHANNEL_WEIGHT_TERM", - os.environ.get( - "NAV_DISCOVERY_CHANNEL_WEIGHT_TERM", str(CHANNEL_WEIGHT_TERM) - ), + os.environ.get("NAV_DISCOVERY_CHANNEL_WEIGHT_TERM", str(CHANNEL_WEIGHT_TERM)), ).strip() or CHANNEL_WEIGHT_TERM ) @@ -478,23 +445,18 @@ def fuse_channel_bm25_dense( dense_vals = [float(dense_by_id.get(uid, 0.0) or 0.0) for uid in unit_ids] bm25_n = _normalize_score_list(bm25_vals) dense_n = _normalize_score_list(dense_vals) - dense_w = float( - os.environ.get("NAV_MAP_CHANNEL_DENSE_WEIGHT", "0.5").strip() or "0.5" - ) + dense_w = float(os.environ.get("NAV_MAP_CHANNEL_DENSE_WEIGHT", "0.5").strip() or "0.5") dense_w = min(1.0, max(0.0, dense_w)) bm25_w = 1.0 - dense_w return { - uid: bm25_w * bm25_n[i] + dense_w * dense_n[i] for i, uid in enumerate(unit_ids) + uid: bm25_w * bm25_n[i] + dense_w * dense_n[i] + for i, uid in enumerate(unit_ids) } def _rank_ids_by_score(score_by_id: Dict[str, float]) -> List[str]: ranked = sorted( - ( - (sid, float(score)) - for sid, score in score_by_id.items() - if float(score) > 0.0 - ), + ((sid, float(score)) for sid, score in score_by_id.items() if float(score) > 0.0), key=lambda item: (-item[1], item[0]), ) return [sid for sid, _ in ranked] @@ -508,7 +470,9 @@ def score_rows_hybrid_all( content_texts: Optional[Dict[str, str]] = None, doc_id: Optional[str] = None, namespace: Optional[str] = None, - dense_scores_by_channel: Optional[Dict[str, Optional[Dict[str, float]]]] = None, + dense_scores_by_channel: Optional[ + Dict[str, Optional[Dict[str, float]]] + ] = None, ) -> List[dict[str, Any]]: """Score every row with path/content/term; optional within-channel dense fuse. @@ -526,9 +490,7 @@ def score_rows_hybrid_all( if not unit_ids: return [dict(row, score=0.0) for row in rows] - row_by_id = { - str(row.get("chunk_id") or ""): dict(row) for row in rows if row.get("chunk_id") - } + row_by_id = {str(row.get("chunk_id") or ""): dict(row) for row in rows if row.get("chunk_id")} path_w, content_w, term_w = map_channel_weights() rrf_k = int( os.environ.get( @@ -539,9 +501,7 @@ def score_rows_hybrid_all( ) if query_tokens: - path_ranked = rank_rows_by_bm25( - list(rows), query_tokens, search_field="path_search_text" - ) + path_ranked = rank_rows_by_bm25(list(rows), query_tokens, search_field="path_search_text") content_ranked = rank_rows_by_bm25( list(rows), query_tokens, search_field="content_search_text" ) @@ -553,8 +513,7 @@ def score_rows_hybrid_all( str(r.get("chunk_id") or ""): float(r.get("score") or 0.0) for r in path_ranked } content_bm25 = { - str(r.get("chunk_id") or ""): float(r.get("score") or 0.0) - for r in content_ranked + str(r.get("chunk_id") or ""): float(r.get("score") or 0.0) for r in content_ranked } term_bm25 = { str(r.get("chunk_id") or ""): float(r.get("score") or 0.0) for r in term_ranked @@ -594,12 +553,19 @@ def score_rows_hybrid_all( namespace=namespace, ) path_dense_by_id = ( - {uid: float(path_dense_scores[i]) for i, uid in enumerate(unit_ids)} - if path_dense_scores is not None and len(path_dense_scores) == len(unit_ids) + { + uid: float(path_dense_scores[i]) + for i, uid in enumerate(unit_ids) + } + if path_dense_scores is not None + and len(path_dense_scores) == len(unit_ids) else None ) content_dense_by_id = ( - {uid: float(content_dense_scores[i]) for i, uid in enumerate(unit_ids)} + { + uid: float(content_dense_scores[i]) + for i, uid in enumerate(unit_ids) + } if content_dense_scores is not None and len(content_dense_scores) == len(unit_ids) else None @@ -609,9 +575,7 @@ def score_rows_hybrid_all( content_dense_by_id = dense_scores_by_channel.get("content") path_channel = fuse_channel_bm25_dense(path_bm25, path_dense_by_id, unit_ids) - content_channel = fuse_channel_bm25_dense( - content_bm25, content_dense_by_id, unit_ids - ) + content_channel = fuse_channel_bm25_dense(content_bm25, content_dense_by_id, unit_ids) term_channel = {uid: float(term_bm25.get(uid, 0.0) or 0.0) for uid in unit_ids} # Convert channel scores to ranked lists for existing RRF merger. @@ -633,9 +597,7 @@ def _rows_from_scores(score_by_id: Dict[str, float]) -> List[dict[str, Any]]: top_k=len(unit_ids), k=rrf_k, ) - fused_by_id = { - str(r.get("chunk_id") or ""): float(r.get("score") or 0.0) for r in fused - } + fused_by_id = {str(r.get("chunk_id") or ""): float(r.get("score") or 0.0) for r in fused} out_rows: List[dict[str, Any]] = [] for uid in unit_ids: row = dict(row_by_id[uid]) @@ -657,58 +619,33 @@ def score_unit_stream_hybrid_all( weighted-RRF implementation, but keeps only token statistics, identifiers, and final scores between bounded provider reads. """ - return score_unit_stream_hybrid_many(unit_factory, [query]).get(query, {}) - - -def score_unit_stream_hybrid_many( - unit_factory: Callable[[], Iterable[ScoreUnitRow]], - queries: Sequence[str], -) -> Dict[str, Dict[str, float]]: - """Score several queries exactly while reading the corpus only once.""" - unique_queries = list(dict.fromkeys(str(query) for query in queries)) - if not unique_queries: - return {} - - query_tokens_by_query = { - query: tokenize_query_for_ranker(query) for query in unique_queries - } - query_token_set = { - token - for query_tokens in query_tokens_by_query.values() - for token in query_tokens - } - query_lower_by_query = {query: query.lower().strip() for query in unique_queries} - units: List[_StreamingManyUnit] = [] + query_tokens = tokenize_query_for_ranker(query) path_stats = _StreamingBm25Stats.empty() content_stats = _StreamingBm25Stats.empty() + units: List[_StreamingUnit] = [] + query_token_set = set(query_tokens) + query_lower = query.lower().strip() for row in unit_factory(): unit_id = str(row.get("chunk_id") or "").strip() if not unit_id: continue path_tokens = _get_search_tokens(row, search_field="path_search_text") content_tokens = _get_search_tokens(row, search_field="content_search_text") - # BM25 corpus statistics are computed over the complete input rows by - # the eager scorer, including rows whose public IDs collide. path_stats.observe(path_tokens) content_stats.observe(content_tokens) path_frequencies = Counter(path_tokens) content_frequencies = Counter(content_tokens) - haystack = str(row.get("term_search_text") or "").lower() - term_scores: List[float] = [] - for query in unique_queries: - query_lower = query_lower_by_query[query] - query_tokens = query_tokens_by_query[query] - term_score = 0.0 - if query_lower: - if query_lower in haystack: - term_score = 100.0 - else: - hit_count = sum(1 for token in query_tokens if token in haystack) - if hit_count > 0: - term_score = float(hit_count) - term_scores.append(term_score) + term_score = 0.0 + if query_lower: + haystack = str(row.get("term_search_text") or "").lower() + if query_lower in haystack: + term_score = 100.0 + else: + hit_count = sum(1 for token in query_tokens if token in haystack) + if hit_count > 0: + term_score = float(hit_count) units.append( - _StreamingManyUnit( + _StreamingUnit( unit_id=unit_id, path_length=len(path_tokens), content_length=len(content_tokens), @@ -722,31 +659,11 @@ def score_unit_stream_hybrid_many( for token in query_token_set if content_frequencies[token] }, - term_scores=tuple(term_scores), + term_score=term_score, ) ) path_stats.finalize() content_stats.finalize() - return { - query: _score_streaming_units( - units, - path_stats=path_stats, - content_stats=content_stats, - query_tokens=query_tokens_by_query[query], - query_index=index, - ) - for index, query in enumerate(unique_queries) - } - - -def _score_streaming_units( - units: Sequence["_StreamingManyUnit"], - *, - path_stats: "_StreamingBm25Stats", - content_stats: "_StreamingBm25Stats", - query_tokens: List[str], - query_index: int, -) -> Dict[str, float]: path_by_id: Dict[str, float] = {} content_by_id: Dict[str, float] = {} term_by_id: Dict[str, float] = {} @@ -760,21 +677,24 @@ def _score_streaming_units( ) path_by_id[unit.unit_id] = path_score content_by_id[unit.unit_id] = content_score - term_by_id[unit.unit_id] = float( - unit.term_scores[query_index] - if query_index < len(unit.term_scores) - else 0.0 - ) + term_by_id[unit.unit_id] = unit.term_score path_rows = [ - (score, unit_id) for unit_id, score in path_by_id.items() if score > 0.0 + (score, unit_id) + for unit_id, score in path_by_id.items() + if score > 0.0 ] content_rows = [ - (score, unit_id) for unit_id, score in content_by_id.items() if score > 0.0 + (score, unit_id) + for unit_id, score in content_by_id.items() + if score > 0.0 ] term_rows = [ - (score, unit_id) for unit_id, score in term_by_id.items() if score > 0.0 + (score, unit_id) + for unit_id, score in term_by_id.items() + if score > 0.0 ] + path_rows.sort(key=lambda item: (-item[0], item[1])) content_rows.sort(key=lambda item: (-item[0], item[1])) term_rows.sort(key=lambda item: (-item[0], item[1])) @@ -797,94 +717,13 @@ def _score_streaming_units( @dataclass(frozen=True) -class _StreamingManyUnit: +class _StreamingUnit: unit_id: str path_length: int content_length: int path_frequencies: Mapping[str, int] content_frequencies: Mapping[str, int] - term_scores: Tuple[float, ...] - - -@dataclass(frozen=True) -class PersistedBm25Stats: - """Corpus statistics needed to reproduce ``BM25Okapi`` exactly.""" - - document_count: int - total_length: int - document_frequency: Mapping[str, int] - average_idf: float - - -@dataclass(frozen=True) -class PersistedScoreUnit: - """Query-specific frequencies for one persisted map unit.""" - - unit_id: str - path_length: int - content_length: int - path_frequencies: Mapping[str, int] - content_frequencies: Mapping[str, int] - term_scores: Tuple[float, ...] - - -@dataclass(frozen=True) -class PersistedScoreCorpus: - """Compact query projection loaded from the map-unit index.""" - - units: Sequence[PersistedScoreUnit] - path_stats: PersistedBm25Stats - content_stats: PersistedBm25Stats - - -def score_persisted_corpus_many( - corpus: PersistedScoreCorpus, - queries: Sequence[str], -) -> Dict[str, Dict[str, float]]: - """Apply the existing BM25/RRF scorer to persisted query projections.""" - unique_queries = list(dict.fromkeys(str(query) for query in queries)) - if not unique_queries: - return {} - path_stats = _restore_bm25_stats(corpus.path_stats) - content_stats = _restore_bm25_stats(corpus.content_stats) - units = [ - _StreamingManyUnit( - unit_id=unit.unit_id, - path_length=unit.path_length, - content_length=unit.content_length, - path_frequencies=unit.path_frequencies, - content_frequencies=unit.content_frequencies, - term_scores=unit.term_scores, - ) - for unit in corpus.units - ] - return { - query: _score_streaming_units( - units, - path_stats=path_stats, - content_stats=content_stats, - query_tokens=tokenize_query_for_ranker(query), - query_index=query_index, - ) - for query_index, query in enumerate(unique_queries) - } - - -def _restore_bm25_stats(source: PersistedBm25Stats) -> "_StreamingBm25Stats": - stats = _StreamingBm25Stats.empty() - stats.document_count = source.document_count - stats.total_length = source.total_length - stats.average_length = ( - source.total_length / source.document_count if source.document_count else 0.0 - ) - idf_by_token: Dict[str, float] = {} - for token, frequency in source.document_frequency.items(): - idf = math.log(source.document_count - frequency + 0.5) - math.log( - frequency + 0.5 - ) - idf_by_token[token] = 0.25 * source.average_idf if idf < 0.0 else idf - stats.idf_by_token = idf_by_token - return stats + term_score: float class _StreamingBm25Stats: diff --git a/packages/shared-python/shared/services/retrieval/nav/nav_hierarchy.py b/packages/shared-python/shared/services/retrieval/nav/nav_hierarchy.py index d9e42b07..0609cbc2 100644 --- a/packages/shared-python/shared/services/retrieval/nav/nav_hierarchy.py +++ b/packages/shared-python/shared/services/retrieval/nav/nav_hierarchy.py @@ -22,22 +22,7 @@ from __future__ import annotations from dataclasses import dataclass, field -from typing import ( - Any, - Dict, - List, - Optional, - Protocol, - Sequence, - Set, - Tuple, - TYPE_CHECKING, - cast, - runtime_checkable, -) - -if TYPE_CHECKING: - from .knowhere_hybrid import PersistedScoreCorpus +from typing import Any, Dict, List, Optional, Protocol, Sequence, Set, Tuple, runtime_checkable @dataclass @@ -185,9 +170,7 @@ def _node_unit_span(self, section_id: str) -> Tuple[str, int, int]: first_order = int(getattr(units[0], "sort_order", 0) or 0) return "\n".join(texts), first_order, len(units) - def _make_chunk( - self, node_id: str, doc_id: str, text: str, order: int, section_id: str - ) -> Any: + def _make_chunk(self, node_id: str, doc_id: str, text: str, order: int, section_id: str) -> Any: from ._compat import Chunk # type: ignore return Chunk( @@ -225,9 +208,7 @@ def _materialize_leaf_path_chunks(self, section_id: str, doc_id: str) -> List[An text = str(self._provider.content(section_id) or "") if not text.strip(): return [] - return [ - self._make_chunk(f"{section_id}__path", doc_id, text, 0, section_id) - ] + return [self._make_chunk(f"{section_id}__path", doc_id, text, 0, section_id)] # One unit per descendant leaf, plus one per interstitial parent, so # node ids line up with the keys nav_map_scores.build_score_units emits. @@ -246,9 +227,7 @@ def _materialize_leaf_path_chunks(self, section_id: str, doc_id: str) -> List[An out.sort(key=lambda c: (min(c.line_ids or (0,)), c.node_id)) return out - def read_chunks( - self, section_id: str, query: str, *, doc_id: str, k: int - ) -> List[Any]: + def read_chunks(self, section_id: str, query: str, *, doc_id: str, k: int) -> List[Any]: del section_id, query, doc_id, k return [] @@ -270,15 +249,6 @@ def prefetch_document_units(self, doc_id: str) -> None: if str(getattr(provider, "doc_id", "")) == str(doc_id): fn() - def prefetch_document_units_batch(self, doc_ids: Sequence[str]) -> None: - """Forward a provider's bounded multi-document prefetch capability.""" - fn = getattr(self._provider, "prefetch_document_units_batch", None) - if callable(fn): - fn(doc_ids) - return - for doc_id in doc_ids: - self.prefetch_document_units(str(doc_id)) - def release_document_units(self, doc_id: str) -> None: """Forward release of one document's prefetched payloads.""" provider = self._provider @@ -291,17 +261,6 @@ def release_document_units(self, doc_id: str) -> None: if str(getattr(provider, "doc_id", "")) == str(doc_id): fn() - def load_persisted_score_corpus( - self, - doc_ids: Sequence[str], - queries: Sequence[str], - ) -> Optional["PersistedScoreCorpus"]: - """Forward the optional revision-pinned map-unit index capability.""" - fn = getattr(self._provider, "load_persisted_score_corpus", None) - if not callable(fn): - return None - return cast(Optional["PersistedScoreCorpus"], fn(doc_ids, queries)) - @dataclass class InMemoryNode: diff --git a/packages/shared-python/shared/services/retrieval/nav/nav_knowhere.py b/packages/shared-python/shared/services/retrieval/nav/nav_knowhere.py index cf1de779..a1387a64 100644 --- a/packages/shared-python/shared/services/retrieval/nav/nav_knowhere.py +++ b/packages/shared-python/shared/services/retrieval/nav/nav_knowhere.py @@ -21,36 +21,16 @@ from __future__ import annotations import os -from hashlib import sha256 from dataclasses import dataclass, field -from typing import ( - Any, - Callable, - Dict, - Iterable, - List, - Mapping, - Optional, - Protocol, - Sequence, - Set, - Tuple, -) +from typing import Any, Callable, Dict, Iterable, List, Optional, Protocol, Sequence, Set, Tuple from .nav_address import NavLevel from .nav_hierarchy import NodeMeta -from .knowhere_hybrid import ( - PersistedBm25Stats, - PersistedScoreCorpus, - PersistedScoreUnit, - tokenize_query_for_ranker, -) _ASSET_TYPES = ("table", "image") # Knowhere sentinel path for the virtual document container (not a collectable leaf). ROOT_SECTION_PATH = "Root" _DEFAULT_DSN = "postgresql://root:root123@127.0.0.1:5433/Knowhere" -_MAP_UNIT_INDEX_FORMAT_VERSION = 1 @dataclass(frozen=True) @@ -161,20 +141,6 @@ def knowhere_database_url() -> str: class ChunkStore(Protocol): - def load_persisted_score_corpus( - self, - document_ids: Sequence[str], - allowed_section_ids_by_document: Mapping[str, Sequence[str]], - queries: Sequence[str], - ) -> Optional[PersistedScoreCorpus]: - raise NotImplementedError - - def load_documents_units( - self, - section_ids_by_document: Mapping[str, Sequence[str]], - ) -> Dict[str, List[UnitRow]]: - raise NotImplementedError - def load_document_units( self, document_id: str, @@ -225,20 +191,11 @@ def load_section_units( doc_id = str(document_id).strip() sid = str(section_id).strip() job_result_id = self._revisions.get(doc_id) - if ( - not doc_id - or not sid - or not job_result_id - or (doc_id, sid) in self._excluded_sections - ): + if not doc_id or not sid or not job_result_id or (doc_id, sid) in self._excluded_sections: return [] cur = self._connection().cursor() try: - ids = [ - str(chunk_id).strip() - for chunk_id in extra_chunk_ids - if str(chunk_id).strip() - ] + ids = [str(chunk_id).strip() for chunk_id in extra_chunk_ids if str(chunk_id).strip()] if ids: cur.execute( "SELECT chunk_id, section_id, chunk_type, content, sort_order, " @@ -262,259 +219,6 @@ def load_section_units( finally: cur.close() - def load_persisted_score_corpus( - self, - document_ids: Sequence[str], - allowed_section_ids_by_document: Mapping[str, Sequence[str]], - queries: Sequence[str], - ) -> Optional[PersistedScoreCorpus]: - """Load query-relevant score inputs when every revision is indexed.""" - revisions = [ - (document_id, self._revisions[document_id]) - for raw_document_id in document_ids - if (document_id := str(raw_document_id).strip()) in self._revisions - ] - if not revisions or len(revisions) != len(document_ids): - return None - values_sql = ", ".join(["(%s, %s)"] * len(revisions)) - revision_params: List[object] = [ - value for revision in revisions for value in revision - ] - cur = self._connection().cursor() - try: - cur.execute( - "SELECT indexes.document_id, indexes.job_result_id, " - "indexes.format_version, indexes.unit_count, indexes.token_count " - "FROM document_map_unit_indexes AS indexes " - f"JOIN (VALUES {values_sql}) AS revisions(document_id, job_result_id) " - "ON indexes.document_id = revisions.document_id " - "AND indexes.job_result_id = revisions.job_result_id", - revision_params, - ) - manifests = list(cur.fetchall()) - if len(manifests) != len(revisions) or any( - int(row[2]) != _MAP_UNIT_INDEX_FORMAT_VERSION for row in manifests - ): - return None - - cur.execute( - "SELECT COUNT(*), COUNT(DISTINCT (units.document_id, units.unit_id)) " - "FROM document_map_units AS units " - f"JOIN (VALUES {values_sql}) AS revisions(document_id, job_result_id) " - "ON units.document_id = revisions.document_id " - "AND units.job_result_id = revisions.job_result_id", - revision_params, - ) - unit_count_row = cur.fetchone() - indexed_unit_count = int(unit_count_row[0]) if unit_count_row else 0 - distinct_unit_count = int(unit_count_row[1]) if unit_count_row else 0 - expected_count = sum(int(row[3]) for row in manifests) - if indexed_unit_count != expected_count or distinct_unit_count != indexed_unit_count: - return None - cur.execute( - "SELECT COUNT(*) FROM document_map_unit_tokens AS tokens " - "JOIN document_map_units AS units ON units.id = tokens.map_unit_id " - f"JOIN (VALUES {values_sql}) AS revisions(document_id, job_result_id) " - "ON units.document_id = revisions.document_id " - "AND units.job_result_id = revisions.job_result_id", - revision_params, - ) - token_row = cur.fetchone() - indexed_token_count = int(token_row[0]) if token_row else 0 - expected_token_count = sum(int(row[4]) for row in manifests) - if indexed_token_count != expected_token_count: - return None - # The public chunk id is content-derived and may repeat within a - # revision. The persisted scorer keys scores by that id, so use - # the legacy payload path whenever ambiguity would change results. - allowed_by_document = { - str(document_id): {str(section_id) for section_id in section_ids} - for document_id, section_ids in allowed_section_ids_by_document.items() - } - allowed_pairs = [ - (document_id, section_id) - for document_id, section_ids in allowed_by_document.items() - for section_id in section_ids - ] - unit_rows: list[Sequence[object]] = [] - if allowed_pairs: - allowed_document_ids = [pair[0] for pair in allowed_pairs] - allowed_section_ids = [pair[1] for pair in allowed_pairs] - cur.execute( - "SELECT units.id, units.document_id, units.unit_id, units.section_id, " - "units.path_token_count, units.content_token_count " - "FROM document_map_units AS units " - f"JOIN (VALUES {values_sql}) AS revisions(document_id, job_result_id) " - "ON units.document_id = revisions.document_id " - "AND units.job_result_id = revisions.job_result_id " - "JOIN UNNEST(%s::text[], %s::text[]) " - "AS allowed(document_id, section_id) " - "ON units.document_id = allowed.document_id " - "AND units.section_id = allowed.section_id " - "ORDER BY units.document_id, units.sort_order, units.unit_id", - [*revision_params, allowed_document_ids, allowed_section_ids], - ) - unit_rows = list(cur.fetchall()) - map_unit_ids = [str(row[0]) for row in unit_rows] - unique_queries = list(dict.fromkeys(str(query) for query in queries)) - query_tokens_by_query = { - query: tokenize_query_for_ranker(query) for query in unique_queries - } - query_tokens = list( - dict.fromkeys( - token - for query in unique_queries - for token in query_tokens_by_query[query] - ) - ) - frequencies: Dict[Tuple[str, str], Dict[str, int]] = {} - if map_unit_ids and query_tokens: - query_token_hashes = [ - sha256(token.encode("utf-8")).hexdigest() for token in query_tokens - ] - cur.execute( - "SELECT map_unit_id, channel, token, frequency " - "FROM document_map_unit_tokens " - "WHERE map_unit_id = ANY(%s) AND token_hash = ANY(%s) " - "AND token = ANY(%s)", - (map_unit_ids, query_token_hashes, query_tokens), - ) - for map_unit_id, channel, token, frequency in cur.fetchall(): - frequencies.setdefault((str(map_unit_id), str(channel)), {})[ - str(token) - ] = int(frequency) - - term_scores = self._load_term_scores( - cur, - map_unit_ids=map_unit_ids, - queries=unique_queries, - query_tokens_by_query=query_tokens_by_query, - ) - path_stats = self._load_persisted_bm25_stats( - cur, - unit_rows=unit_rows, - map_unit_ids=map_unit_ids, - channel="path", - query_tokens=query_tokens, - frequencies=frequencies, - length_index=4, - ) - content_stats = self._load_persisted_bm25_stats( - cur, - unit_rows=unit_rows, - map_unit_ids=map_unit_ids, - channel="content", - query_tokens=query_tokens, - frequencies=frequencies, - length_index=5, - ) - return PersistedScoreCorpus( - units=[ - PersistedScoreUnit( - unit_id=str(row[2]), - path_length=int(row[4]), - content_length=int(row[5]), - path_frequencies=frequencies.get((str(row[0]), "path"), {}), - content_frequencies=frequencies.get( - (str(row[0]), "content"), {} - ), - term_scores=term_scores.get( - str(row[0]), tuple(0.0 for _query in unique_queries) - ), - ) - for row in unit_rows - ], - path_stats=path_stats, - content_stats=content_stats, - ) - finally: - cur.close() - - def _load_term_scores( - self, - cur: "_SyncCursor", - *, - map_unit_ids: Sequence[str], - queries: Sequence[str], - query_tokens_by_query: Mapping[str, Sequence[str]], - ) -> Dict[str, Tuple[float, ...]]: - if not map_unit_ids or not queries: - return {} - expressions: List[str] = [] - params: List[object] = [] - for query in queries: - query_lower = query.lower().strip() - if not query_lower: - expressions.append("0.0") - continue - token_expressions = [ - "CASE WHEN POSITION(%s IN term_search_text_lower) > 0 THEN 1 ELSE 0 END" - for _token in query_tokens_by_query[query] - ] - token_sum = " + ".join(token_expressions) or "0" - expressions.append( - "CASE WHEN POSITION(%s IN term_search_text_lower) > 0 " - f"THEN 100.0 ELSE ({token_sum})::double precision END" - ) - params.append(query_lower) - params.extend(query_tokens_by_query[query]) - params.append(list(map_unit_ids)) - cur.execute( - "SELECT id, " + ", ".join(expressions) + " " - "FROM document_map_units WHERE id = ANY(%s)", - params, - ) - return { - str(row[0]): tuple(float(value) for value in row[1:]) - for row in cur.fetchall() - } - - def _load_persisted_bm25_stats( - self, - cur: "_SyncCursor", - *, - unit_rows: Sequence[Sequence[object]], - map_unit_ids: Sequence[str], - channel: str, - query_tokens: Sequence[str], - frequencies: Mapping[Tuple[str, str], Mapping[str, int]], - length_index: int, - ) -> PersistedBm25Stats: - lengths = [ - int(row[length_index]) for row in unit_rows if int(row[length_index]) > 0 - ] - document_count = len(lengths) - document_frequency = { - token: sum( - 1 - for row in unit_rows - if frequencies.get((str(row[0]), channel), {}).get(token, 0) > 0 - ) - for token in query_tokens - } - needs_average_idf = any( - frequency > document_count / 2 for frequency in document_frequency.values() - ) - average_idf = 0.0 - if needs_average_idf and map_unit_ids and document_count: - cur.execute( - "SELECT COALESCE(AVG(LN((%s - frequencies.document_frequency + 0.5) " - "/ (frequencies.document_frequency + 0.5))), 0.0) " - "FROM (SELECT token, COUNT(*) AS document_frequency " - "FROM document_map_unit_tokens " - "WHERE map_unit_id = ANY(%s) AND channel = %s " - "GROUP BY token) AS frequencies", - (document_count, list(map_unit_ids), channel), - ) - row = cur.fetchone() - average_idf = float(row[0]) if row else 0.0 - return PersistedBm25Stats( - document_count=document_count, - total_length=sum(lengths), - document_frequency=document_frequency, - average_idf=average_idf, - ) - def load_document_units( self, document_id: str, @@ -524,11 +228,7 @@ def load_document_units( """Load one document's section payloads in a single ordered query.""" doc_id = str(document_id).strip() job_result_id = self._revisions.get(doc_id) - section_values = [ - str(section_id).strip() - for section_id in section_ids - if str(section_id).strip() - ] + section_values = [str(section_id).strip() for section_id in section_ids if str(section_id).strip()] extra_ids = [ str(chunk_id).strip() for chunk_ids in (extra_chunk_ids_by_section or {}).values() @@ -568,82 +268,6 @@ def load_document_units( finally: cur.close() - def load_documents_units( - self, - section_ids_by_document: Mapping[str, Sequence[str]], - ) -> Dict[str, List[UnitRow]]: - """Load a bounded group of revisions with keyset-paged SQL queries.""" - requested = [ - (str(document_id).strip(), self._revisions.get(str(document_id).strip())) - for document_id in section_ids_by_document - if str(document_id).strip() - ] - revisions = [ - (document_id, str(job_result_id)) - for document_id, job_result_id in requested - if job_result_id - ] - if not revisions: - return {} - - values_sql = ", ".join(["(%s, %s)"] * len(revisions)) - params: List[object] = [value for revision in revisions for value in revision] - units_by_document: Dict[str, List[UnitRow]] = { - document_id: [] for document_id, _job_result_id in revisions - } - last_key: Optional[Tuple[str, str, int, str, str]] = None - while True: - page_params = list(params) - keyset_sql = "" - if last_key is not None: - keyset_sql = ( - " AND (chunks.document_id, chunks.job_result_id, " - "chunks.sort_order, chunks.chunk_id, chunks.id) > " - "(%s, %s, %s, %s, %s)" - ) - page_params.extend(last_key) - page_params.append(10_000) - cur = self._connection().cursor() - try: - cur.execute( - "SELECT chunks.document_id, chunks.job_result_id, chunks.chunk_id, " - "chunks.section_id, chunks.chunk_type, chunks.content, " - "chunks.sort_order, chunks.source_chunk_path, chunks.file_path, " - "chunks.chunk_metadata, chunks.id " - "FROM document_chunks AS chunks " - f"JOIN (VALUES {values_sql}) AS revisions(document_id, job_result_id) " - "ON chunks.document_id = revisions.document_id " - "AND chunks.job_result_id = revisions.job_result_id" - f"{keyset_sql} " - "ORDER BY chunks.document_id, chunks.job_result_id, " - "chunks.sort_order, chunks.chunk_id, chunks.id LIMIT %s", - page_params, - ) - rows = cur.fetchall() - finally: - cur.close() - if not rows: - break - for row in rows: - document_id = str(row[0]) - section_id = str(row[3]) if row[3] else None - if section_id and (document_id, section_id) in self._excluded_sections: - continue - units_by_document.setdefault(document_id, []).append( - _unit_from_row(row[2:10]) - ) - last = rows[-1] - last_key = ( - str(last[0]), - str(last[1]), - int(last[6] or 0), - str(last[2] or ""), - str(last[10] or ""), - ) - if len(rows) < 10_000: - break - return units_by_document - def close(self) -> None: if self._conn is not None: self._conn.close() @@ -657,9 +281,6 @@ def execute(self, query: str, params: Sequence[object]) -> None: def fetchall(self) -> Sequence[Sequence[object]]: raise NotImplementedError - def fetchone(self) -> Optional[Sequence[object]]: - raise NotImplementedError - def close(self) -> None: raise NotImplementedError @@ -987,14 +608,6 @@ def prefetch_document_units(self) -> None: section_ids, self._remounted_assets_by_section, ) - self.install_prefetched_document_units(loaded) - - def install_prefetched_document_units( - self, - loaded: Sequence[UnitRow], - ) -> None: - """Install rows fetched by either a document or corpus-group query.""" - section_ids = list(self._sections) by_section: Dict[str, List[UnitRow]] = {} for unit in loaded: sid = str(unit.section_id or "").strip() @@ -1182,9 +795,7 @@ def load_namespace_from_db( providers = [load_document_from_db(did, dsn=url) for did in wanted] merged_titles = dict(auto_titles) if titles: - merged_titles.update( - {str(k): str(v) for k, v in titles.items() if str(k).strip()} - ) + merged_titles.update({str(k): str(v) for k, v in titles.items() if str(k).strip()}) return NamespaceKnowhereProvider(providers, titles=merged_titles or None) @@ -1261,77 +872,6 @@ def prefetch_document_units(self, doc_id: str) -> None: if callable(prefetch): prefetch() - def prefetch_document_units_batch(self, doc_ids: Sequence[str]) -> None: - """Load a bounded document group with one query per shared chunk store.""" - providers = [ - self._docs[doc_id] - for raw_doc_id in doc_ids - if (doc_id := str(raw_doc_id).strip()) in self._docs - ] - providers_by_store: Dict[int, List[LazyKnowhereProvider]] = {} - stores_by_id: Dict[int, ChunkStore] = {} - for provider in providers: - if not isinstance(provider, LazyKnowhereProvider): - continue - store = provider._chunk_store - store_id = id(store) - stores_by_id[store_id] = store - providers_by_store.setdefault(store_id, []).append(provider) - - for store_id, lazy_providers in providers_by_store.items(): - store = stores_by_id[store_id] - batch_loader = getattr(store, "load_documents_units", None) - if not callable(batch_loader): - for provider in lazy_providers: - provider.prefetch_document_units() - continue - loaded_by_document = batch_loader( - { - provider.doc_id: list(provider._sections) - for provider in lazy_providers - } - ) - for provider in lazy_providers: - provider.install_prefetched_document_units( - loaded_by_document.get(provider.doc_id, ()) - ) - - def load_persisted_score_corpus( - self, - doc_ids: Sequence[str], - queries: Sequence[str], - ) -> Optional[PersistedScoreCorpus]: - """Return the index projection only when all documents share one store.""" - providers = [ - self._docs[document_id] - for raw_document_id in doc_ids - if (document_id := str(raw_document_id).strip()) in self._docs - ] - if len(providers) != len(doc_ids) or not all( - isinstance(provider, LazyKnowhereProvider) for provider in providers - ): - return None - lazy_providers = [ - provider - for provider in providers - if isinstance(provider, LazyKnowhereProvider) - ] - stores = { - id(provider._chunk_store): provider._chunk_store - for provider in lazy_providers - } - if len(stores) != 1: - return None - store = next(iter(stores.values())) - loader = getattr(store, "load_persisted_score_corpus", None) - if not callable(loader): - return None - return loader( - [provider.doc_id for provider in lazy_providers], - {provider.doc_id: list(provider._sections) for provider in lazy_providers}, - queries, - ) - def release_document_units(self, doc_id: str) -> None: provider = self._docs.get(str(doc_id).strip()) release = getattr(provider, "release_document_units", None) @@ -1379,12 +919,8 @@ def node_meta(self, section_id: str) -> NodeMeta: if sid in self._docs: provider = self._docs[sid] count_fn = getattr(provider, "chunk_count", None) - n_chunks = ( - int(count_fn()) - if callable(count_fn) - else sum( - len(provider.self_units(sec)) for sec in provider.all_section_ids() - ) + n_chunks = int(count_fn()) if callable(count_fn) else sum( + len(provider.self_units(sec)) for sec in provider.all_section_ids() ) return NodeMeta( title=self._titles.get(sid, sid), @@ -1416,9 +952,7 @@ def content(self, section_id: str) -> str: if sid in self._docs: provider = self._docs[sid] return "\n".join( - provider.content(root) - for root in provider.roots(sid) - if provider.content(root) + provider.content(root) for root in provider.roots(sid) if provider.content(root) ) owner = self._section_owner.get(sid) if not owner: diff --git a/packages/shared-python/shared/services/retrieval/nav/nav_map_scores.py b/packages/shared-python/shared/services/retrieval/nav/nav_map_scores.py index 6318f490..0223e232 100644 --- a/packages/shared-python/shared/services/retrieval/nav/nav_map_scores.py +++ b/packages/shared-python/shared/services/retrieval/nav/nav_map_scores.py @@ -9,23 +9,16 @@ build_path_search_text, build_term_search_text, score_rows_hybrid_all, - score_persisted_corpus_many, - score_unit_stream_hybrid_many, + score_unit_stream_hybrid_all, ) -# Keep one bulk read bounded when a namespace contains a very large document, -# while still replacing the per-document N+1 access pattern. -_CORPUS_PREFETCH_GROUP_SIZE = 8 - def _children_ids(ts: Any, section_id: str, doc_id: str) -> List[str]: children_fn = getattr(ts, "_children_for_section_path", None) if not callable(children_fn): st = ts.get_structure(section_id) rows = st.get("children") or [] - return [ - str(r.get("section_id") or "").strip() for r in rows if r.get("section_id") - ] + return [str(r.get("section_id") or "").strip() for r in rows if r.get("section_id")] rows = children_fn(section_id, doc_id, limit=100000) return [str(r.get("section_id") or "").strip() for r in rows if r.get("section_id")] @@ -155,7 +148,8 @@ def _pool_unit_scores_to_tree( ) -> Dict[str, float]: """MAX-pool globally comparable unit scores onto one document tree.""" map_scores = { - leaf_id: float(unit_scores.get(leaf_id, 0.0) or 0.0) for leaf_id in leaves + leaf_id: float(unit_scores.get(leaf_id, 0.0) or 0.0) + for leaf_id in leaves } def score_node(section_id: str) -> float: @@ -166,9 +160,12 @@ def score_node(section_id: str) -> float: score = float(unit_scores.get(section_id, 0.0) or 0.0) map_scores[section_id] = score return score - descendant_leaves = _collect_descendant_leaves(section_id, children_map, leaves) + descendant_leaves = _collect_descendant_leaves( + section_id, children_map, leaves + ) parts = [ - float(unit_scores.get(leaf_id, 0.0) or 0.0) for leaf_id in descendant_leaves + float(unit_scores.get(leaf_id, 0.0) or 0.0) + for leaf_id in descendant_leaves ] self_key = f"{section_id}__self" if self_key in unit_scores: @@ -182,9 +179,7 @@ def score_node(section_id: str) -> float: return map_scores -def build_score_units( - ts: Any, doc_id: str, root_ids: Optional[Sequence[str]] = None -) -> List[dict]: +def build_score_units(ts: Any, doc_id: str, root_ids: Optional[Sequence[str]] = None) -> List[dict]: """Build leaf (+ interstitial self_only) units for hybrid scoring.""" if root_ids is None: root_ids = list(ts.sections_for_doc(doc_id)) @@ -213,9 +208,7 @@ def build_score_units( section_path=path_text, section_title=title or content ), "content_search_text": build_content_search_text(content), - "term_search_text": build_term_search_text( - content, path_text=path_text - ), + "term_search_text": build_term_search_text(content, path_text=path_text), } ) @@ -242,9 +235,7 @@ def build_score_units( section_path=path_text, section_title=titles.get(sid) or "" ), "content_search_text": build_content_search_text(self_text), - "term_search_text": build_term_search_text( - self_text, path_text=path_text - ), + "term_search_text": build_term_search_text(self_text, path_text=path_text), } ) return units @@ -279,9 +270,7 @@ def iter_score_units( section_path=path_text, section_title=title or content ), "content_search_text": build_content_search_text(content), - "term_search_text": build_term_search_text( - content, path_text=path_text - ), + "term_search_text": build_term_search_text(content, path_text=path_text), } finally: if callable(release): @@ -305,9 +294,7 @@ def iter_score_units( section_path=path_text, section_title=titles.get(sid) or "" ), "content_search_text": build_content_search_text(self_text), - "term_search_text": build_term_search_text( - self_text, path_text=path_text - ), + "term_search_text": build_term_search_text(self_text, path_text=path_text), } finally: if callable(release): @@ -358,9 +345,7 @@ def compute_map_and_unit_scores( doc_id=doc_id, namespace=ns, ) - unit_score = { - str(r.get("chunk_id") or ""): float(r.get("score") or 0.0) for r in scored - } + unit_score = {str(r.get("chunk_id") or ""): float(r.get("score") or 0.0) for r in scored} map_scores = _pool_unit_scores_to_tree(children_map, leaves, unit_score) return map_scores, unit_score @@ -377,26 +362,6 @@ def compute_corpus_map_and_unit_scores( All documents share one BM25 corpus, path/content normalization, channel ranking, and RRF pass. Document-level scores are keyed by bare ``document_id``. """ - return compute_corpus_map_and_unit_scores_many( - ts, - doc_ids=doc_ids, - queries=[query], - namespace=namespace, - ).get(query, ({}, {})) - - -def compute_corpus_map_and_unit_scores_many( - ts: Any, - *, - doc_ids: Sequence[str], - queries: Sequence[str], - namespace: Optional[str] = None, -) -> Dict[str, Tuple[Dict[str, float], Dict[str, float]]]: - """Globally score several queries with one replay of the corpus units.""" - unique_queries = list(dict.fromkeys(str(query) for query in queries)) - if not unique_queries: - return {} - valid_doc_ids: List[str] = [] seen_doc_ids: Set[str] = set() for raw in doc_ids: @@ -418,37 +383,10 @@ def compute_corpus_map_and_unit_scores_many( tree_by_doc[doc_id] = (children_map, leaves, titles) def unit_factory() -> Iterator[ScoreUnitRow]: - prefetch_batch = getattr(ts, "prefetch_document_units_batch", None) - release = getattr(ts, "release_document_units", None) - if callable(prefetch_batch): - for group_start in range( - 0, - len(valid_doc_ids), - _CORPUS_PREFETCH_GROUP_SIZE, - ): - document_group = valid_doc_ids[ - group_start : group_start + _CORPUS_PREFETCH_GROUP_SIZE - ] - prefetch_batch(document_group) - try: - for document_id in document_group: - children_map, leaves, titles = tree_by_doc[document_id] - yield from iter_score_units( - ts, - document_id, - children_map=children_map, - leaves=leaves, - titles=titles, - ) - finally: - if callable(release): - for document_id in document_group: - release(document_id) - return - for document_id in valid_doc_ids: children_map, leaves, titles = tree_by_doc[document_id] prefetch = getattr(ts, "prefetch_document_units", None) + release = getattr(ts, "release_document_units", None) if callable(prefetch): prefetch(document_id) try: @@ -463,34 +401,21 @@ def unit_factory() -> Iterator[ScoreUnitRow]: if callable(release): release(document_id) - persisted_loader = getattr(ts, "load_persisted_score_corpus", None) - persisted_corpus = ( - persisted_loader(valid_doc_ids, unique_queries) - if callable(persisted_loader) - else None - ) - unit_scores_by_query = ( - score_persisted_corpus_many(persisted_corpus, unique_queries) - if persisted_corpus is not None - else score_unit_stream_hybrid_many(unit_factory, unique_queries) - ) - results: Dict[str, Tuple[Dict[str, float], Dict[str, float]]] = {} - for query in unique_queries: - unit_scores = unit_scores_by_query.get(query, {}) - map_scores: Dict[str, float] = {} - for doc_id in valid_doc_ids: - children_map, leaves, _titles = tree_by_doc[doc_id] - doc_map_scores = _pool_unit_scores_to_tree( - children_map, leaves, unit_scores - ) - map_scores.update(doc_map_scores) - doc_max = max( - (float(value) for value in doc_map_scores.values()), - default=0.0, - ) - map_scores[doc_id] = doc_max - results[query] = (map_scores, unit_scores) - return results + unit_scores = score_unit_stream_hybrid_all(unit_factory, query) + + map_scores: Dict[str, float] = {} + for doc_id in valid_doc_ids: + children_map, leaves, _titles = tree_by_doc[doc_id] + doc_map_scores = _pool_unit_scores_to_tree( + children_map, leaves, unit_scores + ) + map_scores.update(doc_map_scores) + doc_max = max( + (float(value) for value in doc_map_scores.values()), + default=0.0, + ) + map_scores[doc_id] = doc_max + return map_scores, unit_scores def unit_id_to_section_id(unit_id: str) -> str: @@ -549,44 +474,3 @@ def relight_map_for_query( ts, doc_ids=doc_ids, query=query ) return map_scores, unit_scores, select_map_highlights(unit_scores, k=int(top_k)) - - -def relight_maps_for_queries( - ts: Any, - *, - doc_id: str, - queries: Sequence[str], - top_k: int = 6, -) -> Dict[str, Tuple[Dict[str, float], Dict[str, float], List[str]]]: - """Re-score a shared map for several queries with one corpus replay.""" - unique_queries = list(dict.fromkeys(str(query) for query in queries)) - if not unique_queries: - return {} - doc = str(doc_id or "").strip() - if doc: - return { - query: relight_map_for_query( - ts, - doc_id=doc, - query=query, - top_k=top_k, - ) - for query in unique_queries - } - - doc_ids = [str(value) for value in (ts.document_ids() or ()) if str(value).strip()] - if not doc_ids: - return {} - scored = compute_corpus_map_and_unit_scores_many( - ts, - doc_ids=doc_ids, - queries=unique_queries, - ) - return { - query: ( - map_scores, - unit_scores, - select_map_highlights(unit_scores, k=int(top_k)), - ) - for query, (map_scores, unit_scores) in scored.items() - } diff --git a/packages/shared-python/shared/services/retrieval/nav/nav_orchestrate.py b/packages/shared-python/shared/services/retrieval/nav/nav_orchestrate.py index 91f45e1c..1e30b8f4 100644 --- a/packages/shared-python/shared/services/retrieval/nav/nav_orchestrate.py +++ b/packages/shared-python/shared/services/retrieval/nav/nav_orchestrate.py @@ -126,16 +126,6 @@ def _unbound_retrieval_query(subgoal: Subgoal) -> str: return raw or (subgoal.need or "").strip() or subgoal.retrieval_query -def _resolve_subgoal_query(state: NavState, subgoal: Subgoal) -> str: - query = bind_slots(subgoal.retrieval_query, state.slot_bindings) - if unbound_slots(query): - query = _unbound_retrieval_query(subgoal) - refined = str( - (state.subgoal_refined_queries or {}).get(subgoal.id) or "" - ).strip() - return refined or query - - def _run_navigate_for_query( ts: Any, state: NavState, @@ -190,7 +180,6 @@ def _relit_map( config: NavConfig, *, query: str, - prepared: Optional[Tuple[Dict[str, float], Dict[str, float], List[str]]] = None, ) -> Iterator[None]: """Score the shared map against the harvest ``query`` for one call. @@ -200,9 +189,9 @@ def _relit_map( the query the policy is told to pursue. Scoring failures degrade to the episode lighting. """ - relit = prepared + relit: Optional[Tuple[Dict[str, float], Dict[str, float], List[str]]] = None q = (query or "").strip() - if relit is None and q: + if q: try: from .nav_map_scores import relight_map_for_query @@ -235,10 +224,6 @@ def _execute_subgoal_harvest_once( subgoal: Subgoal, *, steps_out: Optional[List[Any]], - retrieval_query: Optional[str] = None, - prepared_relight: Optional[ - Tuple[Dict[str, float], Dict[str, float], List[str]] - ] = None, ) -> Dict[str, Any]: """One harvest() call for this subgoal this wave — no internal retry loop. @@ -247,21 +232,22 @@ def _execute_subgoal_harvest_once( """ from .nav_harvest import harvest - rq = retrieval_query or _resolve_subgoal_query(state, subgoal) + rq = bind_slots(subgoal.retrieval_query, state.slot_bindings) + if unbound_slots(rq): + # F1: deps may be "settled" (satisfied or dropped) without ever + # producing this subgoal's referenced slot — degrade to a query with + # the unresolved {{...}} braces stripped rather than stalling. + rq = _unbound_retrieval_query(subgoal) refined = str((state.subgoal_refined_queries or {}).get(subgoal.id) or "").strip() + if refined: + rq = refined _set_focus(state, subgoal, rq) # Always enter at namespace/document root; prior dead-ends stay hidden via # subgoal_dismissed_section_ids so the next harvest sees siblings instead. before_sections = set(state.collected_section_ids) before_explicit = set(state.explicit_collect_ids) before_len = len(state.collected) - with _relit_map( - ts, - state, - config, - query=rq, - prepared=prepared_relight, - ): + with _relit_map(ts, state, config, query=rq): harvest_result = harvest( ts, state, @@ -447,39 +433,10 @@ def execute_plan( by_id = {s.id: s for s in plan.subgoals} outputs: List[Dict[str, Any]] = [] - query_by_subgoal = { - sid: _resolve_subgoal_query(state, by_id[sid]) for sid in ready - } - prepared_relights: Dict[ - str, - Tuple[Dict[str, float], Dict[str, float], List[str]], - ] = {} - try: - from .nav_map_scores import relight_maps_for_queries - - prepared_relights = relight_maps_for_queries( - ts, - doc_id=state.doc_id, - queries=list(query_by_subgoal.values()), - top_k=int(config.collect_top_k), - ) - except Exception: - prepared_relights = {} def _run_one(sid: str, working_state: NavState, out_steps: Optional[List[Any]]) -> Dict[str, Any]: - query = query_by_subgoal[sid] - prepared = prepared_relights.get(query) - if prepared is not None and not prepared[0]: - prepared = None return _execute_subgoal_harvest_once( - ts, - working_state, - config, - plan, - by_id[sid], - steps_out=out_steps, - retrieval_query=query, - prepared_relight=prepared, + ts, working_state, config, plan, by_id[sid], steps_out=out_steps ) # Serial wave execution (parallel fan-out retired with ThreadPoolExecutor). diff --git a/packages/shared-python/shared/services/retrieval/publication_content.py b/packages/shared-python/shared/services/retrieval/publication_content.py index 19a7f207..9ebbdd74 100644 --- a/packages/shared-python/shared/services/retrieval/publication_content.py +++ b/packages/shared-python/shared/services/retrieval/publication_content.py @@ -6,13 +6,7 @@ from sqlalchemy import delete from sqlalchemy.orm import Session -from shared.models.database.document import ( - DocumentChunk, - DocumentMapUnit, - DocumentMapUnitIndex, - DocumentSection, -) -from shared.services.retrieval.map_unit_index import replace_document_map_units +from shared.models.database.document import DocumentChunk, DocumentSection from shared.services.retrieval.publication_models import DocumentPublicationScope from shared.services.retrieval.search.lexical_text import ( build_content_lexical_text, @@ -63,9 +57,7 @@ def replace_document_revision_content( """Replace retrieval sections and chunks for one published document revision.""" _delete_existing_revision_content(db, scope=scope) section_publisher = DocumentSectionPublisher( - db=db, - scope=scope, - section_summaries=section_summaries, + db=db, scope=scope, section_summaries=section_summaries, ) for index, chunk in enumerate(chunks): safe_chunk = cast(dict[str, Any], remove_nul_characters(chunk)) @@ -89,8 +81,6 @@ def replace_document_revision_content( fallback_sort_order=index, ) ) - db.flush() - replace_document_map_units(db, scope=scope) class DocumentSectionPublisher: @@ -156,16 +146,6 @@ def _delete_existing_revision_content( *, scope: DocumentPublicationScope, ) -> None: - db.execute( - delete(DocumentMapUnitIndex) - .where(DocumentMapUnitIndex.document_id == scope.document_id) - .where(DocumentMapUnitIndex.job_result_id == scope.job_result_id) - ) - db.execute( - delete(DocumentMapUnit) - .where(DocumentMapUnit.document_id == scope.document_id) - .where(DocumentMapUnit.job_result_id == scope.job_result_id) - ) db.execute( delete(DocumentChunk) .where(DocumentChunk.document_id == scope.document_id) From 37daa76582353a3f9cec63c06cc27739c0c5fb0f Mon Sep 17 00:00:00 2001 From: suguanYang Date: Mon, 31 Aug 2026 11:08:46 +0800 Subject: [PATCH 06/10] Revert "Merge pull request #343 from Ontos-AI/fix/wangbinqi/normalize-native-dsn" This reverts commit a87134fdec545df5bd7d8252d88e8dfaa70aa70d, reversing changes made to 5241307011930d78e8f4e0096c4f6929268b61f3. --- ...etrieval_lazy_snapshot_quality_contract.py | 36 +----- .../services/retrieval/execution/routes.py | 24 ---- .../services/retrieval/nav/nav_hierarchy.py | 24 ---- .../services/retrieval/nav/nav_knowhere.py | 116 ------------------ .../services/retrieval/nav/nav_map_scores.py | 22 ++-- 5 files changed, 10 insertions(+), 212 deletions(-) diff --git a/apps/api/tests/contract/test_retrieval_lazy_snapshot_quality_contract.py b/apps/api/tests/contract/test_retrieval_lazy_snapshot_quality_contract.py index cfb82c9e..13c5bcf4 100644 --- a/apps/api/tests/contract/test_retrieval_lazy_snapshot_quality_contract.py +++ b/apps/api/tests/contract/test_retrieval_lazy_snapshot_quality_contract.py @@ -27,32 +27,6 @@ @dataclass class _FakeChunkStore: units_by_section: dict[str, list[UnitRow]] - document_loads: int = 0 - - def load_document_units( - self, - document_id: str, - section_ids: Sequence[str], - extra_chunk_ids_by_section: dict[str, Sequence[str]] | None = None, - ) -> list[UnitRow]: - del document_id - self.document_loads += 1 - selected = {str(section_id) for section_id in section_ids} - units = [ - unit - for section_id, section_units in self.units_by_section.items() - if section_id in selected - for unit in section_units - ] - known = {unit.chunk_id for unit in units} - for chunk_ids in (extra_chunk_ids_by_section or {}).values(): - for chunk_id in chunk_ids: - for section_units in self.units_by_section.values(): - for unit in section_units: - if unit.chunk_id == chunk_id and unit.chunk_id not in known: - units.append(unit) - known.add(unit.chunk_id) - return units def load_section_units( self, @@ -73,7 +47,7 @@ def close(self) -> None: return None -def _providers() -> tuple[ProviderToolSpace, ProviderToolSpace, _FakeChunkStore]: +def _providers() -> tuple[ProviderToolSpace, ProviderToolSpace]: sections = [ SectionRow("root", None, "Root", "Root", 0, "", 0), SectionRow("section", "root", "Root / Section", "Section", 1, "", 1), @@ -115,11 +89,11 @@ def _providers() -> tuple[ProviderToolSpace, ProviderToolSpace, _FakeChunkStore] titles={"doc": "document"}, chunk_owner_by_id={"duplicate-chunk": "doc", "asset-1": "doc"}, ) - return ProviderToolSpace(eager), ProviderToolSpace(lazy), store + return ProviderToolSpace(eager), ProviderToolSpace(lazy) def test_lazy_provider_preserves_score_units_and_scores() -> None: - eager, lazy, store = _providers() + eager, lazy = _providers() assert build_score_units(eager, "doc") == build_score_units(lazy, "doc") assert compute_corpus_map_and_unit_scores( @@ -129,10 +103,6 @@ def test_lazy_provider_preserves_score_units_and_scores() -> None: ) lazy_provider = lazy._provider - store.document_loads = 0 - prefetch = getattr(lazy_provider, "prefetch_document_units") - prefetch("doc") - assert store.document_loads == 1 self_units = getattr(lazy_provider, "self_units") assert [unit.chunk_id for unit in self_units("leaf")] == [ "duplicate-chunk", diff --git a/packages/shared-python/shared/services/retrieval/execution/routes.py b/packages/shared-python/shared/services/retrieval/execution/routes.py index 97961074..ad2268ab 100644 --- a/packages/shared-python/shared/services/retrieval/execution/routes.py +++ b/packages/shared-python/shared/services/retrieval/execution/routes.py @@ -1,7 +1,6 @@ from __future__ import annotations import asyncio -import time from contextlib import AbstractAsyncContextManager from loguru import logger @@ -181,7 +180,6 @@ async def _run_mapnav_route( episode_token_count, episode_workflow_plan, ) - snapshot_started = time.perf_counter() snapshot = await load_nav_snapshot( context.db, user_id=context.user_id, @@ -190,14 +188,6 @@ async def _run_mapnav_route( exclude_sections=context.exclude_sections, lazy=True, ) - snapshot_seconds = time.perf_counter() - snapshot_started - logger.info( - "retrieval mapnav stage=snapshot_load seconds={:.3f} documents={} refs={}".format( - snapshot_seconds, - len(snapshot.document_ids), - len(snapshot.chunk_ref_index), - ) - ) # Small-corpus count / snapshot reads may leave a checkout; drop it before # the sync LLM episode (same pattern as the retired workflow route). @@ -207,7 +197,6 @@ async def _run_mapnav_route( cfg = build_nav_config() toolspace = ProviderToolSpace(snapshot.provider) - episode_started = time.perf_counter() try: episode = await asyncio.to_thread( run_nav_episode, @@ -222,16 +211,9 @@ async def _run_mapnav_route( ) refs, score_by_chunk_id = build_referenced_chunks(episode, snapshot) - logger.info( - "retrieval mapnav stage=episode seconds={:.3f} refs={}".format( - time.perf_counter() - episode_started, - len(refs), - ) - ) finally: snapshot.close() - hydration_started = time.perf_counter() async with open_fresh_database_context() as final_db: resolved = await resolve_workflow_references( db=final_db, @@ -278,12 +260,6 @@ async def _run_mapnav_route( selected_paths=selected_paths, selected_doc_ids=selected_docs, ) - logger.info( - "retrieval mapnav stage=hydration seconds={:.3f} results={}".format( - time.perf_counter() - hydration_started, - len(assembled_rows), - ) - ) stop_reason = str(getattr(episode, "stop_reason", "") or "completed") evidence_text = str(getattr(episode, "evidence_text", "") or "") diff --git a/packages/shared-python/shared/services/retrieval/nav/nav_hierarchy.py b/packages/shared-python/shared/services/retrieval/nav/nav_hierarchy.py index 0609cbc2..93b35598 100644 --- a/packages/shared-python/shared/services/retrieval/nav/nav_hierarchy.py +++ b/packages/shared-python/shared/services/retrieval/nav/nav_hierarchy.py @@ -237,30 +237,6 @@ def release_section_units(self, section_id: str) -> None: if callable(release): release(section_id) - def prefetch_document_units(self, doc_id: str) -> None: - """Forward a provider's bounded document payload prefetch capability.""" - provider = self._provider - fn = getattr(provider, "prefetch_document_units", None) - if not callable(fn): - return - if callable(getattr(provider, "document_ids", None)): - fn(doc_id) - return - if str(getattr(provider, "doc_id", "")) == str(doc_id): - fn() - - def release_document_units(self, doc_id: str) -> None: - """Forward release of one document's prefetched payloads.""" - provider = self._provider - fn = getattr(provider, "release_document_units", None) - if not callable(fn): - return - if callable(getattr(provider, "document_ids", None)): - fn(doc_id) - return - if str(getattr(provider, "doc_id", "")) == str(doc_id): - fn() - @dataclass class InMemoryNode: diff --git a/packages/shared-python/shared/services/retrieval/nav/nav_knowhere.py b/packages/shared-python/shared/services/retrieval/nav/nav_knowhere.py index a1387a64..c2c13187 100644 --- a/packages/shared-python/shared/services/retrieval/nav/nav_knowhere.py +++ b/packages/shared-python/shared/services/retrieval/nav/nav_knowhere.py @@ -141,14 +141,6 @@ def knowhere_database_url() -> str: class ChunkStore(Protocol): - def load_document_units( - self, - document_id: str, - section_ids: Sequence[str], - extra_chunk_ids_by_section: Optional[Dict[str, Sequence[str]]] = None, - ) -> List[UnitRow]: - raise NotImplementedError - def load_section_units( self, document_id: str, @@ -219,55 +211,6 @@ def load_section_units( finally: cur.close() - def load_document_units( - self, - document_id: str, - section_ids: Sequence[str], - extra_chunk_ids_by_section: Optional[Dict[str, Sequence[str]]] = None, - ) -> List[UnitRow]: - """Load one document's section payloads in a single ordered query.""" - doc_id = str(document_id).strip() - job_result_id = self._revisions.get(doc_id) - section_values = [str(section_id).strip() for section_id in section_ids if str(section_id).strip()] - extra_ids = [ - str(chunk_id).strip() - for chunk_ids in (extra_chunk_ids_by_section or {}).values() - for chunk_id in chunk_ids - if str(chunk_id).strip() - ] - if not doc_id or not job_result_id or (not section_values and not extra_ids): - return [] - - predicates: list[str] = [] - params: list[object] = [doc_id, job_result_id] - if section_values: - predicates.append("section_id = ANY(%s)") - params.append(section_values) - if extra_ids: - predicates.append("chunk_id = ANY(%s)") - params.append(extra_ids) - - cur = self._connection().cursor() - try: - cur.execute( - "SELECT chunk_id, section_id, chunk_type, content, sort_order, " - "source_chunk_path, file_path, chunk_metadata " - "FROM document_chunks " - "WHERE document_id = %s AND job_result_id = %s AND (" - + " OR ".join(predicates) - + ") ORDER BY section_id, sort_order, chunk_id, id", - params, - ) - units: list[UnitRow] = [] - for row in cur.fetchall(): - section_id = str(row[1]) if row[1] else None - if section_id and (doc_id, section_id) in self._excluded_sections: - continue - units.append(_unit_from_row(row)) - return units - finally: - cur.close() - def close(self) -> None: if self._conn is not None: self._conn.close() @@ -595,53 +538,6 @@ def __init__( ) self._chunk_store = chunk_store self._root_asset_ids = {str(chunk_id) for chunk_id in root_asset_ids} - self._remounted_assets_by_section = { - str(section_id): [str(chunk_id) for chunk_id in chunk_ids] - for section_id, chunk_ids in (remounted_assets_by_section or {}).items() - } - - def prefetch_document_units(self) -> None: - """Load this document's section payloads with one bounded SQL query.""" - section_ids = list(self._sections) - loaded = self._chunk_store.load_document_units( - self.doc_id, - section_ids, - self._remounted_assets_by_section, - ) - by_section: Dict[str, List[UnitRow]] = {} - for unit in loaded: - sid = str(unit.section_id or "").strip() - if sid and sid in self._sections: - by_section.setdefault(sid, []).append(unit) - - units_by_id = {unit.chunk_id: unit for unit in loaded if unit.chunk_id} - for section_id, asset_ids in self._remounted_assets_by_section.items(): - target = by_section.setdefault(section_id, []) - known = {unit.chunk_id for unit in target} - for asset_id in asset_ids: - asset = units_by_id.get(asset_id) - if asset is not None and asset.chunk_id not in known: - target.append(asset) - known.add(asset.chunk_id) - - for section_id in section_ids: - units = by_section.get(section_id, []) - units.sort(key=lambda unit: (unit.sort_order, unit.chunk_id)) - self._units_by_section[section_id] = units - self._loaded_sections.add(section_id) - - for section_id in section_ids: - if is_root_section_path(self.section_path(section_id)): - self._units_by_section[section_id] = [ - unit - for unit in self._units_by_section.get(section_id, ()) - if unit.chunk_id not in self._root_asset_ids - ] - - def release_document_units(self) -> None: - """Release all payloads loaded by the document batch.""" - self._units_by_section.clear() - self._loaded_sections.clear() def _ensure_section_loaded(self, section_id: str) -> None: super()._ensure_section_loaded(section_id) @@ -866,18 +762,6 @@ def release_section_units(self, section_id: str) -> None: if callable(release): release(section_id) - def prefetch_document_units(self, doc_id: str) -> None: - provider = self._docs.get(str(doc_id).strip()) - prefetch = getattr(provider, "prefetch_document_units", None) - if callable(prefetch): - prefetch() - - def release_document_units(self, doc_id: str) -> None: - provider = self._docs.get(str(doc_id).strip()) - release = getattr(provider, "release_document_units", None) - if callable(release): - release() - def address_level(self, node_id: str) -> Optional[NavLevel]: sid = str(node_id or "").strip() if not sid: diff --git a/packages/shared-python/shared/services/retrieval/nav/nav_map_scores.py b/packages/shared-python/shared/services/retrieval/nav/nav_map_scores.py index 0223e232..b88df1fa 100644 --- a/packages/shared-python/shared/services/retrieval/nav/nav_map_scores.py +++ b/packages/shared-python/shared/services/retrieval/nav/nav_map_scores.py @@ -385,21 +385,13 @@ def compute_corpus_map_and_unit_scores( def unit_factory() -> Iterator[ScoreUnitRow]: for document_id in valid_doc_ids: children_map, leaves, titles = tree_by_doc[document_id] - prefetch = getattr(ts, "prefetch_document_units", None) - release = getattr(ts, "release_document_units", None) - if callable(prefetch): - prefetch(document_id) - try: - yield from iter_score_units( - ts, - document_id, - children_map=children_map, - leaves=leaves, - titles=titles, - ) - finally: - if callable(release): - release(document_id) + yield from iter_score_units( + ts, + document_id, + children_map=children_map, + leaves=leaves, + titles=titles, + ) unit_scores = score_unit_stream_hybrid_all(unit_factory, query) From 56cdc59aa75e81b9303d44541d03128966c05af7 Mon Sep 17 00:00:00 2001 From: suguanYang Date: Mon, 31 Aug 2026 11:08:57 +0800 Subject: [PATCH 07/10] Revert "Merge pull request #342 from Ontos-AI/fix/wangbinqi/normalize-native-dsn" This reverts commit 5241307011930d78e8f4e0096c4f6929268b61f3, reversing changes made to a8483326f0975b1a305f4e42eb6b5d7aa99ba362. --- ...t_retrieval_lazy_snapshot_quality_contract.py | 16 ---------------- .../services/retrieval/nav/nav_knowhere.py | 6 +----- 2 files changed, 1 insertion(+), 21 deletions(-) diff --git a/apps/api/tests/contract/test_retrieval_lazy_snapshot_quality_contract.py b/apps/api/tests/contract/test_retrieval_lazy_snapshot_quality_contract.py index 13c5bcf4..782f01f9 100644 --- a/apps/api/tests/contract/test_retrieval_lazy_snapshot_quality_contract.py +++ b/apps/api/tests/contract/test_retrieval_lazy_snapshot_quality_contract.py @@ -11,7 +11,6 @@ NamespaceKnowhereProvider, SectionRow, UnitRow, - knowhere_database_url, ) from shared.services.retrieval.nav.nav_map_scores import ( build_score_units, @@ -175,18 +174,3 @@ def test_streaming_scorer_preserves_duplicate_id_eager_semantics() -> None: } assert score_unit_stream_hybrid_all(lambda: rows, "alpha beta") == eager_scores - - -def test_native_chunk_store_strips_async_driver_from_database_url( - monkeypatch: Any, -) -> None: - monkeypatch.setenv( - "DATABASE_URL", - "postgresql+asyncpg://prod-user:prod-password@db.example/knowhere", - ) - monkeypatch.delenv("KNOWHERE_DATABASE_URL", raising=False) - - assert ( - knowhere_database_url() - == "postgresql://prod-user:prod-password@db.example/knowhere" - ) diff --git a/packages/shared-python/shared/services/retrieval/nav/nav_knowhere.py b/packages/shared-python/shared/services/retrieval/nav/nav_knowhere.py index c2c13187..d73eb178 100644 --- a/packages/shared-python/shared/services/retrieval/nav/nav_knowhere.py +++ b/packages/shared-python/shared/services/retrieval/nav/nav_knowhere.py @@ -132,11 +132,7 @@ def knowhere_database_url() -> str: or str(os.environ.get("DATABASE_URL") or "").strip() ) if configured: - # ``ReadOnlyChunkStore`` uses psycopg2's native connector, which - # accepts libpq URLs but not SQLAlchemy's ``+driver`` suffix. - return configured.replace("postgresql+asyncpg://", "postgresql://", 1).replace( - "postgresql+psycopg2://", "postgresql://", 1 - ) + return configured.replace("postgresql+asyncpg", "postgresql+psycopg2") return _DEFAULT_DSN From 2ac1ae5141a1f7619e64dcde15c2dc40e3eb8346 Mon Sep 17 00:00:00 2001 From: suguanYang Date: Mon, 31 Aug 2026 11:09:07 +0800 Subject: [PATCH 08/10] Revert "Merge pull request #341 from Ontos-AI/feat/wangbinqi/optimize-retrieval-snapshot-load" This reverts commit a8483326f0975b1a305f4e42eb6b5d7aa99ba362, reversing changes made to 9ebee653131fd1704965411e3110b5c68080ed2a. --- ..._add_chunk_revision_section_order_index.py | 43 --- ...etrieval_lazy_snapshot_quality_contract.py | 176 ------------ .../tests/migrations/test_schema_contract.py | 22 -- .../shared/models/database/document.py | 9 - .../services/retrieval/execution/routes.py | 28 +- .../services/retrieval/nav/knowhere_hybrid.py | 206 +------------- .../services/retrieval/nav/nav_hierarchy.py | 6 - .../services/retrieval/nav/nav_knowhere.py | 257 +----------------- .../services/retrieval/nav/nav_map_scores.py | 159 +++++------ .../shared/services/retrieval/nav_snapshot.py | 195 ------------- 10 files changed, 95 insertions(+), 1006 deletions(-) delete mode 100644 apps/api/alembic/versions/0c1d2e3f4a5b_add_chunk_revision_section_order_index.py delete mode 100644 apps/api/tests/contract/test_retrieval_lazy_snapshot_quality_contract.py diff --git a/apps/api/alembic/versions/0c1d2e3f4a5b_add_chunk_revision_section_order_index.py b/apps/api/alembic/versions/0c1d2e3f4a5b_add_chunk_revision_section_order_index.py deleted file mode 100644 index 300dc99c..00000000 --- a/apps/api/alembic/versions/0c1d2e3f4a5b_add_chunk_revision_section_order_index.py +++ /dev/null @@ -1,43 +0,0 @@ -"""Add the index used by lazy map-nav section loads.""" - -from __future__ import annotations - -from alembic import op - - -revision = "0c1d2e3f4a5b" -down_revision = "fbf0c1d2e3f4" -branch_labels = None -depends_on = None - -_INDEX_NAME = "idx_document_chunks_revision_section_order" - - -def upgrade() -> None: - external_transaction = bool( - op.get_context().opts.get("knowhere_external_transaction", False) - ) - if external_transaction: - op.execute( - f"CREATE INDEX IF NOT EXISTS {_INDEX_NAME} " - "ON document_chunks " - "(document_id, job_result_id, section_id, sort_order, chunk_id, id)" - ) - return - with op.get_context().autocommit_block(): - op.execute( - f"CREATE INDEX CONCURRENTLY IF NOT EXISTS {_INDEX_NAME} " - "ON document_chunks " - "(document_id, job_result_id, section_id, sort_order, chunk_id, id)" - ) - - -def downgrade() -> None: - external_transaction = bool( - op.get_context().opts.get("knowhere_external_transaction", False) - ) - if external_transaction: - op.execute(f"DROP INDEX IF EXISTS {_INDEX_NAME}") - return - with op.get_context().autocommit_block(): - op.execute(f"DROP INDEX CONCURRENTLY IF EXISTS {_INDEX_NAME}") diff --git a/apps/api/tests/contract/test_retrieval_lazy_snapshot_quality_contract.py b/apps/api/tests/contract/test_retrieval_lazy_snapshot_quality_contract.py deleted file mode 100644 index 782f01f9..00000000 --- a/apps/api/tests/contract/test_retrieval_lazy_snapshot_quality_contract.py +++ /dev/null @@ -1,176 +0,0 @@ -from __future__ import annotations - -from dataclasses import dataclass -from collections.abc import Sequence -from typing import Any - -from shared.services.retrieval.nav.nav_hierarchy import ProviderToolSpace -from shared.services.retrieval.nav.nav_knowhere import ( - KnowhereProvider, - LazyKnowhereProvider, - NamespaceKnowhereProvider, - SectionRow, - UnitRow, -) -from shared.services.retrieval.nav.nav_map_scores import ( - build_score_units, - compute_corpus_map_and_unit_scores, -) -from shared.services.retrieval.nav.knowhere_hybrid import ( - ScoreUnitRow, - score_rows_hybrid_all, - score_unit_stream_hybrid_all, -) - - -@dataclass -class _FakeChunkStore: - units_by_section: dict[str, list[UnitRow]] - - def load_section_units( - self, - document_id: str, - section_id: str, - extra_chunk_ids: Sequence[str] = (), - ) -> list[UnitRow]: - del document_id - units = list(self.units_by_section.get(section_id, ())) - known = {unit.chunk_id for unit in units} - for units_in_section in self.units_by_section.values(): - for unit in units_in_section: - if unit.chunk_id in extra_chunk_ids and unit.chunk_id not in known: - units.append(unit) - return units - - def close(self) -> None: - return None - - -def _providers() -> tuple[ProviderToolSpace, ProviderToolSpace]: - sections = [ - SectionRow("root", None, "Root", "Root", 0, "", 0), - SectionRow("section", "root", "Root / Section", "Section", 1, "", 1), - SectionRow("leaf", "section", "Root / Section / Leaf", "Leaf", 2, "", 2), - ] - text = UnitRow( - "duplicate-chunk", - "leaf", - "text", - "alpha retrieval evidence", - 1, - metadata={"connect_to": [{"target": "asset-1", "relation": "embeds"}]}, - ) - asset = UnitRow( - "asset-1", - "root", - "image", - "", - 2, - file_path="images/asset.png", - metadata={"summary": "supporting image"}, - ) - eager = NamespaceKnowhereProvider( - [KnowhereProvider(doc_id="doc", sections=sections, units=[text, asset])], - titles={"doc": "document"}, - ) - store = _FakeChunkStore({"root": [asset], "leaf": [text]}) - lazy = NamespaceKnowhereProvider( - [ - LazyKnowhereProvider( - doc_id="doc", - sections=sections, - chunk_store=store, - known_chunk_ids=[text.chunk_id, asset.chunk_id], - root_asset_ids=[asset.chunk_id], - remounted_assets_by_section={"leaf": [asset.chunk_id]}, - ) - ], - titles={"doc": "document"}, - chunk_owner_by_id={"duplicate-chunk": "doc", "asset-1": "doc"}, - ) - return ProviderToolSpace(eager), ProviderToolSpace(lazy) - - -def test_lazy_provider_preserves_score_units_and_scores() -> None: - eager, lazy = _providers() - - assert build_score_units(eager, "doc") == build_score_units(lazy, "doc") - assert compute_corpus_map_and_unit_scores( - eager, doc_ids=["doc"], query="alpha retrieval" - ) == compute_corpus_map_and_unit_scores( - lazy, doc_ids=["doc"], query="alpha retrieval" - ) - - lazy_provider = lazy._provider - self_units = getattr(lazy_provider, "self_units") - assert [unit.chunk_id for unit in self_units("leaf")] == [ - "duplicate-chunk", - "asset-1", - ] - - -def test_streaming_scorer_preserves_exact_eager_scores() -> None: - rows: list[ScoreUnitRow] = [ - { - "chunk_id": "unit-a", - "path_search_text": "root alpha", - "content_search_text": "alpha alpha evidence", - "term_search_text": "alpha alpha evidence root", - }, - { - "chunk_id": "unit-b", - "path_search_text": "root beta", - "content_search_text": "beta evidence", - "term_search_text": "beta evidence root", - }, - { - "chunk_id": "unit-c", - "path_search_text": "root common", - "content_search_text": "common evidence", - "term_search_text": "common evidence root", - }, - ] - eager_rows: list[dict[str, Any]] = [dict(row) for row in rows] - eager_scores = { - str(row["chunk_id"]): float(row["score"]) - for row in score_rows_hybrid_all(eager_rows, "alpha evidence") - } - replay_count: int = 0 - - def unit_factory() -> Sequence[ScoreUnitRow]: - nonlocal replay_count - replay_count += 1 - return rows - - assert score_unit_stream_hybrid_all(unit_factory, "alpha evidence") == eager_scores - assert replay_count == 1 - - -def test_streaming_scorer_preserves_duplicate_id_eager_semantics() -> None: - rows: list[ScoreUnitRow] = [ - { - "chunk_id": "duplicate", - "path_search_text": "alpha", - "content_search_text": "alpha", - "term_search_text": "alpha", - }, - { - "chunk_id": "duplicate", - "path_search_text": "beta", - "content_search_text": "beta", - "term_search_text": "beta", - }, - { - "chunk_id": "other", - "path_search_text": "alpha beta", - "content_search_text": "alpha beta", - "term_search_text": "alpha beta", - }, - ] - eager_rows: list[dict[str, Any]] = [dict(row) for row in rows] - eager_scores = { - str(row["chunk_id"]): float(row["score"]) - for row in score_rows_hybrid_all(eager_rows, "alpha beta") - } - - assert score_unit_stream_hybrid_all(lambda: rows, "alpha beta") == eager_scores diff --git a/apps/api/tests/migrations/test_schema_contract.py b/apps/api/tests/migrations/test_schema_contract.py index 3ead1217..24047dcc 100644 --- a/apps/api/tests/migrations/test_schema_contract.py +++ b/apps/api/tests/migrations/test_schema_contract.py @@ -233,28 +233,6 @@ def test_should_index_document_chunks_in_snapshot_pagination_order( ) -def test_should_index_document_chunks_in_lazy_section_order( - migrated_head_engine: Engine, -) -> None: - with migrated_head_engine.begin() as connection: - index_definition = connection.execute( - text( - """ - SELECT indexdef - FROM pg_indexes - WHERE schemaname = current_schema() - AND tablename = 'document_chunks' - AND indexname = 'idx_document_chunks_revision_section_order' - """ - ) - ).scalar_one() - - assert ( - "(document_id, job_result_id, section_id, sort_order, chunk_id, id)" - in str(index_definition) - ) - - def test_should_upgrade_with_a_caller_owned_connection( alembic_engine: Engine, ) -> None: diff --git a/packages/shared-python/shared/models/database/document.py b/packages/shared-python/shared/models/database/document.py index b2381996..f9681e17 100644 --- a/packages/shared-python/shared/models/database/document.py +++ b/packages/shared-python/shared/models/database/document.py @@ -210,15 +210,6 @@ class DocumentChunk(Base): "chunk_id", "id", ), - Index( - "idx_document_chunks_revision_section_order", - "document_id", - "job_result_id", - "section_id", - "sort_order", - "chunk_id", - "id", - ), Index("idx_document_chunks_section", "section_id"), Index( "idx_chunk_content_search_tsv", diff --git a/packages/shared-python/shared/services/retrieval/execution/routes.py b/packages/shared-python/shared/services/retrieval/execution/routes.py index ad2268ab..e67a1ead 100644 --- a/packages/shared-python/shared/services/retrieval/execution/routes.py +++ b/packages/shared-python/shared/services/retrieval/execution/routes.py @@ -186,7 +186,6 @@ async def _run_mapnav_route( namespace=context.namespace, exclude_document_ids=context.exclude_document_ids, exclude_sections=context.exclude_sections, - lazy=True, ) # Small-corpus count / snapshot reads may leave a checkout; drop it before @@ -197,22 +196,19 @@ async def _run_mapnav_route( cfg = build_nav_config() toolspace = ProviderToolSpace(snapshot.provider) - try: - episode = await asyncio.to_thread( - run_nav_episode, - None, - context.query, - corpus_doc_ids=list(snapshot.document_ids), - budget_chars=budget, - compose_answer=False, - policy="llm", - config=cfg, - toolspace=toolspace, - ) + episode = await asyncio.to_thread( + run_nav_episode, + None, + context.query, + corpus_doc_ids=list(snapshot.document_ids), + budget_chars=budget, + compose_answer=False, + policy="llm", + config=cfg, + toolspace=toolspace, + ) - refs, score_by_chunk_id = build_referenced_chunks(episode, snapshot) - finally: - snapshot.close() + refs, score_by_chunk_id = build_referenced_chunks(episode, snapshot) async with open_fresh_database_context() as final_db: resolved = await resolve_workflow_references( diff --git a/packages/shared-python/shared/services/retrieval/nav/knowhere_hybrid.py b/packages/shared-python/shared/services/retrieval/nav/knowhere_hybrid.py index 7d8b195b..d56c72fb 100644 --- a/packages/shared-python/shared/services/retrieval/nav/knowhere_hybrid.py +++ b/packages/shared-python/shared/services/retrieval/nav/knowhere_hybrid.py @@ -9,11 +9,7 @@ import os import re -import math -from collections import Counter -from collections.abc import Callable, Iterable -from dataclasses import dataclass -from typing import Any, Dict, List, Mapping, Optional, Sequence, Tuple, TypedDict +from typing import Any, Dict, List, Optional, Sequence, Tuple RRF_K = 60 CHANNEL_WEIGHT_PATH = 1.0 @@ -22,19 +18,6 @@ INTERNAL_RECALL_K_MULTIPLIER = 2 -class ScoreUnitRow(TypedDict, total=False): - """Compact scoring-unit shape shared by eager and streaming scorers.""" - - chunk_id: str - section_id: str - kind: str - content: str - path_text: str - path_search_text: str - content_search_text: str - term_search_text: str - - def tokenize_for_retrieval(text: str, *, dedupe: bool = True) -> List[str]: tokens = re.findall(r"[a-z0-9_]+|[\u4e00-\u9fff]", str(text or "").lower()) if not dedupe: @@ -85,7 +68,7 @@ def build_term_search_text(content: str, *, path_text: Optional[str] = None) -> return combined -def _get_search_tokens(row: Mapping[str, object], *, search_field: str) -> List[str]: +def _get_search_tokens(row: dict[str, Any], *, search_field: str) -> List[str]: return [token for token in str(row.get(search_field) or "").split() if token] @@ -607,188 +590,3 @@ def _rows_from_scores(score_by_id: Dict[str, float]) -> List[dict[str, Any]]: row["term_channel_score"] = float(term_channel.get(uid, 0.0) or 0.0) out_rows.append(row) return out_rows - - -def score_unit_stream_hybrid_all( - unit_factory: Callable[[], Iterable[ScoreUnitRow]], - query: str, -) -> Dict[str, float]: - """Score replayable units without retaining their payloads. - - This is the corpus scorer used by map-nav. It mirrors the active BM25 and - weighted-RRF implementation, but keeps only token statistics, identifiers, - and final scores between bounded provider reads. - """ - query_tokens = tokenize_query_for_ranker(query) - path_stats = _StreamingBm25Stats.empty() - content_stats = _StreamingBm25Stats.empty() - units: List[_StreamingUnit] = [] - query_token_set = set(query_tokens) - query_lower = query.lower().strip() - for row in unit_factory(): - unit_id = str(row.get("chunk_id") or "").strip() - if not unit_id: - continue - path_tokens = _get_search_tokens(row, search_field="path_search_text") - content_tokens = _get_search_tokens(row, search_field="content_search_text") - path_stats.observe(path_tokens) - content_stats.observe(content_tokens) - path_frequencies = Counter(path_tokens) - content_frequencies = Counter(content_tokens) - term_score = 0.0 - if query_lower: - haystack = str(row.get("term_search_text") or "").lower() - if query_lower in haystack: - term_score = 100.0 - else: - hit_count = sum(1 for token in query_tokens if token in haystack) - if hit_count > 0: - term_score = float(hit_count) - units.append( - _StreamingUnit( - unit_id=unit_id, - path_length=len(path_tokens), - content_length=len(content_tokens), - path_frequencies={ - token: path_frequencies[token] - for token in query_token_set - if path_frequencies[token] - }, - content_frequencies={ - token: content_frequencies[token] - for token in query_token_set - if content_frequencies[token] - }, - term_score=term_score, - ) - ) - path_stats.finalize() - content_stats.finalize() - path_by_id: Dict[str, float] = {} - content_by_id: Dict[str, float] = {} - term_by_id: Dict[str, float] = {} - unit_ids = list(dict.fromkeys(unit.unit_id for unit in units)) - for unit in units: - path_score = path_stats.score( - unit.path_length, unit.path_frequencies, query_tokens - ) - content_score = content_stats.score( - unit.content_length, unit.content_frequencies, query_tokens - ) - path_by_id[unit.unit_id] = path_score - content_by_id[unit.unit_id] = content_score - term_by_id[unit.unit_id] = unit.term_score - - path_rows = [ - (score, unit_id) - for unit_id, score in path_by_id.items() - if score > 0.0 - ] - content_rows = [ - (score, unit_id) - for unit_id, score in content_by_id.items() - if score > 0.0 - ] - term_rows = [ - (score, unit_id) - for unit_id, score in term_by_id.items() - if score > 0.0 - ] - - path_rows.sort(key=lambda item: (-item[0], item[1])) - content_rows.sort(key=lambda item: (-item[0], item[1])) - term_rows.sort(key=lambda item: (-item[0], item[1])) - path_weight, content_weight, term_weight = map_channel_weights() - rrf_k = int( - os.environ.get( - "NAV_MAP_RRF_K", - os.environ.get("NAV_DISCOVERY_RRF_K", str(RRF_K)), - ).strip() - or RRF_K - ) - fused: Dict[str, float] = {unit_id: 0.0 for unit_id in unit_ids} - for rank, (_score, unit_id) in enumerate(path_rows): - fused[unit_id] = fused.get(unit_id, 0.0) + path_weight / (rrf_k + rank + 1) - for rank, (_score, unit_id) in enumerate(content_rows): - fused[unit_id] = fused.get(unit_id, 0.0) + content_weight / (rrf_k + rank + 1) - for rank, (_score, unit_id) in enumerate(term_rows): - fused[unit_id] = fused.get(unit_id, 0.0) + term_weight / (rrf_k + rank + 1) - return {unit_id: round(score, 6) for unit_id, score in fused.items()} - - -@dataclass(frozen=True) -class _StreamingUnit: - unit_id: str - path_length: int - content_length: int - path_frequencies: Mapping[str, int] - content_frequencies: Mapping[str, int] - term_score: float - - -class _StreamingBm25Stats: - """Exact BM25Okapi corpus statistics collected without row retention.""" - - def __init__(self) -> None: - self.document_count: int = 0 - self.document_frequency: Counter[str] = Counter() - self.total_length: int = 0 - self.average_length: float = 0.0 - self.idf_by_token: Dict[str, float] = {} - - @classmethod - def empty(cls) -> "_StreamingBm25Stats": - return cls() - - def observe(self, tokens: List[str]) -> None: - if not tokens: - return - self.document_count += 1 - self.total_length += len(tokens) - self.document_frequency.update(set(tokens)) - - def finalize(self) -> None: - self.average_length = ( - self.total_length / self.document_count if self.document_count else 0.0 - ) - idf_by_token: Dict[str, float] = {} - idf_sum = 0.0 - negative_tokens: List[str] = [] - for token, frequency in self.document_frequency.items(): - idf = math.log(self.document_count - frequency + 0.5) - math.log( - frequency + 0.5 - ) - idf_by_token[token] = idf - idf_sum += idf - if idf < 0.0: - negative_tokens.append(token) - average_idf = idf_sum / len(idf_by_token) if idf_by_token else 0.0 - epsilon_floor = 0.25 * average_idf - for token in negative_tokens: - idf_by_token[token] = epsilon_floor - self.idf_by_token = idf_by_token - - def score( - self, - document_length: int, - frequencies: Dict[str, int], - query_tokens: List[str], - ) -> float: - if ( - not frequencies - or not query_tokens - or not self.document_count - or self.average_length <= 0.0 - ): - return 0.0 - denominator_base = 1.5 * ( - 1.0 - 0.75 + 0.75 * document_length / self.average_length - ) - score = 0.0 - for token in query_tokens: - frequency = frequencies.get(token, 0) - if not frequency: - continue - idf = self.idf_by_token.get(token, 0.0) - score += idf * (frequency * 2.5 / (frequency + denominator_base)) - return score diff --git a/packages/shared-python/shared/services/retrieval/nav/nav_hierarchy.py b/packages/shared-python/shared/services/retrieval/nav/nav_hierarchy.py index 93b35598..fafa66d2 100644 --- a/packages/shared-python/shared/services/retrieval/nav/nav_hierarchy.py +++ b/packages/shared-python/shared/services/retrieval/nav/nav_hierarchy.py @@ -231,12 +231,6 @@ def read_chunks(self, section_id: str, query: str, *, doc_id: str, k: int) -> Li del section_id, query, doc_id, k return [] - def release_section_units(self, section_id: str) -> None: - """Release one lazy section without discarding the hierarchy.""" - release = getattr(self._provider, "release_section_units", None) - if callable(release): - release(section_id) - @dataclass class InMemoryNode: diff --git a/packages/shared-python/shared/services/retrieval/nav/nav_knowhere.py b/packages/shared-python/shared/services/retrieval/nav/nav_knowhere.py index d73eb178..19caf893 100644 --- a/packages/shared-python/shared/services/retrieval/nav/nav_knowhere.py +++ b/packages/shared-python/shared/services/retrieval/nav/nav_knowhere.py @@ -22,7 +22,7 @@ import os from dataclasses import dataclass, field -from typing import Any, Callable, Dict, Iterable, List, Optional, Protocol, Sequence, Set, Tuple +from typing import Any, Dict, List, Optional, Sequence, Set, Tuple from .nav_address import NavLevel from .nav_hierarchy import NodeMeta @@ -127,125 +127,7 @@ def _connect_to_targets(metadata: Dict[str, Any]) -> List[str]: def knowhere_database_url() -> str: - configured = ( - str(os.environ.get("KNOWHERE_DATABASE_URL") or "").strip() - or str(os.environ.get("DATABASE_URL") or "").strip() - ) - if configured: - return configured.replace("postgresql+asyncpg", "postgresql+psycopg2") - return _DEFAULT_DSN - - -class ChunkStore(Protocol): - def load_section_units( - self, - document_id: str, - section_id: str, - extra_chunk_ids: Sequence[str] = (), - ) -> List[UnitRow]: - raise NotImplementedError - - def close(self) -> None: - raise NotImplementedError - - -class ReadOnlyChunkStore: - """Episode-local, revision-pinned loader for lazy map-nav chunks.""" - - def __init__( - self, - *, - dsn: str, - revisions: Dict[str, str], - excluded_sections: Optional[Iterable[Tuple[str, str]]] = None, - ) -> None: - self._dsn = str(dsn) - self._revisions = dict(revisions) - self._excluded_sections = set(excluded_sections or ()) - self._conn: Optional[_SyncConnection] = None - - def _connection(self) -> "_SyncConnection": - if self._conn is None: - self._conn = _connect(self._dsn) - self._conn.set_session(readonly=True, autocommit=True) - return self._conn - - def load_section_units( - self, - document_id: str, - section_id: str, - extra_chunk_ids: Sequence[str] = (), - ) -> List[UnitRow]: - doc_id = str(document_id).strip() - sid = str(section_id).strip() - job_result_id = self._revisions.get(doc_id) - if not doc_id or not sid or not job_result_id or (doc_id, sid) in self._excluded_sections: - return [] - cur = self._connection().cursor() - try: - ids = [str(chunk_id).strip() for chunk_id in extra_chunk_ids if str(chunk_id).strip()] - if ids: - cur.execute( - "SELECT chunk_id, section_id, chunk_type, content, sort_order, " - "source_chunk_path, file_path, chunk_metadata " - "FROM document_chunks " - "WHERE document_id = %s AND job_result_id = %s " - "AND (section_id = %s OR chunk_id = ANY(%s)) " - "ORDER BY sort_order, chunk_id, id", - (doc_id, job_result_id, sid, ids), - ) - else: - cur.execute( - "SELECT chunk_id, section_id, chunk_type, content, sort_order, " - "source_chunk_path, file_path, chunk_metadata " - "FROM document_chunks " - "WHERE document_id = %s AND job_result_id = %s AND section_id = %s " - "ORDER BY sort_order, chunk_id, id", - (doc_id, job_result_id, sid), - ) - return [_unit_from_row(row) for row in cur.fetchall()] - finally: - cur.close() - - def close(self) -> None: - if self._conn is not None: - self._conn.close() - self._conn = None - - -class _SyncCursor(Protocol): - def execute(self, query: str, params: Sequence[object]) -> None: - raise NotImplementedError - - def fetchall(self) -> Sequence[Sequence[object]]: - raise NotImplementedError - - def close(self) -> None: - raise NotImplementedError - - -class _SyncConnection(Protocol): - def set_session(self, *, readonly: bool, autocommit: bool) -> None: - raise NotImplementedError - - def cursor(self) -> _SyncCursor: - raise NotImplementedError - - def close(self) -> None: - raise NotImplementedError - - -def _unit_from_row(row: Sequence[object]) -> UnitRow: - return UnitRow( - chunk_id=str(row[0] or ""), - section_id=str(row[1]) if row[1] else None, - chunk_type=str(row[2] or "text"), - content=str(row[3] or ""), - sort_order=int(row[4] or 0), - source_chunk_path=str(row[5] or ""), - file_path=str(row[6] or ""), - metadata=_as_meta(row[7]), - ) + return str(os.environ.get("KNOWHERE_DATABASE_URL") or "").strip() or _DEFAULT_DSN class KnowhereProvider: @@ -257,12 +139,8 @@ def __init__( doc_id: str, sections: Sequence[SectionRow], units: Sequence[UnitRow], - lazy_loader: Optional[Callable[[str], Sequence[UnitRow]]] = None, - known_chunk_ids: Optional[Sequence[str]] = None, ) -> None: self.doc_id = str(doc_id) - self._lazy_loader = lazy_loader - self._loaded_sections: Set[str] = set() self._sections: Dict[str, SectionRow] = {s.section_id: s for s in sections} self._children: Dict[str, List[str]] = {} self._roots: List[str] = [] @@ -279,12 +157,6 @@ def __init__( self._units_by_section: Dict[str, List[UnitRow]] = {} self._chunk_ids: Set[str] = set() - if known_chunk_ids: - self._chunk_ids.update( - str(chunk_id).strip() - for chunk_id in known_chunk_ids - if str(chunk_id).strip() - ) for unit in sorted(units, key=lambda u: (u.sort_order, u.chunk_id)): sid = unit.section_id if not sid or sid not in self._sections: @@ -294,21 +166,6 @@ def __init__( self._chunk_ids.add(unit.chunk_id) self._remount_root_assets() - def _ensure_section_loaded(self, section_id: str) -> None: - if self._lazy_loader is None or section_id in self._loaded_sections: - return - loaded = list(self._lazy_loader(section_id) or ()) - self._loaded_sections.add(section_id) - if not loaded: - return - current = self._units_by_section.setdefault(section_id, []) - known = {unit.chunk_id for unit in current} - for unit in loaded: - if unit.chunk_id and unit.chunk_id not in known: - current.append(unit) - known.add(unit.chunk_id) - current.sort(key=lambda unit: (unit.sort_order, unit.chunk_id)) - def _remount_root_assets(self) -> None: """Reattach Root-FK image|table units to host sections via ``connect_to``. @@ -430,23 +287,12 @@ def content(self, section_id: str) -> str: return "\n".join(self.unit_text(u) for u in units if self.unit_text(u)) def self_units(self, section_id: str) -> List[UnitRow]: - self._ensure_section_loaded(section_id) return list(self._units_by_section.get(section_id, ())) - def release_section_units(self, section_id: str) -> None: - """Drop one section's loaded payload while keeping its structure.""" - if self._lazy_loader is None: - return - sid = str(section_id or "").strip() - if not sid: - return - self._units_by_section.pop(sid, None) - self._loaded_sections.discard(sid) - def subtree_units(self, section_id: str) -> List[UnitRow]: - out = list(self.self_units(section_id)) + out = list(self._units_by_section.get(section_id, ())) for cid in self.relations(section_id)[1]: - out.extend(self.self_units(cid)) + out.extend(self._units_by_section.get(cid, ())) out.sort(key=lambda u: (u.sort_order, u.chunk_id)) return out @@ -504,59 +350,8 @@ def summaries(self) -> Dict[str, str]: def all_section_ids(self) -> List[str]: return list(self._sections) - def chunk_count(self) -> int: - return len(self._chunk_ids) - - -class LazyKnowhereProvider(KnowhereProvider): - """Hierarchy provider that loads full chunk rows only on first access.""" - - def __init__( - self, - *, - doc_id: str, - sections: Sequence[SectionRow], - chunk_store: ChunkStore, - known_chunk_ids: Sequence[str], - root_asset_ids: Sequence[str] = (), - remounted_assets_by_section: Optional[Dict[str, Sequence[str]]] = None, - ) -> None: - super().__init__( - doc_id=doc_id, - sections=sections, - units=(), - lazy_loader=lambda section_id: chunk_store.load_section_units( - doc_id, - section_id, - (remounted_assets_by_section or {}).get(section_id, ()), - ), - known_chunk_ids=known_chunk_ids, - ) - self._chunk_store = chunk_store - self._root_asset_ids = {str(chunk_id) for chunk_id in root_asset_ids} - - def _ensure_section_loaded(self, section_id: str) -> None: - super()._ensure_section_loaded(section_id) - if ( - section_id in self._units_by_section - and self._root_asset_ids - and is_root_section_path(self.section_path(section_id)) - ): - self._units_by_section[section_id] = [ - unit - for unit in self._units_by_section[section_id] - if unit.chunk_id not in self._root_asset_ids - ] - - def close(self) -> None: - self._chunk_store.close() - def release_loaded_units(self) -> None: - self._units_by_section.clear() - self._loaded_sections.clear() - - -def _connect(dsn: str) -> _SyncConnection: +def _connect(dsn: str): import psycopg2 return psycopg2.connect(dsn) @@ -704,7 +499,6 @@ def __init__( providers: Sequence[KnowhereProvider], *, titles: Optional[Dict[str, str]] = None, - chunk_owner_by_id: Optional[Dict[str, str]] = None, ) -> None: self._docs: Dict[str, KnowhereProvider] = { p.doc_id: p for p in providers if p.doc_id @@ -720,44 +514,14 @@ def __init__( for doc_id, provider in self._docs.items(): for sid in provider.all_section_ids(): self._section_owner[sid] = doc_id - if chunk_owner_by_id: - self._chunk_owner.update( - { - str(chunk_id): str(doc_id) - for chunk_id, doc_id in chunk_owner_by_id.items() - if str(chunk_id).strip() and str(doc_id).strip() - } - ) - else: - for doc_id, provider in self._docs.items(): - for sid in provider.all_section_ids(): - for unit in provider.self_units(sid): - if unit.chunk_id: - self._chunk_owner[unit.chunk_id] = doc_id + for sid in provider.all_section_ids(): + for unit in provider.self_units(sid): + if unit.chunk_id: + self._chunk_owner[unit.chunk_id] = doc_id def document_ids(self) -> List[str]: return list(self._docs) - def close(self) -> None: - for provider in self._docs.values(): - close = getattr(provider, "close", None) - if callable(close): - close() - - def release_loaded_units(self) -> None: - for provider in self._docs.values(): - release = getattr(provider, "release_loaded_units", None) - if callable(release): - release() - - def release_section_units(self, section_id: str) -> None: - owner = self._section_owner.get(str(section_id or "").strip()) - if not owner: - return - release = getattr(self._docs[owner], "release_section_units", None) - if callable(release): - release(section_id) - def address_level(self, node_id: str) -> Optional[NavLevel]: sid = str(node_id or "").strip() if not sid: @@ -798,8 +562,7 @@ def node_meta(self, section_id: str) -> NodeMeta: sid = str(section_id or "").strip() if sid in self._docs: provider = self._docs[sid] - count_fn = getattr(provider, "chunk_count", None) - n_chunks = int(count_fn()) if callable(count_fn) else sum( + n_chunks = sum( len(provider.self_units(sec)) for sec in provider.all_section_ids() ) return NodeMeta( diff --git a/packages/shared-python/shared/services/retrieval/nav/nav_map_scores.py b/packages/shared-python/shared/services/retrieval/nav/nav_map_scores.py index b88df1fa..d1049adf 100644 --- a/packages/shared-python/shared/services/retrieval/nav/nav_map_scores.py +++ b/packages/shared-python/shared/services/retrieval/nav/nav_map_scores.py @@ -1,15 +1,13 @@ from __future__ import annotations -from collections.abc import Iterator from typing import Any, Dict, List, Optional, Sequence, Set, Tuple from .knowhere_hybrid import ( - ScoreUnitRow, build_content_search_text, build_path_search_text, build_term_search_text, + score_dense_channel, score_rows_hybrid_all, - score_unit_stream_hybrid_all, ) @@ -28,10 +26,6 @@ def _line_content(ts: Any, section_id: str, doc_id: str) -> str: idx = getattr(ts, "_idx", None) b = getattr(idx, "_bundles", {}).get(doc_id) if idx is not None else None if b is None: - path_fn = getattr(ts, "path_titles", None) - if callable(path_fn): - path = str(path_fn(section_id, doc_id) or "").strip() - return path.rsplit(" / ", 1)[-1] if path else "" st = ts.get_structure(section_id) return str(st.get("preview") or "").strip() loc = getattr(idx, "_node_to_doc_line", {}).get(section_id) @@ -179,6 +173,47 @@ def score_node(section_id: str) -> float: return map_scores +def _score_dense_units_by_doc( + units_by_doc: Sequence[Tuple[str, List[dict]]], + query: str, + *, + namespace: Optional[str], +) -> Dict[str, Optional[Dict[str, float]]]: + """Read per-doc vector caches, returning raw cosine scores for global fusion. + + Dense cosine is independently comparable across documents. Partitioning only + preserves the existing per-doc disk cache; no ranking or normalization occurs + here. If any partition fails, that whole channel falls back to global BM25. + """ + dense_by_channel: Dict[str, Optional[Dict[str, float]]] = {} + for channel, text_field in (("path", "path_text"), ("content", "content")): + score_by_id: Dict[str, float] = {} + complete = True + for doc_id, units in units_by_doc: + if not units: + continue + unit_ids = [str(unit["chunk_id"]) for unit in units] + scores = score_dense_channel( + [str(unit.get(text_field) or "") for unit in units], + query, + unit_ids=unit_ids, + doc_id=doc_id, + channel=channel, + namespace=namespace, + ) + if scores is None or len(scores) != len(unit_ids): + complete = False + break + score_by_id.update( + { + unit_id: float(scores[index]) + for index, unit_id in enumerate(unit_ids) + } + ) + dense_by_channel[channel] = score_by_id if complete else None + return dense_by_channel + + def build_score_units(ts: Any, doc_id: str, root_ids: Optional[Sequence[str]] = None) -> List[dict]: """Build leaf (+ interstitial self_only) units for hybrid scoring.""" if root_ids is None: @@ -241,66 +276,6 @@ def build_score_units(ts: Any, doc_id: str, root_ids: Optional[Sequence[str]] = return units -def iter_score_units( - ts: Any, - doc_id: str, - *, - children_map: Dict[str, List[str]], - leaves: Set[str], - titles: Dict[str, str], -) -> Iterator[ScoreUnitRow]: - """Yield scoring units while releasing each lazy section after use.""" - release = getattr(ts, "release_section_units", None) - - for leaf_id in sorted(leaves): - try: - content = _section_body_text(ts, leaf_id, doc_id) or ( - titles.get(leaf_id) or _line_content(ts, leaf_id, doc_id) - ) - if content: - path_text = _ancestor_path_titles(ts, leaf_id, doc_id) - title = titles.get(leaf_id) or _line_content(ts, leaf_id, doc_id) - yield { - "chunk_id": leaf_id, - "section_id": leaf_id, - "kind": "leaf", - "content": content, - "path_text": path_text, - "path_search_text": build_path_search_text( - section_path=path_text, section_title=title or content - ), - "content_search_text": build_content_search_text(content), - "term_search_text": build_term_search_text(content, path_text=path_text), - } - finally: - if callable(release): - release(leaf_id) - - for sid, kids in children_map.items(): - if not kids: - continue - try: - self_text, has_interstitial = _self_only_text(ts, sid, doc_id) - if not has_interstitial or not self_text: - continue - path_text = _ancestor_path_titles(ts, sid, doc_id) - yield { - "chunk_id": f"{sid}__self", - "section_id": sid, - "kind": "self_only", - "content": self_text, - "path_text": path_text, - "path_search_text": build_path_search_text( - section_path=path_text, section_title=titles.get(sid) or "" - ), - "content_search_text": build_content_search_text(self_text), - "term_search_text": build_term_search_text(self_text, path_text=path_text), - } - finally: - if callable(release): - release(sid) - - def compute_map_scores( ts: Any, *, @@ -371,33 +346,41 @@ def compute_corpus_map_and_unit_scores( seen_doc_ids.add(doc_id) valid_doc_ids.append(doc_id) - del namespace # Dense scoring is intentionally disabled for the corpus path. + ns = namespace + if not ns: + import os + + ns = os.environ.get("NAV_MAP_UNIT_CACHE_NS", "").strip() or None - tree_by_doc: Dict[ - str, - Tuple[Dict[str, List[str]], Set[str], Dict[str, str]], - ] = {} + tree_by_doc: Dict[str, Tuple[Dict[str, List[str]], Set[str]]] = {} + units_by_doc: List[Tuple[str, List[dict]]] = [] + all_units: List[dict] = [] for doc_id in valid_doc_ids: root_ids = list(ts.sections_for_doc(doc_id)) - children_map, leaves, titles = _walk_tree(ts, doc_id, root_ids) - tree_by_doc[doc_id] = (children_map, leaves, titles) - - def unit_factory() -> Iterator[ScoreUnitRow]: - for document_id in valid_doc_ids: - children_map, leaves, titles = tree_by_doc[document_id] - yield from iter_score_units( - ts, - document_id, - children_map=children_map, - leaves=leaves, - titles=titles, - ) - - unit_scores = score_unit_stream_hybrid_all(unit_factory, query) + children_map, leaves, _titles = _walk_tree(ts, doc_id, root_ids) + units = build_score_units(ts, doc_id, root_ids=root_ids) + tree_by_doc[doc_id] = (children_map, leaves) + units_by_doc.append((doc_id, units)) + all_units.extend(units) + + dense_scores = _score_dense_units_by_doc( + units_by_doc, + query, + namespace=ns, + ) + scored = score_rows_hybrid_all( + all_units, + query, + dense_scores_by_channel=dense_scores, + ) + unit_scores = { + str(row.get("chunk_id") or ""): float(row.get("score") or 0.0) + for row in scored + } map_scores: Dict[str, float] = {} for doc_id in valid_doc_ids: - children_map, leaves, _titles = tree_by_doc[doc_id] + children_map, leaves = tree_by_doc[doc_id] doc_map_scores = _pool_unit_scores_to_tree( children_map, leaves, unit_scores ) diff --git a/packages/shared-python/shared/services/retrieval/nav_snapshot.py b/packages/shared-python/shared/services/retrieval/nav_snapshot.py index 1edfd374..7dc60c6f 100644 --- a/packages/shared-python/shared/services/retrieval/nav_snapshot.py +++ b/packages/shared-python/shared/services/retrieval/nav_snapshot.py @@ -7,7 +7,6 @@ from __future__ import annotations -import json from dataclasses import dataclass from typing import Any, Protocol @@ -17,13 +16,10 @@ from shared.models.database.document import Document, DocumentChunk, DocumentSection from shared.models.database.job_result import JobResult from shared.services.retrieval.nav.nav_knowhere import ( - LazyKnowhereProvider, KnowhereProvider, NamespaceKnowhereProvider, - ReadOnlyChunkStore, SectionRow, UnitRow, - knowhere_database_url, ) from shared.services.retrieval.search.section_filters import is_excluded_section @@ -52,11 +48,6 @@ class NavSnapshot: document_ids: list[str] document_titles: dict[str, str] - def close(self) -> None: - close = getattr(self.provider, "close", None) - if callable(close): - close() - def build_nav_snapshot( *, @@ -95,7 +86,6 @@ async def load_nav_snapshot( namespace: str, exclude_document_ids: list[str] | None = None, exclude_sections: list[dict[str, str]] | None = None, - lazy: bool = False, ) -> NavSnapshot: """Preload namespace current revision into a sync map-nav snapshot.""" excluded_docs = [str(x).strip() for x in (exclude_document_ids or ()) if str(x).strip()] @@ -146,71 +136,6 @@ async def load_nav_snapshot( document_revisions=document_revisions, exclude_sections=excluded_secs, ) - if lazy: - chunk_ids_by_doc, chunk_ref_index, remounted_assets = await _load_chunk_index( - db, - document_revisions=document_revisions, - exclude_sections=excluded_secs, - section_path_by_id=section_path_by_id, - job_id_by_result_id=job_id_by_result_id, - ) - kept_titles = { - did: title for did, title in document_titles.items() if sections_by_doc.get(did) - } - if not kept_titles: - raise ValueError( - f"nav snapshot empty after excludes for " - f"user_id={user_id!r} namespace={namespace!r}" - ) - revisions = {did: result_id for did, result_id in document_revisions if did in kept_titles} - store = ReadOnlyChunkStore( - dsn=knowhere_database_url(), - revisions=revisions, - excluded_sections={ - ( - str(item.get("document_id") or "").strip(), - section_id, - ) - for section_id, section_path in section_path_by_id.items() - for item in excluded_secs - if isinstance(item, dict) - and str(item.get("document_id") or "").strip() - and str(item.get("section_path") or "").strip() == section_path - }, - ) - try: - providers = [ - LazyKnowhereProvider( - doc_id=did, - sections=sections_by_doc.get(did, ()), - chunk_store=store, - known_chunk_ids=chunk_ids_by_doc.get(did, ()), - root_asset_ids=remounted_assets.get(did, {}).get("root", ()), - remounted_assets_by_section=remounted_assets.get(did, {}).get("owners", {}), - ) - for did in kept_titles - ] - provider = NamespaceKnowhereProvider( - providers, - titles=kept_titles, - chunk_owner_by_id={ - chunk_id: meta["document_id"] - for chunk_id, meta in chunk_ref_index.items() - if ":" not in chunk_id - and isinstance(meta, dict) - and meta.get("document_id") - }, - ) - except Exception: - store.close() - raise - return NavSnapshot( - provider=provider, - chunk_ref_index=dict(chunk_ref_index), - document_ids=list(provider.document_ids()), - document_titles={did: kept_titles.get(did, did) for did in provider.document_ids()}, - ) - units_by_doc, chunk_ref_index = await _load_chunks( db, document_revisions=document_revisions, @@ -239,126 +164,6 @@ async def load_nav_snapshot( ) -async def _load_chunk_index( - db: SnapshotSession, - *, - document_revisions: list[tuple[str, str]], - exclude_sections: list[dict[str, str]], - section_path_by_id: dict[str, str], - job_id_by_result_id: dict[str, str], -) -> tuple[dict[str, list[str]], dict[str, dict[str, Any]], dict[str, dict[str, Any]]]: - """Load only IDs/reference metadata; content remains lazy.""" - ids_by_doc: dict[str, list[str]] = {} - ref_index: dict[str, dict[str, Any]] = {} - root_assets_by_doc: dict[str, set[str]] = {} - text_connections_by_doc: dict[str, list[tuple[str, str]]] = {} - for group_start in range(0, len(document_revisions), _REVISION_GROUP_SIZE): - revision_group = document_revisions[group_start : group_start + _REVISION_GROUP_SIZE] - last_key: tuple[str, str, int, str, str] | None = None - while True: - stmt = ( - select( - DocumentChunk.document_id, - DocumentChunk.job_result_id, - DocumentChunk.chunk_id, - DocumentChunk.section_id, - DocumentChunk.chunk_type, - DocumentChunk.file_path, - DocumentChunk.chunk_metadata["connect_to"].label("connect_to"), - DocumentChunk.sort_order, - DocumentChunk.id, - ) - .where(tuple_(DocumentChunk.document_id, DocumentChunk.job_result_id).in_(revision_group)) - .order_by( - DocumentChunk.document_id, - DocumentChunk.job_result_id, - DocumentChunk.sort_order, - DocumentChunk.chunk_id, - DocumentChunk.id, - ) - .limit(_CHUNK_BATCH_SIZE) - ) - if last_key is not None: - stmt = stmt.where( - tuple_( - DocumentChunk.document_id, - DocumentChunk.job_result_id, - DocumentChunk.sort_order, - DocumentChunk.chunk_id, - DocumentChunk.id, - ) - > tuple_(*[literal(value) for value in last_key]) - ) - rows = (await db.execute(stmt)).all() - if not rows: - break - for row in rows: - document_id = str(row[0]) - job_result_id = str(row[1]) - chunk_id = str(row[2] or "").strip() - section_id = str(row[3]) if row[3] else None - section_path = section_path_by_id.get(section_id) if section_id else None - if is_excluded_section( - document_id=document_id, - section_path=section_path, - exclude_sections=exclude_sections, - ) or (section_id and section_id not in section_path_by_id): - continue - if not chunk_id: - continue - chunk_type = str(row[4] or "text") - meta = { - "document_id": document_id, - "section_path": section_path, - "chunk_type": chunk_type, - "file_path": str(row[5] or "") or None, - "job_id": job_id_by_result_id.get(job_result_id), - } - ids_by_doc.setdefault(document_id, []).append(chunk_id) - ref_index[chunk_id] = meta - ref_index[f"{document_id}:{chunk_id}"] = meta - if ( - chunk_type in {"image", "table"} - and section_id - and section_path == "Root" - ): - root_assets_by_doc.setdefault(document_id, set()).add(chunk_id) - connections = row[6] - if isinstance(connections, str) and connections.strip(): - try: - connections = json.loads(connections) - except json.JSONDecodeError: - connections = None - if chunk_type == "text" and isinstance(connections, list): - for connection in connections: - if not isinstance(connection, dict): - continue - target = str(connection.get("target") or "").strip() - if not target: - continue - text_connections_by_doc.setdefault(document_id, []).append( - (section_id or "", target) - ) - last = rows[-1] - last_key = ( - str(last[0]), - str(last[1]), - int(last[7] or 0), - str(last[2] or ""), - str(last[8]), - ) - if len(rows) < _CHUNK_BATCH_SIZE: - break - remounted: dict[str, dict[str, Any]] = {} - for document_id, asset_ids in root_assets_by_doc.items(): - owners: dict[str, list[str]] = {} - for section_id, target in text_connections_by_doc.get(document_id, ()): - if target in asset_ids: - owners.setdefault(section_id, []).append(target) - remounted[document_id] = {"root": sorted(asset_ids), "owners": owners} - return ids_by_doc, ref_index, remounted - - async def _load_sections( db: SnapshotSession, *, From dfc7c79a59715a928eb4d2fcf7c8dce123f8bc37 Mon Sep 17 00:00:00 2001 From: suguanYang Date: Mon, 31 Aug 2026 11:09:19 +0800 Subject: [PATCH 09/10] Revert "Merge pull request #337 from Ontos-AI/feat/wangbinqi/optimize-retrieval-snapshot-load" This reverts commit 9ebee653131fd1704965411e3110b5c68080ed2a, reversing changes made to 79317ab42a6afb67a78a35726597f2148b6ff9d4. --- apps/api/alembic/env.py | 20 +- ...c1d2e3f4_add_chunk_snapshot_order_index.py | 110 ------ apps/api/app/core/exception_handlers.py | 18 +- .../test_exception_handlers_contract.py | 55 --- ...st_retrieval_snapshot_batching_contract.py | 194 ---------- ...retrieval_snapshot_consistency_contract.py | 182 --------- ...etrieval_snapshot_large_corpus_contract.py | 344 ------------------ .../tests/migrations/test_schema_contract.py | 80 ---- .../support/retrieval_snapshot_support.py | 19 - .../shared/core/response/ErrorCode.py | 3 - .../shared/models/database/document.py | 8 - .../shared/services/retrieval/nav_snapshot.py | 266 ++++++-------- .../shared/tests/test_nav_snapshot.py | 12 + 13 files changed, 120 insertions(+), 1191 deletions(-) delete mode 100644 apps/api/alembic/versions/fbf0c1d2e3f4_add_chunk_snapshot_order_index.py delete mode 100644 apps/api/tests/contract/test_exception_handlers_contract.py delete mode 100644 apps/api/tests/contract/test_retrieval_snapshot_batching_contract.py delete mode 100644 apps/api/tests/contract/test_retrieval_snapshot_consistency_contract.py delete mode 100644 apps/api/tests/contract/test_retrieval_snapshot_large_corpus_contract.py delete mode 100644 apps/api/tests/support/retrieval_snapshot_support.py diff --git a/apps/api/alembic/env.py b/apps/api/alembic/env.py index 8e3a5d5a..7b9e5022 100644 --- a/apps/api/alembic/env.py +++ b/apps/api/alembic/env.py @@ -101,10 +101,6 @@ def run_migrations_offline() -> None: include_object=include_object, literal_binds=True, dialect_opts={"paramstyle": "named"}, - # Keep each migration in its own transaction so migrations that - # require an autocommit block (for example CREATE INDEX - # CONCURRENTLY) can safely commit only their own predecessor. - transaction_per_migration=True, # Pass through configured SSL connect args. connect_args=ssl_connect_args, ) @@ -123,25 +119,13 @@ def run_migrations_online() -> None: configured_connection = config.attributes.get("connection") def run_with_connection(connection: Connection) -> None: - caller_owned_transaction = connection.in_transaction() if settings.API_STANDALONE_MODE_ENABLED: ensure_better_auth_user_table(connection) - # The standalone bootstrap query starts SQLAlchemy's implicit - # transaction before Alembic begins tracking migration - # transactions. End only that transaction; never commit a - # transaction supplied by the caller. - if not caller_owned_transaction: - connection.commit() context.configure( connection=connection, target_metadata=target_metadata, include_object=include_object, - # Required for migrations that use autocommit_block(). - transaction_per_migration=True, - # Concurrent DDL cannot run inside a transaction owned by the - # caller. Migrations use regular DDL for that compatibility path. - knowhere_external_transaction=caller_owned_transaction, ) with context.begin_transaction(): @@ -152,7 +136,7 @@ def run_with_connection(connection: Connection) -> None: return if isinstance(configured_connection, Engine): - with configured_connection.connect() as connection: + with configured_connection.begin() as connection: run_with_connection(connection) return @@ -165,7 +149,7 @@ def run_with_connection(connection: Connection) -> None: connect_args=ssl_connect_args, ) - with connectable.connect() as connection: + with connectable.begin() as connection: run_with_connection(connection) diff --git a/apps/api/alembic/versions/fbf0c1d2e3f4_add_chunk_snapshot_order_index.py b/apps/api/alembic/versions/fbf0c1d2e3f4_add_chunk_snapshot_order_index.py deleted file mode 100644 index 56b6f4e3..00000000 --- a/apps/api/alembic/versions/fbf0c1d2e3f4_add_chunk_snapshot_order_index.py +++ /dev/null @@ -1,110 +0,0 @@ -"""Add the chunk snapshot pagination order index. - -Revision ID: fbf0c1d2e3f4 -Revises: f0d85d209e68, fbe1c2d3e4f5 -Create Date: 2026-08-27 00:00:00.000000 -""" - -from __future__ import annotations - -from typing import Sequence, Union - -from alembic import op -from sqlalchemy import text - - -revision: str = "fbf0c1d2e3f4" -down_revision: Union[str, Sequence[str], None] = ( - "f0d85d209e68", - "fbe1c2d3e4f5", -) -branch_labels: Union[str, Sequence[str], None] = None -depends_on: Union[str, Sequence[str], None] = None - -_INDEX_NAME = "idx_document_chunks_revision_snapshot_order" -_INDEX_COLUMNS = "(document_id, job_result_id, sort_order, chunk_id, id)" -_INDEX_METHOD_AND_COLUMNS = f"using btree {_INDEX_COLUMNS}" - - -def _read_index_definition() -> tuple[str, bool] | None: - row = op.get_bind().execute( - text( - """ - SELECT indexdef, index_metadata.indisvalid - FROM pg_indexes AS indexes - JOIN pg_class AS index_class - ON index_class.relname = indexes.indexname - JOIN pg_namespace AS index_namespace - ON index_namespace.oid = index_class.relnamespace - AND index_namespace.nspname = indexes.schemaname - JOIN pg_index AS index_metadata - ON index_metadata.indexrelid = index_class.oid - WHERE indexes.schemaname = current_schema() - AND indexes.tablename = 'document_chunks' - AND indexes.indexname = :index_name - """ - ), - {"index_name": _INDEX_NAME}, - ).first() - if row is None: - return None - return str(row[0]), bool(row[1]) - - -def _index_is_intended() -> bool: - definition = _read_index_definition() - if definition is None or not definition[1]: - return False - normalized_definition = " ".join(definition[0].lower().split()) - return ( - normalized_definition.startswith("create index ") - and _INDEX_METHOD_AND_COLUMNS in normalized_definition - and " where " not in normalized_definition - ) - - -def _drop_index(*, concurrently: bool) -> None: - concurrent_clause = "CONCURRENTLY " if concurrently else "" - op.execute(f"DROP INDEX {concurrent_clause}IF EXISTS {_INDEX_NAME}") - - -def _create_index(*, concurrently: bool) -> None: - concurrent_clause = "CONCURRENTLY " if concurrently else "" - op.execute( - f""" - CREATE INDEX {concurrent_clause}IF NOT EXISTS {_INDEX_NAME} - ON document_chunks {_INDEX_COLUMNS} - """ - ) - - -def upgrade() -> None: - external_transaction = bool( - op.get_context().opts.get("knowhere_external_transaction", False) - ) - if external_transaction: - if _read_index_definition() is not None and not _index_is_intended(): - _drop_index(concurrently=False) - if _read_index_definition() is None: - _create_index(concurrently=False) - return - - # Index creation must not hold a write lock on document_chunks while the - # production corpus is being indexed. CONCURRENTLY cannot run inside the - # transaction Alembic normally opens, so switch to an autocommit block. - with op.get_context().autocommit_block(): - if _read_index_definition() is not None and not _index_is_intended(): - _drop_index(concurrently=True) - if _read_index_definition() is None: - _create_index(concurrently=True) - - -def downgrade() -> None: - external_transaction = bool( - op.get_context().opts.get("knowhere_external_transaction", False) - ) - if external_transaction: - _drop_index(concurrently=False) - else: - with op.get_context().autocommit_block(): - _drop_index(concurrently=True) diff --git a/apps/api/app/core/exception_handlers.py b/apps/api/app/core/exception_handlers.py index 7d5a7268..1df2f085 100644 --- a/apps/api/app/core/exception_handlers.py +++ b/apps/api/app/core/exception_handlers.py @@ -48,7 +48,7 @@ """ import uuid -from collections.abc import Awaitable, Callable, Mapping +from collections.abc import Awaitable, Callable from typing import List, cast from fastapi import FastAPI, HTTPException, Request @@ -88,10 +88,7 @@ def _get_request_id(request: Request) -> str: async def knowhere_exception_handler( - request: Request, - exc: KnowhereException, - *, - response_headers: Mapping[str, str] | None = None, + request: Request, exc: KnowhereException ) -> JSONResponse: """ This handler enforces the separation between: @@ -122,10 +119,7 @@ async def knowhere_exception_handler( exc.logging(request_id=request_id) # Always include request ID header for client-side correlation - # Preserve framework-provided headers such as ``Allow: POST`` for 405 - # responses and ``WWW-Authenticate`` for authentication challenges. - headers = dict(response_headers or {}) - headers["X-Request-ID"] = request_id + headers = {"X-Request-ID": request_id} retry_after = exc.details.get("retry_after") if retry_after: headers["Retry-After"] = str(retry_after) @@ -214,11 +208,7 @@ async def http_exception_handler(request: Request, exc: HTTPException) -> JSONRe ) # Delegate to central handler - return await knowhere_exception_handler( - request, - knowhere_exc, - response_headers=exc.headers, - ) + return await knowhere_exception_handler(request, knowhere_exc) async def validation_exception_handler( diff --git a/apps/api/tests/contract/test_exception_handlers_contract.py b/apps/api/tests/contract/test_exception_handlers_contract.py deleted file mode 100644 index 40eff106..00000000 --- a/apps/api/tests/contract/test_exception_handlers_contract.py +++ /dev/null @@ -1,55 +0,0 @@ -from __future__ import annotations - -import sys -from pathlib import Path -from typing import cast - -from fastapi import FastAPI -from httpx import ASGITransport, AsyncClient - -from tests.support.import_environment import ( - configure_import_environment, - ensure_import_paths, -) - -configure_import_environment() -ensure_import_paths() - - -def _prepare_api_app_imports() -> None: - api_root = str(Path(__file__).resolve().parents[2]) - if api_root in sys.path: - sys.path.remove(api_root) - sys.path.insert(0, api_root) - - -def _create_post_only_app() -> FastAPI: - _prepare_api_app_imports() - - from app.core.exception_handlers import setup_exception_handlers - - app = FastAPI() - - @app.post("/v2/retrieval/query") - async def query_retrieval() -> dict[str, bool]: - return {"ok": True} - - setup_exception_handlers(app) - return app - - -async def test_get_to_post_only_route_returns_method_not_allowed() -> None: - app = _create_post_only_app() - transport = ASGITransport(app=app) - - async with AsyncClient(transport=transport, base_url="http://test") as client: - response = await client.get("/v2/retrieval/query") - - assert response.status_code == 405 - assert response.headers["allow"] == "POST" - - response_json = cast(dict[str, object], response.json()) - error = cast(dict[str, object], response_json["error"]) - assert response_json["success"] is False - assert error["code"] == "METHOD_NOT_ALLOWED" - assert error["message"] == "Method not allowed" diff --git a/apps/api/tests/contract/test_retrieval_snapshot_batching_contract.py b/apps/api/tests/contract/test_retrieval_snapshot_batching_contract.py deleted file mode 100644 index a7239100..00000000 --- a/apps/api/tests/contract/test_retrieval_snapshot_batching_contract.py +++ /dev/null @@ -1,194 +0,0 @@ -from __future__ import annotations - -from collections.abc import Callable -from contextlib import AbstractAsyncContextManager -from typing import cast -from uuid import uuid4 - -from httpx import AsyncClient -from sqlalchemy import Executable, Result -from sqlalchemy.sql.selectable import Select - -from shared.services.retrieval.nav_snapshot import SnapshotSession, load_nav_snapshot -from tests.support.retrieval_snapshot_support import contract_db_session -from tests.support.contract_database import ContractDatabase - - -_USER_ID = "local-dev-user" - - -class _CountingSession: - def __init__(self, session: SnapshotSession) -> None: - self._session = session - self.chunk_query_count = 0 - - async def execute(self, statement: Executable) -> Result[tuple[object, ...]]: - selected_tables = ( - { - getattr(table, "name", "") - for table in statement.get_final_froms() - } - if isinstance(statement, Select) - else set() - ) - if "document_chunks" in selected_tables: - self.chunk_query_count += 1 - result = await self._session.execute(statement) - return cast(Result[tuple[object, ...]], result) - - -async def _seed_many_small_documents(namespace: str, *, document_count: int) -> None: - identifier = uuid4().hex[:8] - await ContractDatabase.execute( - """ - INSERT INTO jobs ( - job_id, user_id, job_type, status, source_type, version, - webhook_enabled, created_at, updated_at, credits_charged, billing_status - ) - SELECT - 'job_batch_' || :identifier || '_' || i, - :user_id, - 'document_ingestion', - 'done', - 'file', - 0, - false, - NOW(), - NOW(), - 0, - 'skipped' - FROM generate_series(1, :document_count) AS values(i) - """, - { - "identifier": identifier, - "user_id": _USER_ID, - "document_count": document_count, - }, - ) - await ContractDatabase.execute( - """ - INSERT INTO documents ( - document_id, user_id, namespace, status, current_job_result_id, - source_file_name, parse_track, created_at, updated_at - ) - SELECT - 'doc_batch_' || :identifier || '_' || i, - :user_id, - :namespace, - 'active', - NULL, - 'batch-' || i || '.pdf', - 'chunk', - NOW(), - NOW() - FROM generate_series(1, :document_count) AS values(i) - """, - { - "identifier": identifier, - "user_id": _USER_ID, - "namespace": namespace, - "document_count": document_count, - }, - ) - await ContractDatabase.execute( - """ - INSERT INTO job_results ( - id, job_id, document_id, delivery_mode, created_at, updated_at - ) - SELECT - 'result_batch_' || :identifier || '_' || i, - 'job_batch_' || :identifier || '_' || i, - 'doc_batch_' || :identifier || '_' || i, - 'inline', - NOW(), - NOW() - FROM generate_series(1, :document_count) AS values(i) - """, - {"identifier": identifier, "document_count": document_count}, - ) - await ContractDatabase.execute( - """ - UPDATE documents - SET current_job_result_id = 'result_batch_' || :identifier || '_' || i - FROM generate_series(1, :document_count) AS values(i) - WHERE document_id = 'doc_batch_' || :identifier || '_' || i - """, - {"identifier": identifier, "document_count": document_count}, - ) - await ContractDatabase.execute( - """ - INSERT INTO document_sections ( - section_id, user_id, namespace, document_id, job_result_id, - section_path, section_title, section_level, sort_order, created_at - ) - SELECT - 'section_batch_' || :identifier || '_' || i, - :user_id, - :namespace, - 'doc_batch_' || :identifier || '_' || i, - 'result_batch_' || :identifier || '_' || i, - 'batch-' || i || '.pdf/section', - 'section', - 1, - 1, - NOW() - FROM generate_series(1, :document_count) AS values(i) - """, - { - "identifier": identifier, - "user_id": _USER_ID, - "namespace": namespace, - "document_count": document_count, - }, - ) - await ContractDatabase.execute( - """ - INSERT INTO document_chunks ( - id, chunk_id, user_id, namespace, document_id, job_result_id, - section_id, chunk_type, content, source_chunk_path, - chunk_metadata, sort_order, created_at - ) - SELECT - 'dchunk_batch_' || :identifier || '_' || i, - 'chunk_batch_' || :identifier || '_' || i, - :user_id, - :namespace, - 'doc_batch_' || :identifier || '_' || i, - 'result_batch_' || :identifier || '_' || i, - 'section_batch_' || :identifier || '_' || i, - 'text', - 'content-' || i, - 'batch-' || i || '.pdf/section/chunk', - '{}'::json, - 1, - NOW() - FROM generate_series(1, :document_count) AS values(i) - """, - { - "identifier": identifier, - "user_id": _USER_ID, - "namespace": namespace, - "document_count": document_count, - }, - ) - - -async def test_snapshot_batches_chunks_across_many_small_documents( - developer_api_client_factory: Callable[ - [], AbstractAsyncContextManager[AsyncClient] - ], -) -> None: - document_count = 25 - namespace = f"batch-documents-{uuid4().hex[:8]}" - async with developer_api_client_factory(): - await _seed_many_small_documents(namespace, document_count=document_count) - async with contract_db_session() as db: - counting_db = _CountingSession(db) - snapshot = await load_nav_snapshot( - counting_db, - user_id=_USER_ID, - namespace=namespace, - ) - - assert len(snapshot.document_ids) == document_count - assert counting_db.chunk_query_count == 1 diff --git a/apps/api/tests/contract/test_retrieval_snapshot_consistency_contract.py b/apps/api/tests/contract/test_retrieval_snapshot_consistency_contract.py deleted file mode 100644 index 64aaf4c0..00000000 --- a/apps/api/tests/contract/test_retrieval_snapshot_consistency_contract.py +++ /dev/null @@ -1,182 +0,0 @@ -from __future__ import annotations - -from collections.abc import Awaitable, Callable -from contextlib import AbstractAsyncContextManager -from typing import cast -from uuid import uuid4 - -from httpx import AsyncClient -from sqlalchemy import Executable, Result - -from shared.services.retrieval.nav_snapshot import SnapshotSession, load_nav_snapshot -from tests.support.retrieval_snapshot_support import contract_db_session -from tests.support.contract_database import ContractDatabase - - -_USER_ID = "local-dev-user" - - -class _PublishingSession: - def __init__( - self, - session: SnapshotSession, - publish_revision: Callable[[], Awaitable[None]], - ) -> None: - self._session = session - self._publish_revision = publish_revision - self._has_published = False - - async def execute(self, statement: Executable) -> Result[tuple[object, ...]]: - result = await self._session.execute(statement) - if not self._has_published: - self._has_published = True - await self._publish_revision() - return cast(Result[tuple[object, ...]], result) - - -async def _seed_republished_document(namespace: str) -> tuple[str, str]: - identifier = uuid4().hex[:8] - document_id = f"doc_race_{identifier}" - old_result_id = f"result_race_old_{identifier}" - new_result_id = f"result_race_new_{identifier}" - await ContractDatabase.execute( - """ - INSERT INTO jobs ( - job_id, user_id, job_type, status, source_type, version, - webhook_enabled, created_at, updated_at, credits_charged, billing_status - ) VALUES - (:old_job_id, :user_id, 'document_ingestion', 'done', 'file', 0, - false, NOW(), NOW(), 0, 'skipped'), - (:new_job_id, :user_id, 'document_ingestion', 'done', 'file', 0, - false, NOW(), NOW(), 0, 'skipped') - """, - { - "old_job_id": f"job_race_old_{identifier}", - "new_job_id": f"job_race_new_{identifier}", - "user_id": _USER_ID, - }, - ) - await ContractDatabase.execute( - """ - INSERT INTO documents ( - document_id, user_id, namespace, status, current_job_result_id, - source_file_name, parse_track, created_at, updated_at - ) VALUES ( - :document_id, :user_id, :namespace, 'active', NULL, - 'republished.pdf', 'chunk', NOW(), NOW() - ) - """, - {"document_id": document_id, "user_id": _USER_ID, "namespace": namespace}, - ) - await ContractDatabase.execute( - """ - INSERT INTO job_results ( - id, job_id, document_id, delivery_mode, created_at, updated_at - ) VALUES - (:old_result_id, :old_job_id, :document_id, 'inline', NOW(), NOW()), - (:new_result_id, :new_job_id, :document_id, 'inline', NOW(), NOW()) - """, - { - "old_result_id": old_result_id, - "new_result_id": new_result_id, - "old_job_id": f"job_race_old_{identifier}", - "new_job_id": f"job_race_new_{identifier}", - "document_id": document_id, - }, - ) - await ContractDatabase.execute( - """ - INSERT INTO document_sections ( - section_id, user_id, namespace, document_id, job_result_id, - section_path, section_title, section_level, sort_order, created_at - ) VALUES - (:old_section_id, :user_id, :namespace, :document_id, :old_result_id, - 'republished.pdf/old', 'old', 1, 1, NOW()), - (:new_section_id, :user_id, :namespace, :document_id, :new_result_id, - 'republished.pdf/new', 'new', 1, 1, NOW()) - """, - { - "old_section_id": f"section_race_old_{identifier}", - "new_section_id": f"section_race_new_{identifier}", - "user_id": _USER_ID, - "namespace": namespace, - "document_id": document_id, - "old_result_id": old_result_id, - "new_result_id": new_result_id, - }, - ) - await ContractDatabase.execute( - """ - INSERT INTO document_chunks ( - id, chunk_id, user_id, namespace, document_id, job_result_id, - section_id, chunk_type, content, source_chunk_path, - chunk_metadata, sort_order, created_at - ) VALUES - (:old_id, :old_chunk_id, :user_id, :namespace, :document_id, - :old_result_id, :old_section_id, 'text', 'old content', - 'republished.pdf/old/chunk', '{}'::json, 1, NOW()), - (:new_id, :new_chunk_id, :user_id, :namespace, :document_id, - :new_result_id, :new_section_id, 'text', 'new content', - 'republished.pdf/new/chunk', '{}'::json, 1, NOW()) - """, - { - "old_id": f"dchunk_race_old_{identifier}", - "new_id": f"dchunk_race_new_{identifier}", - "old_chunk_id": f"chunk_race_old_{identifier}", - "new_chunk_id": f"chunk_race_new_{identifier}", - "user_id": _USER_ID, - "namespace": namespace, - "document_id": document_id, - "old_result_id": old_result_id, - "new_result_id": new_result_id, - "old_section_id": f"section_race_old_{identifier}", - "new_section_id": f"section_race_new_{identifier}", - }, - ) - await ContractDatabase.execute( - """ - UPDATE documents - SET current_job_result_id = :old_result_id - WHERE document_id = :document_id - """, - {"old_result_id": old_result_id, "document_id": document_id}, - ) - return document_id, new_result_id - - -async def test_snapshot_keeps_sections_and_chunks_on_the_captured_revision( - developer_api_client_factory: Callable[ - [], AbstractAsyncContextManager[AsyncClient] - ], -) -> None: - namespace = f"revision-race-{uuid4().hex[:8]}" - async with developer_api_client_factory(): - document_id, new_result_id = await _seed_republished_document(namespace) - - async def publish_new_revision() -> None: - await ContractDatabase.execute( - """ - UPDATE documents - SET current_job_result_id = :new_result_id - WHERE document_id = :document_id - """, - {"new_result_id": new_result_id, "document_id": document_id}, - ) - - async with contract_db_session() as db: - snapshot = await load_nav_snapshot( - _PublishingSession(db, publish_new_revision), - user_id=_USER_ID, - namespace=namespace, - ) - - section_ids = list(snapshot.provider.children(document_id)) - chunks = [ - chunk - for section_id in section_ids - for chunk in snapshot.provider.self_units(section_id) - ] - assert [chunk.content for chunk in chunks] == ["old content"] - assert snapshot.chunk_ref_index[chunks[0].chunk_id]["section_path"] == ( - "republished.pdf/old" - ) diff --git a/apps/api/tests/contract/test_retrieval_snapshot_large_corpus_contract.py b/apps/api/tests/contract/test_retrieval_snapshot_large_corpus_contract.py deleted file mode 100644 index a414e44b..00000000 --- a/apps/api/tests/contract/test_retrieval_snapshot_large_corpus_contract.py +++ /dev/null @@ -1,344 +0,0 @@ -from __future__ import annotations - -from collections.abc import Callable -from contextlib import AbstractAsyncContextManager -from typing import cast -from uuid import uuid4 - -from httpx import AsyncClient - -from shared.models.database.document import Document, DocumentChunk, DocumentSection -from shared.models.database.job_result import JobResult -from shared.services.retrieval.nav_snapshot import ( - SnapshotSession, - _CHUNK_BATCH_SIZE, - _REVISION_GROUP_SIZE, - load_nav_snapshot, -) -from sqlalchemy import Executable, Result, select -from sqlalchemy.engine import Row -from sqlalchemy.sql.selectable import Select -from tests.support.retrieval_snapshot_support import contract_db_session -from tests.support.contract_database import ContractDatabase - - -_USER_ID = "local-dev-user" -_DOCUMENT_COUNT = 100 -_CHUNKS_PER_DOCUMENT = 600 -_SECTIONS_PER_DOCUMENT = 8 -_TOTAL_CHUNKS = _DOCUMENT_COUNT * _CHUNKS_PER_DOCUMENT -LegacySnapshotRow = Row[ - tuple[ - str, - str, - str | None, - str, - str, - int, - str, - str | None, - dict[str, object], - str | None, - str | None, - ] -] - - -class _CountingSession: - def __init__(self, session: SnapshotSession) -> None: - self._session = session - self.chunk_query_count = 0 - - async def execute(self, statement: Executable) -> Result[tuple[object, ...]]: - if isinstance(statement, Select): - selected_tables = { - getattr(table, "name", "") - for table in statement.get_final_froms() - } - else: - selected_tables = set() - if "document_chunks" in selected_tables: - self.chunk_query_count += 1 - result = await self._session.execute(statement) - return cast(Result[tuple[object, ...]], result) - - -async def _seed_large_retrieval_corpus(namespace: str) -> None: - await ContractDatabase.execute( - """ - INSERT INTO jobs ( - job_id, user_id, job_type, status, source_type, version, - webhook_enabled, created_at, updated_at, credits_charged, billing_status - ) - SELECT - 'job_lg_' || i, - :user_id, - 'document_ingestion', - 'done', - 'file', - 0, - false, - NOW(), - NOW(), - 0, - 'skipped' - FROM generate_series(1, :document_count) AS values(i) - """, - {"user_id": _USER_ID, "document_count": _DOCUMENT_COUNT}, - ) - await ContractDatabase.execute( - """ - INSERT INTO documents ( - document_id, user_id, namespace, status, current_job_result_id, - source_file_name, parse_track, created_at, updated_at - ) - SELECT - 'doc_lg_' || i, - :user_id, - :namespace, - 'active', - NULL, - 'large-corpus-' || i || '.pdf', - 'chunk', - NOW(), - NOW() - FROM generate_series(1, :document_count) AS values(i) - """, - { - "user_id": _USER_ID, - "namespace": namespace, - "document_count": _DOCUMENT_COUNT, - }, - ) - await ContractDatabase.execute( - """ - INSERT INTO job_results ( - id, job_id, document_id, delivery_mode, created_at, updated_at - ) - SELECT - 'result_lg_' || i, - 'job_lg_' || i, - 'doc_lg_' || i, - 'inline', - NOW(), - NOW() - FROM generate_series(1, :document_count) AS values(i) - """, - {"document_count": _DOCUMENT_COUNT}, - ) - await ContractDatabase.execute( - """ - UPDATE documents - SET current_job_result_id = 'result_lg_' || i - FROM generate_series(1, :document_count) AS values(i) - WHERE documents.document_id = 'doc_lg_' || i - """, - {"document_count": _DOCUMENT_COUNT}, - ) - await ContractDatabase.execute( - """ - INSERT INTO document_sections ( - section_id, user_id, namespace, document_id, job_result_id, - section_path, section_title, section_level, sort_order, created_at - ) - SELECT - 'section_lg_' || i || '_' || section_number, - :user_id, - :namespace, - 'doc_lg_' || i, - 'result_lg_' || i, - 'large-corpus-' || i || '/section/' || section_number, - 'section-' || section_number, - 1, - section_number, - NOW() - FROM generate_series(1, :document_count) AS values(i) - CROSS JOIN generate_series(1, :sections_per_document) AS sections(section_number) - """, - { - "user_id": _USER_ID, - "namespace": namespace, - "document_count": _DOCUMENT_COUNT, - "sections_per_document": _SECTIONS_PER_DOCUMENT, - }, - ) - await ContractDatabase.execute( - """ - INSERT INTO document_chunks ( - id, chunk_id, user_id, namespace, document_id, job_result_id, - section_id, chunk_type, content, content_lexical_text, - path_lexical_text, content_search_text, path_search_text, - term_search_text, source_chunk_path, chunk_metadata, sort_order, - created_at - ) - SELECT - 'dchunk_lg_' || document_number || '_' || chunk_number, - 'chunk_lg_' || document_number || '_' || chunk_number, - :user_id, - :namespace, - 'doc_lg_' || document_number, - 'result_lg_' || document_number, - 'section_lg_' || document_number || '_' || section_number, - 'text', - repeat(md5(document_number::text || ':' || chunk_number::text), 64), - repeat(md5(document_number::text || ':' || chunk_number::text), 64), - 'large-corpus-' || document_number || '/section/' || section_number || '/' || chunk_number, - repeat(md5(document_number::text || ':' || chunk_number::text), 64), - 'large-corpus-' || document_number || '/section/' || section_number || '/' || chunk_number, - repeat(md5(document_number::text || ':' || chunk_number::text), 64), - 'large-corpus-' || document_number || '/section/' || section_number || '/' || chunk_number, - json_build_object( - 'tokens', ARRAY['benchmark', 'retrieval', 'document', document_number::text], - 'keywords', ARRAY['benchmark', 'production-shaped'], - 'summary', repeat('payload ', 32) - )::json, - chunk_number, - NOW() - FROM generate_series(1, :document_count) AS documents(document_number) - CROSS JOIN generate_series(1, :chunks_per_document) AS chunks(chunk_number) - CROSS JOIN LATERAL ( - SELECT ((chunk_number - 1) % :sections_per_document) + 1 AS section_number - ) AS section_values - """, - { - "user_id": _USER_ID, - "namespace": namespace, - "document_count": _DOCUMENT_COUNT, - "chunks_per_document": _CHUNKS_PER_DOCUMENT, - "sections_per_document": _SECTIONS_PER_DOCUMENT, - }, - ) - - -async def _load_legacy_rows( - namespace: str, -) -> list[LegacySnapshotRow]: - stmt = ( - select( - Document.document_id, - DocumentChunk.chunk_id, - DocumentChunk.section_id, - DocumentChunk.chunk_type, - DocumentChunk.content, - DocumentChunk.sort_order, - DocumentChunk.source_chunk_path, - DocumentChunk.file_path, - DocumentChunk.chunk_metadata, - DocumentSection.section_path, - JobResult.job_id, - ) - .join( - DocumentChunk, - (DocumentChunk.document_id == Document.document_id) - & (DocumentChunk.job_result_id == Document.current_job_result_id), - ) - .outerjoin( - DocumentSection, DocumentSection.section_id == DocumentChunk.section_id - ) - .outerjoin(JobResult, JobResult.id == DocumentChunk.job_result_id) - .where(Document.user_id == _USER_ID) - .where(Document.namespace == namespace) - .where(Document.status == "active") - .order_by( - Document.document_id, DocumentChunk.sort_order, DocumentChunk.chunk_id - ) - ) - async with contract_db_session() as db: - rows: list[LegacySnapshotRow] = list((await db.execute(stmt)).all()) - return rows - - -async def test_large_snapshot_keeps_all_retrieval_inputs_after_bounded_sql_load( - developer_api_client_factory: Callable[ - [], AbstractAsyncContextManager[AsyncClient] - ], -) -> None: - namespace = f"large-corpus-{uuid4().hex[:8]}" - async with developer_api_client_factory(): - await _seed_large_retrieval_corpus(namespace) - legacy_rows = await _load_legacy_rows(namespace) - async with contract_db_session() as db: - counting_db = _CountingSession(db) - snapshot = await load_nav_snapshot( - counting_db, - user_id=_USER_ID, - namespace=namespace, - ) - expected_chunk_query_count = sum( - ( - min(_REVISION_GROUP_SIZE, _DOCUMENT_COUNT - group_start) - * _CHUNKS_PER_DOCUMENT - + _CHUNK_BATCH_SIZE - - 1 - ) - // _CHUNK_BATCH_SIZE - for group_start in range(0, _DOCUMENT_COUNT, _REVISION_GROUP_SIZE) - ) - assert counting_db.chunk_query_count == expected_chunk_query_count - - async with contract_db_session() as db: - bounded_snapshot = await load_nav_snapshot( - db, - user_id=_USER_ID, - namespace=namespace, - ) - - assert len(legacy_rows) == _TOTAL_CHUNKS - assert len(snapshot.document_ids) == _DOCUMENT_COUNT - assert len(bounded_snapshot.document_ids) == _DOCUMENT_COUNT - - optimized_rows = [ - ( - document_id, - chunk.chunk_id, - chunk.section_id, - chunk.chunk_type, - chunk.content, - chunk.sort_order, - chunk.source_chunk_path, - chunk.file_path, - chunk.metadata, - snapshot.chunk_ref_index[f"{document_id}:{chunk.chunk_id}"][ - "section_path" - ], - snapshot.chunk_ref_index[f"{document_id}:{chunk.chunk_id}"]["job_id"], - ) - for document_id in snapshot.document_ids - for section_id in snapshot.provider.children(document_id) - for chunk in snapshot.provider.self_units(section_id) - ] - optimized_rows.sort(key=lambda row: (row[0], row[5], row[1], row[2] or "")) - legacy_rows_projected = [ - ( - str(row[0]), - str(row[1]), - str(row[2]) if row[2] else None, - str(row[3]), - str(row[4]), - int(row[5]), - str(row[6]), - str(row[7] or ""), - row[8] if isinstance(row[8], dict) else {}, - str(row[9] or ""), - str(row[10]) if row[10] else None, - ) - for row in legacy_rows - ] - def row_order_key(row: tuple[object, ...]) -> tuple[str, ...]: - return ( - str(row[0]), - str(row[5]), - str(row[1]), - str(row[2] or ""), - str(row[3]), - str(row[4]), - str(row[6]), - str(row[7] or ""), - str(row[8]), - str(row[9] or ""), - str(row[10] or ""), - ) - legacy_rows_projected.sort(key=row_order_key) - optimized_rows.sort(key=row_order_key) - assert len(optimized_rows) == _TOTAL_CHUNKS - assert optimized_rows == legacy_rows_projected diff --git a/apps/api/tests/migrations/test_schema_contract.py b/apps/api/tests/migrations/test_schema_contract.py index 24047dcc..830a2cbd 100644 --- a/apps/api/tests/migrations/test_schema_contract.py +++ b/apps/api/tests/migrations/test_schema_contract.py @@ -31,24 +31,11 @@ def _build_alembic_command_config(*, engine: Engine) -> Config: def _upgrade_to_heads(*, engine: Engine) -> None: config = _build_alembic_command_config(engine=engine) - # Let Alembic create its own connection. This is required for migrations - # that use PostgreSQL autocommit (for example CREATE INDEX CONCURRENTLY). - command.upgrade(config, "heads") - - -def _upgrade_to_heads_with_external_connection(*, engine: Engine) -> None: - config = _build_alembic_command_config(engine=engine) with engine.begin() as connection: config.attributes["connection"] = connection command.upgrade(config, "heads") -def _upgrade_to_snapshot_parents(*, engine: Engine) -> None: - config = _build_alembic_command_config(engine=engine) - command.upgrade(config, "f0d85d209e68") - command.upgrade(config, "fbe1c2d3e4f5") - - def _insert_job( connection: Connection, *, @@ -212,73 +199,6 @@ def test_should_seed_v2_job_polling_system_limit( assert result["description"] == "Job queries - prevent polling" -def test_should_index_document_chunks_in_snapshot_pagination_order( - migrated_head_engine: Engine, -) -> None: - with migrated_head_engine.begin() as connection: - index_definition = connection.execute( - text( - """ - SELECT indexdef - FROM pg_indexes - WHERE schemaname = current_schema() - AND tablename = 'document_chunks' - AND indexname = 'idx_document_chunks_revision_snapshot_order' - """ - ) - ).scalar_one() - - assert "(document_id, job_result_id, sort_order, chunk_id, id)" in str( - index_definition - ) - - -def test_should_upgrade_with_a_caller_owned_connection( - alembic_engine: Engine, -) -> None: - _upgrade_to_heads_with_external_connection(engine=alembic_engine) - - -def test_standalone_upgrade_should_preserve_a_caller_owned_connection( - standalone_alembic_engine: Engine, -) -> None: - _upgrade_to_heads_with_external_connection(engine=standalone_alembic_engine) - - -def test_should_replace_a_same_named_index_with_the_wrong_definition( - alembic_engine: Engine, -) -> None: - _upgrade_to_snapshot_parents(engine=alembic_engine) - with alembic_engine.begin() as connection: - connection.execute( - text( - """ - CREATE INDEX idx_document_chunks_revision_snapshot_order - ON document_chunks (document_id) - """ - ) - ) - - _upgrade_to_heads(engine=alembic_engine) - - with alembic_engine.begin() as connection: - index_definition = connection.execute( - text( - """ - SELECT indexdef - FROM pg_indexes - WHERE schemaname = current_schema() - AND tablename = 'document_chunks' - AND indexname = 'idx_document_chunks_revision_snapshot_order' - """ - ) - ).scalar_one() - - assert "(document_id, job_result_id, sort_order, chunk_id, id)" in str( - index_definition - ) - - def test_api_standalone_mode_should_create_auth_user_table_before_migrations( standalone_alembic_engine: Engine, ) -> None: diff --git a/apps/api/tests/support/retrieval_snapshot_support.py b/apps/api/tests/support/retrieval_snapshot_support.py deleted file mode 100644 index b06b55f1..00000000 --- a/apps/api/tests/support/retrieval_snapshot_support.py +++ /dev/null @@ -1,19 +0,0 @@ -from __future__ import annotations - -from collections.abc import AsyncGenerator -from contextlib import asynccontextmanager - -from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker, create_async_engine - -from shared.testing.contract_runtime import get_contract_database_url - - -@asynccontextmanager -async def contract_db_session() -> AsyncGenerator[AsyncSession, None]: - engine = create_async_engine(get_contract_database_url(), future=True) - session_factory = async_sessionmaker(engine, expire_on_commit=False) - try: - async with session_factory() as session: - yield session - finally: - await engine.dispose() diff --git a/packages/shared-python/shared/core/response/ErrorCode.py b/packages/shared-python/shared/core/response/ErrorCode.py index cf5dd2f8..4c982093 100644 --- a/packages/shared-python/shared/core/response/ErrorCode.py +++ b/packages/shared-python/shared/core/response/ErrorCode.py @@ -55,7 +55,6 @@ class ErrorCode(str, Enum): ) PERMISSION_DENIED = "PERMISSION_DENIED" # 403 - Caller lacks permission NOT_FOUND = "NOT_FOUND" # 404 - Resource does not exist - METHOD_NOT_ALLOWED = "METHOD_NOT_ALLOWED" # 405 - HTTP method is not supported ABORTED = "ABORTED" # 409 - Concurrency conflict ALREADY_EXISTS = "ALREADY_EXISTS" # 409 - Resource already exists RESOURCE_EXHAUSTED = ( @@ -90,7 +89,6 @@ class ErrorCodeMapper: ErrorCode.PAYMENT_REQUIRED: 402, ErrorCode.PERMISSION_DENIED: 403, ErrorCode.NOT_FOUND: 404, - ErrorCode.METHOD_NOT_ALLOWED: 405, ErrorCode.ABORTED: 409, ErrorCode.ALREADY_EXISTS: 409, ErrorCode.RESOURCE_EXHAUSTED: 429, @@ -112,7 +110,6 @@ class ErrorCodeMapper: 402: ErrorCode.PAYMENT_REQUIRED, 403: ErrorCode.PERMISSION_DENIED, 404: ErrorCode.NOT_FOUND, - 405: ErrorCode.METHOD_NOT_ALLOWED, 409: ErrorCode.ALREADY_EXISTS, 422: ErrorCode.INVALID_ARGUMENT, # Pydantic validation 429: ErrorCode.RESOURCE_EXHAUSTED, diff --git a/packages/shared-python/shared/models/database/document.py b/packages/shared-python/shared/models/database/document.py index f9681e17..c9af7140 100644 --- a/packages/shared-python/shared/models/database/document.py +++ b/packages/shared-python/shared/models/database/document.py @@ -202,14 +202,6 @@ class DocumentChunk(Base): Index("idx_document_chunks_scope", "user_id", "namespace"), Index("idx_document_chunks_chunk_id", "chunk_id"), Index("idx_document_chunks_doc_revision", "document_id", "job_result_id"), - Index( - "idx_document_chunks_revision_snapshot_order", - "document_id", - "job_result_id", - "sort_order", - "chunk_id", - "id", - ), Index("idx_document_chunks_section", "section_id"), Index( "idx_chunk_content_search_tsv", diff --git a/packages/shared-python/shared/services/retrieval/nav_snapshot.py b/packages/shared-python/shared/services/retrieval/nav_snapshot.py index 7dc60c6f..964d698b 100644 --- a/packages/shared-python/shared/services/retrieval/nav_snapshot.py +++ b/packages/shared-python/shared/services/retrieval/nav_snapshot.py @@ -8,10 +8,10 @@ from __future__ import annotations from dataclasses import dataclass -from typing import Any, Protocol +from typing import Any, Optional -from sqlalchemy import Executable, literal, select, tuple_ -from sqlalchemy.engine import Result +from sqlalchemy import select +from sqlalchemy.ext.asyncio import AsyncSession from shared.models.database.document import Document, DocumentChunk, DocumentSection from shared.models.database.job_result import JobResult @@ -24,21 +24,6 @@ from shared.services.retrieval.search.section_filters import is_excluded_section -# Keep each payload query bounded under the API's 30-second statement timeout. -# Ten-thousand-row keyset pages avoid OFFSET scans while keeping each payload -# statement bounded. The contract benchmark verifies this page size against -# the full 2 KiB content and metadata payload. -_CHUNK_BATCH_SIZE = 10_000 -_REVISION_GROUP_SIZE = 32 - - -class SnapshotSession(Protocol): - """Minimal database interface required by the snapshot loader.""" - - async def execute(self, statement: Executable) -> Result[tuple[object, ...]]: - raise NotImplementedError - - @dataclass(frozen=True) class NavSnapshot: """In-memory corpus for one map-nav episode.""" @@ -80,7 +65,7 @@ def build_nav_snapshot( async def load_nav_snapshot( - db: SnapshotSession, + db: AsyncSession, *, user_id: str, namespace: str, @@ -109,39 +94,25 @@ async def load_nav_snapshot( ) document_titles: dict[str, str] = {} - current_job_result_ids: set[str] = set() - document_revisions: list[tuple[str, str]] = [] - for document_id, source_file_name, current_job_result_id in doc_rows: + for document_id, source_file_name, _job in doc_rows: did = str(document_id) title = str(source_file_name or "").strip() or did document_titles[did] = title - if current_job_result_id: - job_result_id = str(current_job_result_id) - current_job_result_ids.add(job_result_id) - document_revisions.append((did, job_result_id)) - - job_result_rows = await db.execute( - select(JobResult.id, JobResult.job_id).where( - JobResult.id.in_(list(current_job_result_ids)) - ) - ) - job_id_by_result_id = { - str(job_result_id): str(job_id) - for job_result_id, job_id in job_result_rows.all() - if job_result_id and job_id - } sections_by_doc, section_path_by_id = await _load_sections( db, - document_revisions=document_revisions, + user_id=user_id, + namespace=namespace, + exclude_document_ids=excluded_docs, exclude_sections=excluded_secs, ) units_by_doc, chunk_ref_index = await _load_chunks( db, - document_revisions=document_revisions, + user_id=user_id, + namespace=namespace, + exclude_document_ids=excluded_docs, exclude_sections=excluded_secs, section_path_by_id=section_path_by_id, - job_id_by_result_id=job_id_by_result_id, ) # Keep only documents that still have sections after exclude filters. @@ -165,15 +136,16 @@ async def load_nav_snapshot( async def _load_sections( - db: SnapshotSession, + db: AsyncSession, *, - document_revisions: list[tuple[str, str]], + user_id: str, + namespace: str, + exclude_document_ids: list[str], exclude_sections: list[dict[str, str]], ) -> tuple[dict[str, list[SectionRow]], dict[str, str]]: - # Captured pairs replace DocumentSection.job_result_id == Document.current_job_result_id. stmt = ( select( - DocumentSection.document_id, + Document.document_id, DocumentSection.section_id, DocumentSection.parent_section_id, DocumentSection.section_path, @@ -182,18 +154,18 @@ async def _load_sections( DocumentSection.summary, DocumentSection.sort_order, ) - .where( - tuple_( - DocumentSection.document_id, - DocumentSection.job_result_id, - ).in_(document_revisions) - ) - .order_by( - DocumentSection.document_id, - DocumentSection.sort_order, - DocumentSection.section_id, + .join( + DocumentSection, + (DocumentSection.document_id == Document.document_id) + & (DocumentSection.job_result_id == Document.current_job_result_id), ) + .where(Document.user_id == user_id) + .where(Document.namespace == namespace) + .where(Document.status == "active") + .order_by(Document.document_id, DocumentSection.sort_order, DocumentSection.section_id) ) + if exclude_document_ids: + stmt = stmt.where(Document.document_id.notin_(list(exclude_document_ids))) by_doc: dict[str, list[SectionRow]] = {} path_by_id: dict[str, str] = {} @@ -222,123 +194,89 @@ async def _load_sections( async def _load_chunks( - db: SnapshotSession, + db: AsyncSession, *, - document_revisions: list[tuple[str, str]], + user_id: str, + namespace: str, + exclude_document_ids: list[str], exclude_sections: list[dict[str, str]], section_path_by_id: dict[str, str], - job_id_by_result_id: dict[str, str], ) -> tuple[dict[str, list[UnitRow]], dict[str, dict[str, Any]]]: - # Captured pairs replace DocumentChunk.document_id == document_id and - # DocumentChunk.job_result_id == job_result_id. + stmt = ( + select( + Document.document_id, + DocumentChunk.chunk_id, + DocumentChunk.section_id, + DocumentChunk.chunk_type, + DocumentChunk.content, + DocumentChunk.sort_order, + DocumentChunk.source_chunk_path, + DocumentChunk.file_path, + DocumentChunk.chunk_metadata, + DocumentSection.section_path, + JobResult.job_id, + ) + .join( + DocumentChunk, + (DocumentChunk.document_id == Document.document_id) + & (DocumentChunk.job_result_id == Document.current_job_result_id), + ) + .outerjoin( + DocumentSection, + DocumentSection.section_id == DocumentChunk.section_id, + ) + .outerjoin(JobResult, JobResult.id == DocumentChunk.job_result_id) + .where(Document.user_id == user_id) + .where(Document.namespace == namespace) + .where(Document.status == "active") + .order_by(Document.document_id, DocumentChunk.sort_order, DocumentChunk.chunk_id) + ) + if exclude_document_ids: + stmt = stmt.where(Document.document_id.notin_(list(exclude_document_ids))) + by_doc: dict[str, list[UnitRow]] = {} ref_index: dict[str, dict[str, Any]] = {} - for group_start in range(0, len(document_revisions), _REVISION_GROUP_SIZE): - revision_group = document_revisions[ - group_start : group_start + _REVISION_GROUP_SIZE - ] - last_key: tuple[str, str, int, str, str] | None = None - while True: - stmt = ( - select( - DocumentChunk.document_id, - DocumentChunk.job_result_id, - DocumentChunk.chunk_id, - DocumentChunk.section_id, - DocumentChunk.chunk_type, - DocumentChunk.content, - DocumentChunk.sort_order, - DocumentChunk.source_chunk_path, - DocumentChunk.file_path, - DocumentChunk.chunk_metadata, - DocumentChunk.id, - ) - .where( - tuple_( - DocumentChunk.document_id, - DocumentChunk.job_result_id, - ).in_(revision_group) - ) - .order_by( - DocumentChunk.document_id, - DocumentChunk.job_result_id, - DocumentChunk.sort_order, - DocumentChunk.chunk_id, - DocumentChunk.id, - ) - .limit(_CHUNK_BATCH_SIZE) - ) - if last_key is not None: - stmt = stmt.where( - tuple_( - DocumentChunk.document_id, - DocumentChunk.job_result_id, - DocumentChunk.sort_order, - DocumentChunk.chunk_id, - DocumentChunk.id, - ) - > tuple_( - literal(last_key[0]), - literal(last_key[1]), - literal(last_key[2]), - literal(last_key[3]), - literal(last_key[4]), - ) - ) - - rows = (await db.execute(stmt)).all() - if not rows: - break - for row in rows: - document_id = str(row[0]) - job_result_id = str(row[1]) - chunk_id = str(row[2] or "").strip() - section_id = str(row[3]) if row[3] else None - section_path: str | None = ( - section_path_by_id.get(section_id) if section_id else None - ) - if is_excluded_section( - document_id=document_id, - section_path=section_path, - exclude_sections=exclude_sections, - ): - continue - if section_id and section_id not in section_path_by_id: - continue + for row in (await db.execute(stmt)).all(): + document_id = str(row[0]) + chunk_id = str(row[1] or "").strip() + section_id = str(row[2]) if row[2] else None + # Prefer joined section_path; fall back to kept section map. + section_path: Optional[str] = str(row[9]) if row[9] is not None else None + if section_path is None and section_id: + section_path = section_path_by_id.get(section_id) + if is_excluded_section( + document_id=document_id, + section_path=section_path, + exclude_sections=exclude_sections, + ): + continue + # Drop units whose section was filtered out of the tree. + if section_id and section_id not in section_path_by_id: + continue - unit = UnitRow( - chunk_id=chunk_id, - section_id=section_id, - chunk_type=str(row[4] or "text"), - content=str(row[5] or ""), - sort_order=int(row[6] or 0), - source_chunk_path=str(row[7] or ""), - file_path=str(row[8] or ""), - metadata=_as_meta(row[9]), - ) - by_doc.setdefault(document_id, []).append(unit) - if chunk_id: - meta = { - "document_id": document_id, - "section_path": section_path, - "chunk_type": unit.chunk_type, - "file_path": unit.file_path or None, - "job_id": job_id_by_result_id.get(job_result_id), - } - # Bare chunk_id (last-wins) plus doc-scoped key so the same - # chunk_id can appear under multiple documents. - ref_index[chunk_id] = meta - ref_index[f"{document_id}:{chunk_id}"] = meta - last_row = rows[-1] - last_key = ( - str(last_row[0]), - str(last_row[1]), - int(last_row[6] or 0), - str(last_row[2] or ""), - str(last_row[10]), - ) - if len(rows) < _CHUNK_BATCH_SIZE: - break + unit = UnitRow( + chunk_id=chunk_id, + section_id=section_id, + chunk_type=str(row[3] or "text"), + content=str(row[4] or ""), + sort_order=int(row[5] or 0), + source_chunk_path=str(row[6] or ""), + file_path=str(row[7] or ""), + metadata=_as_meta(row[8]), + ) + by_doc.setdefault(document_id, []).append(unit) + if chunk_id: + meta = { + "document_id": document_id, + "section_path": section_path, + "chunk_type": unit.chunk_type, + "file_path": unit.file_path or None, + "job_id": str(row[10]) if row[10] else None, + } + # Bare chunk_id (last-wins) plus doc-scoped key so the same + # chunk_id can appear under multiple documents. + ref_index[chunk_id] = meta + ref_index[f"{document_id}:{chunk_id}"] = meta return by_doc, ref_index diff --git a/packages/shared-python/shared/tests/test_nav_snapshot.py b/packages/shared-python/shared/tests/test_nav_snapshot.py index f8bf76f6..340c7068 100644 --- a/packages/shared-python/shared/tests/test_nav_snapshot.py +++ b/packages/shared-python/shared/tests/test_nav_snapshot.py @@ -93,3 +93,15 @@ def test_build_nav_snapshot_rejects_empty_corpus() -> None: units_by_doc={}, chunk_ref_index={}, ) + + +def test_load_nav_snapshot_joins_current_revision_only() -> None: + """Section/chunk loaders must bind rows to Document.current_job_result_id.""" + import inspect + + from shared.services.retrieval import nav_snapshot as nav_snapshot_mod + + sections_src = "".join(inspect.getsource(nav_snapshot_mod._load_sections).split()) + chunks_src = "".join(inspect.getsource(nav_snapshot_mod._load_chunks).split()) + assert "DocumentSection.job_result_id==Document.current_job_result_id" in sections_src + assert "DocumentChunk.job_result_id==Document.current_job_result_id" in chunks_src From 4a5bfe84a537824293a92578c9f9ba392f2940e2 Mon Sep 17 00:00:00 2001 From: suguanYang Date: Mon, 31 Aug 2026 11:10:18 +0800 Subject: [PATCH 10/10] Revert "Merge pull request #348 from Ontos-AI/fix/wangbinqi/backfill-script-container-path" This reverts commit d796ff475161890323cee42b260d27a076a74a05, reversing changes made to 8420fad48db54a9c0ea4495208d6007a95a7fe57. --- ...test_backfill_map_unit_indexes_contract.py | 28 ------------------- 1 file changed, 28 deletions(-) delete mode 100644 apps/api/tests/contract/test_backfill_map_unit_indexes_contract.py diff --git a/apps/api/tests/contract/test_backfill_map_unit_indexes_contract.py b/apps/api/tests/contract/test_backfill_map_unit_indexes_contract.py deleted file mode 100644 index 72c61f91..00000000 --- a/apps/api/tests/contract/test_backfill_map_unit_indexes_contract.py +++ /dev/null @@ -1,28 +0,0 @@ -from __future__ import annotations - -from pathlib import Path - - -def test_backfill_script_resolves_shared_package_from_runtime_image_layout( - tmp_path: Path, -) -> None: - from scripts.backfill_map_unit_indexes import _resolve_shared_root - - api_root = tmp_path / "app" - shared_root = api_root / "packages" / "shared-python" - shared_root.mkdir(parents=True) - - assert _resolve_shared_root(api_root) == shared_root - - -def test_backfill_script_resolves_shared_package_from_source_checkout_layout( - tmp_path: Path, -) -> None: - from scripts.backfill_map_unit_indexes import _resolve_shared_root - - repository_root = tmp_path / "repository" - api_root = repository_root / "apps" / "api" - shared_root = repository_root / "packages" / "shared-python" - shared_root.mkdir(parents=True) - - assert _resolve_shared_root(api_root) == shared_root