From 04099252b8f733cd3555ff51f77b26efb546fcf8 Mon Sep 17 00:00:00 2001 From: chengke <404835780@qq.com> Date: Mon, 31 Aug 2026 12:40:37 +0800 Subject: [PATCH 1/6] feat: add conversation_id to retrieval queries and enhance namespace map snapshot handling --- ...b_add_retrieval_namespace_map_snapshots.py | 39 ++++ apps/api/app/api/v1/routes/retrieval.py | 9 + .../services/documents/lifecycle_service.py | 11 + ...etrieval_lazy_snapshot_quality_contract.py | 24 +-- .../test_retrieval_map_unit_index_contract.py | 24 ++- .../test_retrieval_term_score_contract.py | 75 ------- .../shared/models/database/document.py | 32 +++ .../shared/services/retrieval/app_service.py | 2 + .../services/retrieval/execution/plan.py | 2 + .../retrieval/execution/query_request.py | 4 + .../retrieval/execution/route_types.py | 1 + .../services/retrieval/execution/routes.py | 4 +- .../retrieval/namespace_map_snapshot.py | 150 +++++++++++++ .../retrieval/namespace_map_snapshot_cache.py | 43 ++++ .../services/retrieval/nav/knowhere_hybrid.py | 46 +--- .../services/retrieval/nav/nav_knowhere.py | 107 ---------- .../shared/services/retrieval/nav_snapshot.py | 199 ++++++++++++++---- .../services/retrieval/publication_content.py | 6 +- .../services/retrieval/publication_service.py | 9 + .../services/retrieval/serving_manifest.py | 9 +- 20 files changed, 519 insertions(+), 277 deletions(-) create mode 100644 apps/api/alembic/versions/6c7d8e9f0a1b_add_retrieval_namespace_map_snapshots.py delete mode 100644 apps/api/tests/contract/test_retrieval_term_score_contract.py create mode 100644 packages/shared-python/shared/services/retrieval/namespace_map_snapshot.py create mode 100644 packages/shared-python/shared/services/retrieval/namespace_map_snapshot_cache.py diff --git a/apps/api/alembic/versions/6c7d8e9f0a1b_add_retrieval_namespace_map_snapshots.py b/apps/api/alembic/versions/6c7d8e9f0a1b_add_retrieval_namespace_map_snapshots.py new file mode 100644 index 000000000..1f7b7c019 --- /dev/null +++ b/apps/api/alembic/versions/6c7d8e9f0a1b_add_retrieval_namespace_map_snapshots.py @@ -0,0 +1,39 @@ +"""Add persisted namespace-level MAP snapshot table.""" + +from __future__ import annotations + +from collections.abc import Sequence + +from alembic import op +import sqlalchemy as sa + + +revision: str = "6c7d8e9f0a1b" +down_revision: str | None = "5b6c7d8e9f0a" +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_map_snapshots"): + op.create_table( + "retrieval_namespace_map_snapshots", + 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("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("updated_at", sa.DateTime(), nullable=False), + sa.PrimaryKeyConstraint("id"), + sa.UniqueConstraint( + "user_id", + "namespace", + name="uq_retrieval_namespace_map_snapshots_scope", + ), + ) + + +def downgrade() -> None: + op.drop_table("retrieval_namespace_map_snapshots", if_exists=True) diff --git a/apps/api/app/api/v1/routes/retrieval.py b/apps/api/app/api/v1/routes/retrieval.py index 758f111e0..3c9202dfa 100644 --- a/apps/api/app/api/v1/routes/retrieval.py +++ b/apps/api/app/api/v1/routes/retrieval.py @@ -77,6 +77,14 @@ class RetrievalQueryRequest(BaseModel): "Set false to force classic 3-channel top-K retrieval." ), ) + conversation_id: str | None = Field( + None, + max_length=255, + description=( + "Caller-supplied conversation identifier, threaded through for " + "retrieval tracing. Does not affect caching or result content." + ), + ) @field_validator("channels") @classmethod @@ -172,6 +180,7 @@ async def execute_retrieval_query( threshold=payload.threshold, internal_recall_k=payload.internal_recall_k, use_agentic=payload.use_agentic, + conversation_id=payload.conversation_id, llm_config=llm_config, ) diff --git a/apps/api/app/services/documents/lifecycle_service.py b/apps/api/app/services/documents/lifecycle_service.py index 6b39507d9..3ea08f05d 100644 --- a/apps/api/app/services/documents/lifecycle_service.py +++ b/apps/api/app/services/documents/lifecycle_service.py @@ -20,6 +20,9 @@ invalidate_retrieval_cache_namespaces, ) from shared.services.retrieval.graph.service import DocumentGraphService, GraphScope +from shared.services.retrieval.namespace_map_snapshot import ( + remove_document_from_namespace_map_snapshot, +) from shared.services.retrieval.serving_generation import ( advance_namespace_generation, lock_namespace_generation, @@ -483,6 +486,14 @@ async def archive_document( namespace=previous_namespace, ) ) + await db.run_sync( + lambda sync_db: remove_document_from_namespace_map_snapshot( + sync_db, + user_id=user_id, + namespace=previous_namespace, + document_id=document_id, + ) + ) await db.run_sync( lambda sync_db: advance_namespace_generation( sync_db, 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 0752a3b97..a56260d56 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 @@ -220,7 +220,11 @@ def test_lazy_provider_preserves_score_units_and_scores() -> None: ] -def test_streaming_scorer_preserves_exact_eager_scores() -> None: +def test_streaming_scorer_preserves_exact_eager_scores(monkeypatch: Any) -> None: + # Corpus-wide map-nav scoring retired the term channel (see the + # unify-bm25-persistent-map plan); zero its weight so the eager oracle + # is directly comparable to the path+content-only streaming scorer. + monkeypatch.setenv("NAV_MAP_CHANNEL_WEIGHT_TERM", "0") rows: list[ScoreUnitRow] = [ { "chunk_id": "unit-a", @@ -257,7 +261,12 @@ def unit_factory() -> Sequence[ScoreUnitRow]: assert replay_count == 1 -def test_streaming_scorer_preserves_duplicate_id_eager_semantics() -> None: +def test_streaming_scorer_preserves_duplicate_id_eager_semantics( + monkeypatch: Any, +) -> None: + # See test_streaming_scorer_preserves_exact_eager_scores: term is retired + # from corpus-wide scoring, so the eager oracle's term weight is zeroed. + monkeypatch.setenv("NAV_MAP_CHANNEL_WEIGHT_TERM", "0") rows: list[ScoreUnitRow] = [ { "chunk_id": "duplicate", @@ -383,17 +392,6 @@ def build_stats(search_field: str) -> PersistedBm25Stats: 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 ], 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 eb124fb2c..6d30731cf 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 @@ -2,6 +2,7 @@ from collections.abc import Callable, Mapping, Sequence from contextlib import AbstractAsyncContextManager +from typing import Any from uuid import uuid4 from httpx import AsyncClient @@ -33,6 +34,7 @@ ) from shared.services.retrieval.nav_bridge import build_referenced_chunks from shared.services.retrieval.publication_models import DocumentPublicationScope +from shared.services.retrieval.serving_generation import lock_namespace_generation from tests.support.contract_database import ContractDatabase from tests.support.retrieval_snapshot_support import contract_db_session @@ -161,7 +163,7 @@ async def test_published_map_units_preserve_scores_without_chunk_payload_reads( ] async with contract_db_session() as db: await db.run_sync( - lambda sync_db: replace_document_revision_content( + lambda sync_db: _publish_revision_with_generation_lock( sync_db, scope=scope, chunks=chunks, @@ -346,7 +348,7 @@ async def test_lazy_snapshot_defers_selected_asset_reference_metadata( ] async with contract_db_session() as db: await db.run_sync( - lambda sync_db: replace_document_revision_content( + lambda sync_db: _publish_revision_with_generation_lock( sync_db, scope=scope, chunks=chunks, @@ -515,6 +517,24 @@ def test_titleless_leaf_has_identical_eager_and_lazy_path_scoring() -> None: ) +def _publish_revision_with_generation_lock( + sync_db: Any, + *, + scope: DocumentPublicationScope, + chunks: list[dict[str, Any]], +) -> None: + """Mirror the production publish sequence: lock, then patch the snapshot. + + ``replace_document_revision_content`` requires the namespace generation + row to already exist (``publication_service.py`` locks it before every + real publish); this test helper reproduces that precondition. + """ + lock_namespace_generation( + sync_db, user_id=scope.user_id, namespace=scope.namespace + ) + replace_document_revision_content(sync_db, scope=scope, chunks=chunks) + + async def _seed_revision( *, namespace: str, 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 425f4de03..000000000 --- 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/packages/shared-python/shared/models/database/document.py b/packages/shared-python/shared/models/database/document.py index c396283d2..b85dd84f3 100644 --- a/packages/shared-python/shared/models/database/document.py +++ b/packages/shared-python/shared/models/database/document.py @@ -518,6 +518,38 @@ class RetrievalNamespaceTokenStat(Base): ) +class RetrievalNamespaceMapSnapshot(Base): + """Persisted namespace-level MAP (sections + chunk index + map units). + + Incrementally patched at publish/archive time (one document's subtree at + a time); query time reads this row directly instead of merging per-file + manifests. Overwritten in place -- no generation history is retained. + """ + + __tablename__ = "retrieval_namespace_map_snapshots" + + id: Mapped[str] = mapped_column( + String(100), primary_key=True, default=lambda: f"rnmap_{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) + 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) + 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_map_snapshots_scope", + ), + ) + + class GraphNode(Base): """Persisted derived graph node used for routing and expansion.""" diff --git a/packages/shared-python/shared/services/retrieval/app_service.py b/packages/shared-python/shared/services/retrieval/app_service.py index 83ce65b78..69727d1e5 100644 --- a/packages/shared-python/shared/services/retrieval/app_service.py +++ b/packages/shared-python/shared/services/retrieval/app_service.py @@ -32,6 +32,7 @@ async def run_retrieval_query( threshold: float = 0.0, internal_recall_k: int | None = None, use_agentic: bool | None = None, + conversation_id: str | None = None, llm_config: LLMConfig | None = None, ) -> dict[str, Any]: return await execute_retrieval_query( @@ -51,5 +52,6 @@ async def run_retrieval_query( threshold=threshold, internal_recall_k=internal_recall_k, use_agentic=use_agentic, + conversation_id=conversation_id, llm_config=llm_config, ) diff --git a/packages/shared-python/shared/services/retrieval/execution/plan.py b/packages/shared-python/shared/services/retrieval/execution/plan.py index adb2a6da4..8e9a9c7fe 100644 --- a/packages/shared-python/shared/services/retrieval/execution/plan.py +++ b/packages/shared-python/shared/services/retrieval/execution/plan.py @@ -48,6 +48,7 @@ async def run_retrieval_query( threshold: float = 0.0, internal_recall_k: int | None = None, use_agentic: bool | None = None, + conversation_id: str | None = None, llm_config: LLMConfig | None = None, ) -> dict[str, Any]: """Run retrieval through the plan module.""" @@ -69,6 +70,7 @@ async def run_retrieval_query( threshold=threshold, internal_recall_k=internal_recall_k, use_agentic=use_agentic, + conversation_id=conversation_id, llm_config=llm_config, ) ).execute() diff --git a/packages/shared-python/shared/services/retrieval/execution/query_request.py b/packages/shared-python/shared/services/retrieval/execution/query_request.py index 53363d185..9e831bce0 100644 --- a/packages/shared-python/shared/services/retrieval/execution/query_request.py +++ b/packages/shared-python/shared/services/retrieval/execution/query_request.py @@ -31,6 +31,7 @@ class RetrievalQuery: threshold: float = 0.0 internal_recall_k: int | None = None use_agentic: bool | None = None + conversation_id: str | None = None llm_config: LLMConfig | None = None @classmethod @@ -53,6 +54,7 @@ def from_parameters( threshold: float = 0.0, internal_recall_k: int | None = None, use_agentic: bool | None = None, + conversation_id: str | None = None, llm_config: LLMConfig | None = None, ) -> "RetrievalQuery": return cls( @@ -72,6 +74,7 @@ def from_parameters( threshold=threshold, internal_recall_k=internal_recall_k, use_agentic=use_agentic, + conversation_id=conversation_id, llm_config=llm_config, ) @@ -125,4 +128,5 @@ def build_route_context(self) -> RetrievalRouteContext: internal_recall_k=self.internal_recall_k, effective_recall_k=self.resolve_effective_recall_k(), use_agentic=self.use_agentic, + conversation_id=self.conversation_id, ) 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 13c21c46a..339180904 100644 --- a/packages/shared-python/shared/services/retrieval/execution/route_types.py +++ b/packages/shared-python/shared/services/retrieval/execution/route_types.py @@ -28,6 +28,7 @@ class RetrievalRouteContext: internal_recall_k: int | None effective_recall_k: int use_agentic: bool | None + conversation_id: str | None = None revision_pins: RetrievalRevisionPins | None = None diff --git a/packages/shared-python/shared/services/retrieval/execution/routes.py b/packages/shared-python/shared/services/retrieval/execution/routes.py index 0998a8df9..fe47be194 100644 --- a/packages/shared-python/shared/services/retrieval/execution/routes.py +++ b/packages/shared-python/shared/services/retrieval/execution/routes.py @@ -234,10 +234,12 @@ async def _run_mapnav_route( ) snapshot_seconds = time.perf_counter() - snapshot_started logger.info( - "retrieval mapnav stage=snapshot_load seconds={:.3f} documents={} refs={}".format( + "retrieval mapnav stage=snapshot_load seconds={:.3f} documents={} refs={} " + "conversation_id={}".format( snapshot_seconds, len(snapshot.document_ids), len(snapshot.chunk_ref_index), + context.conversation_id or "", ) ) diff --git a/packages/shared-python/shared/services/retrieval/namespace_map_snapshot.py b/packages/shared-python/shared/services/retrieval/namespace_map_snapshot.py new file mode 100644 index 000000000..bd8b7767f --- /dev/null +++ b/packages/shared-python/shared/services/retrieval/namespace_map_snapshot.py @@ -0,0 +1,150 @@ +"""Incrementally-patched namespace-level MAP snapshot (sections + chunk index). + +Callers must already hold the namespace generation lock (see +``serving_generation.lock_namespace_generation``) before calling either +function here, exactly as ``rebuild_namespace_serving_statistics`` requires. +Each call only touches one document's subtree; every other document's +subtree in the payload is left byte-for-byte unchanged. +""" + +from __future__ import annotations + +from typing import Any + +from sqlalchemy import select +from sqlalchemy.orm import Session + +from shared.models.database.document import ( + RetrievalNamespaceGeneration, + RetrievalNamespaceMapSnapshot, +) +from shared.services.retrieval.publication_models import DocumentPublicationScope +from shared.services.retrieval.serving_manifest import ( + decode_serving_manifest, + encode_serving_manifest, +) + + +def patch_namespace_map_snapshot( + db: Session, + *, + scope: DocumentPublicationScope, + manifest_payload: dict[str, Any], +) -> None: + """Replace one document's subtree in the namespace MAP snapshot.""" + row = _load_snapshot_row(db, user_id=scope.user_id, namespace=scope.namespace) + documents = _decode_documents(row) + documents[scope.document_id] = { + "job_result_id": manifest_payload.get("job_result_id"), + "job_id": manifest_payload.get("job_id"), + "source_file_name": manifest_payload.get("source_file_name"), + "sections": manifest_payload.get("sections") or [], + "chunks": manifest_payload.get("chunks") or [], + } + target_generation = _target_generation( + db, user_id=scope.user_id, namespace=scope.namespace + ) + _write_snapshot( + db, + row=row, + user_id=scope.user_id, + namespace=scope.namespace, + documents=documents, + target_generation=target_generation, + ) + + +def remove_document_from_namespace_map_snapshot( + db: Session, + *, + user_id: str, + namespace: str, + document_id: str, +) -> None: + """Drop one document's subtree from the namespace MAP snapshot (archive path).""" + row = _load_snapshot_row(db, user_id=user_id, namespace=namespace) + documents = _decode_documents(row) + if document_id not in documents: + return + del documents[document_id] + target_generation = _target_generation(db, user_id=user_id, namespace=namespace) + _write_snapshot( + db, + row=row, + user_id=user_id, + namespace=namespace, + documents=documents, + target_generation=target_generation, + ) + + +def _target_generation(db: Session, *, user_id: str, namespace: str) -> int: + """Namespace generation this snapshot is prepared for (current + 1). + + Mirrors ``rebuild_namespace_serving_statistics``: callers advance the + generation after this write, in the same transaction. + """ + generation = db.execute( + select(RetrievalNamespaceGeneration) + .where(RetrievalNamespaceGeneration.user_id == user_id) + .where(RetrievalNamespaceGeneration.namespace == namespace) + .with_for_update() + ).scalar_one() + return int(generation.generation) + 1 + + +def _load_snapshot_row( + db: Session, *, user_id: str, namespace: str +) -> RetrievalNamespaceMapSnapshot | None: + return db.execute( + select(RetrievalNamespaceMapSnapshot) + .where(RetrievalNamespaceMapSnapshot.user_id == user_id) + .where(RetrievalNamespaceMapSnapshot.namespace == namespace) + ).scalar_one_or_none() + + +def _decode_documents( + row: RetrievalNamespaceMapSnapshot | None, +) -> dict[str, dict[str, Any]]: + if row is None: + return {} + try: + payload = decode_serving_manifest( + row.payload_zlib, + checksum=row.checksum, + format_version=row.format_version, + ) + except ValueError: + return {} + documents = payload.get("documents") + return dict(documents) if isinstance(documents, dict) else {} + + +def _write_snapshot( + db: Session, + *, + row: RetrievalNamespaceMapSnapshot | None, + user_id: str, + namespace: str, + documents: dict[str, dict[str, Any]], + target_generation: int, +) -> None: + encoded, checksum, format_version = encode_serving_manifest( + {"documents": documents} + ) + if row is None: + db.add( + RetrievalNamespaceMapSnapshot( + user_id=user_id, + namespace=namespace, + generation=target_generation, + format_version=format_version, + payload_zlib=encoded, + checksum=checksum, + ) + ) + else: + row.generation = target_generation + row.format_version = format_version + row.payload_zlib = encoded + row.checksum = checksum diff --git a/packages/shared-python/shared/services/retrieval/namespace_map_snapshot_cache.py b/packages/shared-python/shared/services/retrieval/namespace_map_snapshot_cache.py new file mode 100644 index 000000000..6cc5cab29 --- /dev/null +++ b/packages/shared-python/shared/services/retrieval/namespace_map_snapshot_cache.py @@ -0,0 +1,43 @@ +"""Process-local cache for decoded namespace MAP snapshot documents. + +Keyed by ``(user_id, namespace, generation)`` so a publish/archive that bumps +the namespace generation invalidates stale entries automatically -- no manual +invalidation call is needed. Bounded LRU keeps memory use predictable across +many namespaces sharing one worker process. +""" + +from __future__ import annotations + +import threading +from collections import OrderedDict +from typing import Any + +_MAX_ENTRIES = 64 +_lock = threading.Lock() +_cache: "OrderedDict[tuple[str, str, int], dict[str, dict[str, Any]]]" = OrderedDict() + + +def get_cached_namespace_documents( + *, user_id: str, namespace: str, generation: int +) -> dict[str, dict[str, Any]] | None: + key = (user_id, namespace, generation) + with _lock: + documents = _cache.get(key) + if documents is not None: + _cache.move_to_end(key) + return documents + + +def cache_namespace_documents( + *, + user_id: str, + namespace: str, + generation: int, + documents: dict[str, dict[str, Any]], +) -> None: + key = (user_id, namespace, generation) + with _lock: + _cache[key] = documents + _cache.move_to_end(key) + while len(_cache) > _MAX_ENTRIES: + _cache.popitem(last=False) diff --git a/packages/shared-python/shared/services/retrieval/nav/knowhere_hybrid.py b/packages/shared-python/shared/services/retrieval/nav/knowhere_hybrid.py index 369da962c..3c237aef7 100644 --- a/packages/shared-python/shared/services/retrieval/nav/knowhere_hybrid.py +++ b/packages/shared-python/shared/services/retrieval/nav/knowhere_hybrid.py @@ -664,7 +664,11 @@ 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.""" + """Score several queries exactly while reading the corpus only once. + + Corpus-wide map-nav scoring is path+content only; the term channel was + retired here (see the unify-bm25-persistent-map plan). + """ unique_queries = list(dict.fromkeys(str(query) for query in queries)) if not unique_queries: return {} @@ -677,7 +681,6 @@ def score_unit_stream_hybrid_many( 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] = [] path_stats = _StreamingBm25Stats.empty() content_stats = _StreamingBm25Stats.empty() @@ -693,20 +696,6 @@ def score_unit_stream_hybrid_many( 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) units.append( _StreamingManyUnit( unit_id=unit_id, @@ -722,7 +711,6 @@ def score_unit_stream_hybrid_many( for token in query_token_set if content_frequencies[token] }, - term_scores=tuple(term_scores), ) ) path_stats.finalize() @@ -733,9 +721,8 @@ def score_unit_stream_hybrid_many( 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) + for query in unique_queries } @@ -745,11 +732,9 @@ def _score_streaming_units( 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] = {} unit_ids = list(dict.fromkeys(unit.unit_id for unit in units)) for unit in units: path_score = path_stats.score( @@ -760,11 +745,6 @@ 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 - ) path_rows = [ (score, unit_id) for unit_id, score in path_by_id.items() if score > 0.0 @@ -772,13 +752,9 @@ def _score_streaming_units( 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() + path_weight, content_weight, _term_weight = map_channel_weights() rrf_k = int( os.environ.get( "NAV_MAP_RRF_K", @@ -791,8 +767,6 @@ def _score_streaming_units( 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()} @@ -803,7 +777,6 @@ class _StreamingManyUnit: content_length: int path_frequencies: Mapping[str, int] content_frequencies: Mapping[str, int] - term_scores: Tuple[float, ...] @dataclass(frozen=True) @@ -825,7 +798,6 @@ class PersistedScoreUnit: content_length: int path_frequencies: Mapping[str, int] content_frequencies: Mapping[str, int] - term_scores: Tuple[float, ...] @dataclass(frozen=True) @@ -854,7 +826,6 @@ def score_persisted_corpus_many( 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 ] @@ -864,9 +835,8 @@ def score_persisted_corpus_many( 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) + for query in unique_queries } 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 e25234eac..757f32121 100644 --- a/packages/shared-python/shared/services/retrieval/nav/nav_knowhere.py +++ b/packages/shared-python/shared/services/retrieval/nav/nav_knowhere.py @@ -55,11 +55,6 @@ _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__) @@ -236,9 +231,6 @@ def __init__( 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: @@ -559,16 +551,6 @@ def load_persisted_score_corpus( 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), @@ -602,9 +584,6 @@ def load_persisted_score_corpus( 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 ], @@ -614,92 +593,6 @@ def load_persisted_score_corpus( 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 {} - # Keep candidate selection in PostgreSQL so the trigram index can - # discard non-matching units, but compute the exact score once per - # returned row in Python. The previous query generated one POSITION - # expression and one OR predicate for every query token. Long planner - # subgoals therefore produced very large SQL statements and repeated - # substring evaluation for the same row. LIKE ANY keeps the SQL shape - # constant while preserving the existing literal substring semantics in - # the final Python scoring pass (LIKE may over-select wildcard matches, - # which are rejected by the literal checks below). - candidate_patterns: list[str] = [] - for query in queries: - query_lower = query.lower().strip() - if not query_lower: - 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), - ) - 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, - ) - scores_by_unit: Dict[str, Tuple[float, ...]] = {} - for row in rows: - if len(row) < 2: - continue - unit_id = str(row[0]) - haystack = str(row[1] or "").lower() - scores: list[float] = [] - for query in queries: - query_lower = query.lower().strip() - if not query_lower: - scores.append(0.0) - elif query_lower in haystack: - scores.append(100.0) - else: - scores.append( - float( - sum( - 1 - for token in query_tokens_by_query[query] - if token in haystack - ) - ) - ) - if any(score > 0.0 for score in scores): - scores_by_unit[unit_id] = tuple(scores) - return scores_by_unit - def _load_persisted_bm25_stats( self, cur: "_SyncCursor", diff --git a/packages/shared-python/shared/services/retrieval/nav_snapshot.py b/packages/shared-python/shared/services/retrieval/nav_snapshot.py index 5d34e4510..95d24078e 100644 --- a/packages/shared-python/shared/services/retrieval/nav_snapshot.py +++ b/packages/shared-python/shared/services/retrieval/nav_snapshot.py @@ -31,6 +31,7 @@ Document, DocumentChunk, DocumentSection, + RetrievalNamespaceMapSnapshot, RetrievalServingRevisionManifest, ) from shared.models.database.job_result import JobResult @@ -46,6 +47,10 @@ 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 +from shared.services.retrieval.namespace_map_snapshot_cache import ( + cache_namespace_documents, + get_cached_namespace_documents, +) # Keep each payload query bounded under the API's 30-second statement timeout. @@ -212,15 +217,40 @@ 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, + snapshot_entries = await _resolve_namespace_snapshot_entries( + db, + user_id=user_id, + namespace=namespace, + document_revisions=document_revisions, + ) + if snapshot_entries is not None: + manifest_sections = _parse_manifest_entries( + snapshot_entries, exclude_sections=excluded_secs, job_id_by_result_id=job_id_by_result_id, ) + else: + _logger.warning( + "retrieval snapshot fallback=manifest_merge user_id=%s namespace=%s documents=%d", + user_id, + namespace, + len(document_revisions), + ) + 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, + ) if manifest_sections is None: + _logger.warning( + "retrieval snapshot fallback=table_scan user_id=%s namespace=%s documents=%d", + user_id, + namespace, + len(document_revisions), + ) sections_by_doc, section_path_by_id = await _load_sections( db, document_revisions=document_revisions, @@ -340,18 +370,72 @@ async def load_nav_snapshot( ) -async def _load_manifest_sections( +async def _resolve_namespace_snapshot_entries( db: SnapshotSession, *, + user_id: str, + namespace: str, 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.""" +) -> list[tuple[object, ...]] | None: + """Return manifest-shaped entries from the persisted namespace snapshot. + + Returns ``None`` (triggering the exact per-revision fallback) when the + snapshot row is missing, corrupt, or stale for any requested revision. + """ + statement = select( + RetrievalNamespaceMapSnapshot.generation, + RetrievalNamespaceMapSnapshot.payload_zlib, + RetrievalNamespaceMapSnapshot.checksum, + RetrievalNamespaceMapSnapshot.format_version, + ).where( + RetrievalNamespaceMapSnapshot.user_id == user_id, + RetrievalNamespaceMapSnapshot.namespace == namespace, + ) + try: + row = (await db.execute(statement)).first() + except Exception: + await db.rollback() + return None + if row is None: + return None + generation, payload_zlib, checksum, format_version = row + documents = get_cached_namespace_documents( + user_id=user_id, namespace=namespace, generation=int(generation) + ) + if documents is None: + try: + payload = decode_serving_manifest( + bytes(payload_zlib), + checksum=str(checksum), + format_version=int(format_version), + ) + except (ValueError, TypeError): + return None + decoded_documents = payload.get("documents") + if not isinstance(decoded_documents, dict): + return None + documents = decoded_documents + cache_namespace_documents( + user_id=user_id, + namespace=namespace, + generation=int(generation), + documents=documents, + ) + entries: list[tuple[object, ...]] = [] + for document_id, job_result_id in document_revisions: + entry = documents.get(document_id) + if not isinstance(entry, dict) or str(entry.get("job_result_id") or "") != job_result_id: + return None + entries.append((document_id, job_result_id, entry, None, None)) + return entries + + +async def _resolve_manifest_entries( + db: SnapshotSession, + *, + document_revisions: list[tuple[str, str]], +) -> list[tuple[object, ...]] | None: + """Resolve per-revision manifest rows from the request cache or table.""" cached_payloads = get_cached_manifest_payloads( db, revisions=dict(document_revisions), @@ -365,32 +449,71 @@ 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( + return manifest_entries + + 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( + 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_(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) + return manifest_entries + + +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.""" + manifest_entries = await _resolve_manifest_entries( + db, document_revisions=document_revisions + ) + if manifest_entries is None: + return None + return _parse_manifest_entries( + manifest_entries, + exclude_sections=exclude_sections, + job_id_by_result_id=job_id_by_result_id, + ) + + +def _parse_manifest_entries( + manifest_entries: list[tuple[object, ...]], + *, + 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: + """Parse manifest-shaped entries (decoded dict or raw compressed row) into sections.""" 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/publication_content.py b/packages/shared-python/shared/services/retrieval/publication_content.py index 340924367..7ce75b808 100644 --- a/packages/shared-python/shared/services/retrieval/publication_content.py +++ b/packages/shared-python/shared/services/retrieval/publication_content.py @@ -13,6 +13,9 @@ DocumentSection, ) from shared.services.retrieval.map_unit_index import replace_document_map_units +from shared.services.retrieval.namespace_map_snapshot import ( + patch_namespace_map_snapshot, +) 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 ( @@ -93,7 +96,8 @@ def replace_document_revision_content( db.flush() replace_document_map_units(db, scope=scope) db.flush() - persist_revision_serving_state(db, scope=scope) + manifest_payload = persist_revision_serving_state(db, scope=scope) + patch_namespace_map_snapshot(db, scope=scope, manifest_payload=manifest_payload) 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 d7c879dec..58fa98dde 100644 --- a/packages/shared-python/shared/services/retrieval/publication_service.py +++ b/packages/shared-python/shared/services/retrieval/publication_service.py @@ -22,6 +22,9 @@ from shared.models.schemas.job_metadata import JobMetadataHelper from shared.models.schemas.retrieval_namespace import normalize_retrieval_namespace from shared.services.retrieval.graph.service import DocumentGraphService, GraphScope +from shared.services.retrieval.namespace_map_snapshot import ( + remove_document_from_namespace_map_snapshot, +) from shared.services.retrieval.publication_content import ( deduplicate_chunks_by_source_path, replace_document_revision_content, @@ -191,6 +194,12 @@ def _publish_document_state_for_job( user_id=scope.user_id, namespace=str(existing_namespace), ) + remove_document_from_namespace_map_snapshot( + db, + user_id=scope.user_id, + namespace=str(existing_namespace), + document_id=document.document_id, + ) advance_namespace_generation( db, user_id=scope.user_id, diff --git a/packages/shared-python/shared/services/retrieval/serving_manifest.py b/packages/shared-python/shared/services/retrieval/serving_manifest.py index b2dd40c5e..5fe979e08 100644 --- a/packages/shared-python/shared/services/retrieval/serving_manifest.py +++ b/packages/shared-python/shared/services/retrieval/serving_manifest.py @@ -190,8 +190,12 @@ def persist_revision_serving_state( db: Session, *, scope: DocumentPublicationScope, -) -> None: - """Replace manifest and statistics rows for one revision atomically.""" +) -> dict[str, Any]: + """Replace manifest and statistics rows for one revision atomically. + + Returns the manifest payload so callers can patch the namespace-level MAP + snapshot without rebuilding it. + """ manifest_payload = build_revision_serving_payload(db, scope=scope) statistics_payload = build_revision_statistics_payload(db, scope=scope) manifest_bytes, manifest_checksum, manifest_version = encode_serving_manifest( @@ -232,6 +236,7 @@ def persist_revision_serving_state( checksum=statistics_checksum, ) ) + return manifest_payload def rebuild_namespace_serving_statistics( From 6ab4f8e356790d5d63faffb1751f475b005290a8 Mon Sep 17 00:00:00 2001 From: chengke <404835780@qq.com> Date: Mon, 31 Aug 2026 15:18:10 +0800 Subject: [PATCH 2/6] feat: enhance map unit index backfill process and remove obsolete tests - Updated the backfill_map_unit_indexes script to include patching of the namespace map snapshot after persisting the revision serving state. - Removed outdated contract tests related to BM25 FTS prefilter and unit tests for BM25 channel token preparation. - Refactored retrieval logic to improve handling of document map units, including new fields for asset presence. --- ...d8e9f0a1b2c_add_map_unit_asset_presence.py | 63 +++ apps/api/scripts/backfill_map_unit_indexes.py | 15 +- .../test_bm25_fts_prefilter_contract.py | 231 -------- ...est_retrieval_classic_map_unit_contract.py | 282 ++++++++++ .../test_retrieval_map_unit_index_contract.py | 13 + .../test_retrieval_revision_races_contract.py | 63 +-- .../tests/unit/test_bm25_channel_tsquery.py | 45 -- .../shared/models/database/document.py | 18 + .../services/retrieval/execution/routes.py | 7 +- .../services/retrieval/map_unit_index.py | 7 + .../services/retrieval/nav/knowhere_hybrid.py | 3 +- .../services/retrieval/search/channels.py | 529 ------------------ .../services/retrieval/search/discovery.py | 216 ------- .../retrieval/search/lexical_ranker.py | 74 --- .../retrieval/search/map_unit_discovery.py | 500 +++++++++++++++++ .../services/retrieval/search/scoring.py | 29 - .../tests/test_retrieval_search_channels.py | 216 ------- 17 files changed, 900 insertions(+), 1411 deletions(-) create mode 100644 apps/api/alembic/versions/7d8e9f0a1b2c_add_map_unit_asset_presence.py delete mode 100644 apps/api/tests/contract/test_bm25_fts_prefilter_contract.py create mode 100644 apps/api/tests/contract/test_retrieval_classic_map_unit_contract.py delete mode 100644 apps/api/tests/unit/test_bm25_channel_tsquery.py delete mode 100644 packages/shared-python/shared/services/retrieval/search/channels.py delete mode 100644 packages/shared-python/shared/services/retrieval/search/discovery.py delete mode 100644 packages/shared-python/shared/services/retrieval/search/lexical_ranker.py create mode 100644 packages/shared-python/shared/services/retrieval/search/map_unit_discovery.py delete mode 100644 packages/shared-python/shared/tests/test_retrieval_search_channels.py diff --git a/apps/api/alembic/versions/7d8e9f0a1b2c_add_map_unit_asset_presence.py b/apps/api/alembic/versions/7d8e9f0a1b2c_add_map_unit_asset_presence.py new file mode 100644 index 000000000..b50e947ae --- /dev/null +++ b/apps/api/alembic/versions/7d8e9f0a1b2c_add_map_unit_asset_presence.py @@ -0,0 +1,63 @@ +"""Add has_image/has_table presence flags to document_map_units.""" + +from __future__ import annotations + +import sqlalchemy as sa +from alembic import op + +revision = "7d8e9f0a1b2c" +down_revision = "6c7d8e9f0a1b" +branch_labels = None +depends_on = None + + +def upgrade() -> None: + bind = op.get_bind() + inspector = sa.inspect(bind) + columns = {col["name"] for col in inspector.get_columns("document_map_units")} + if "has_image" not in columns: + op.add_column( + "document_map_units", + sa.Column( + "has_image", sa.Boolean(), nullable=False, server_default=sa.false() + ), + ) + if "has_table" not in columns: + op.add_column( + "document_map_units", + sa.Column( + "has_table", sa.Boolean(), nullable=False, server_default=sa.false() + ), + ) + + inspector = sa.inspect(bind) + indexes = {item["name"] for item in inspector.get_indexes("document_map_units")} + if "idx_document_map_units_has_image" not in indexes: + op.create_index( + "idx_document_map_units_has_image", + "document_map_units", + ["document_id", "job_result_id"], + postgresql_where=sa.text("has_image = true"), + ) + if "idx_document_map_units_has_table" not in indexes: + op.create_index( + "idx_document_map_units_has_table", + "document_map_units", + ["document_id", "job_result_id"], + postgresql_where=sa.text("has_table = true"), + ) + + +def downgrade() -> None: + inspector = sa.inspect(op.get_bind()) + indexes = {item["name"] for item in inspector.get_indexes("document_map_units")} + if "idx_document_map_units_has_table" in indexes: + op.drop_index("idx_document_map_units_has_table", table_name="document_map_units") + if "idx_document_map_units_has_image" in indexes: + op.drop_index("idx_document_map_units_has_image", table_name="document_map_units") + + columns = {col["name"] for col in inspector.get_columns("document_map_units")} + if "has_table" in columns: + op.drop_column("document_map_units", "has_table") + if "has_image" in columns: + op.drop_column("document_map_units", "has_image") diff --git a/apps/api/scripts/backfill_map_unit_indexes.py b/apps/api/scripts/backfill_map_unit_indexes.py index 2bfd41771..be2b4dbbf 100644 --- a/apps/api/scripts/backfill_map_unit_indexes.py +++ b/apps/api/scripts/backfill_map_unit_indexes.py @@ -1,7 +1,10 @@ """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 +Rebuilds, per active revision: the map-unit index, the revision serving +manifest, that document's subtree in the namespace MAP snapshot, namespace +statistics, and the namespace generation. The migrations that create these +derived tables leave them empty intentionally. Run this command after +deployment with ``--apply`` so each revision is rebuilt and committed independently; without ``--apply`` it is a read-only inventory. """ @@ -44,6 +47,9 @@ def _bootstrap_python_path() -> None: 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.namespace_map_snapshot import ( + patch_namespace_map_snapshot, +) from shared.services.retrieval.publication_models import DocumentPublicationScope from shared.services.retrieval.serving_generation import ( advance_namespace_generation, @@ -130,7 +136,10 @@ def backfill_map_unit_indexes(*, apply: bool, document_id: str = "") -> int: source_file_name=str(locked_document.source_file_name or ""), ) replace_document_map_units(db, scope=scope) - persist_revision_serving_state(db, scope=scope) + manifest_payload = persist_revision_serving_state(db, scope=scope) + patch_namespace_map_snapshot( + db, scope=scope, manifest_payload=manifest_payload + ) rebuild_namespace_serving_statistics( db, user_id=scope.user_id, diff --git a/apps/api/tests/contract/test_bm25_fts_prefilter_contract.py b/apps/api/tests/contract/test_bm25_fts_prefilter_contract.py deleted file mode 100644 index 56add8e5b..000000000 --- a/apps/api/tests/contract/test_bm25_fts_prefilter_contract.py +++ /dev/null @@ -1,231 +0,0 @@ -"""Contract tests for the BM25 channel Postgres FTS prefilter. - -These run against a real Postgres so the prefilter is validated against the -same generated tsvector columns and GIN indexes production uses. A pure-Python -fake would not catch a mismatch between the query configuration and the one -the columns were generated with. -""" - -from __future__ import annotations - -from collections.abc import AsyncGenerator - -import pytest -import pytest_asyncio -from shared.services.retrieval.search.channels import content_channel, path_channel -from shared.testing.contract_runtime import PostgreSQLProcess -from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker, create_async_engine -from sqlalchemy import text - -_SCHEMA = """ -CREATE TABLE documents ( - document_id TEXT PRIMARY KEY, - user_id TEXT, - namespace TEXT, - status TEXT, - current_job_result_id TEXT, - source_file_name TEXT -); -CREATE TABLE job_results (id TEXT PRIMARY KEY, job_id TEXT); -CREATE TABLE document_sections (section_id TEXT PRIMARY KEY, section_path TEXT); -CREATE TABLE document_chunks ( - id SERIAL PRIMARY KEY, - chunk_id TEXT, - document_id TEXT, - section_id TEXT, - chunk_type TEXT, - content TEXT, - source_chunk_path TEXT, - file_path TEXT, - chunk_metadata JSONB, - job_result_id TEXT, - sort_order INTEGER, - content_search_text TEXT, - content_search_tsv TSVECTOR GENERATED ALWAYS AS - (to_tsvector('simple', COALESCE(content_search_text, ''))) STORED, - path_search_text TEXT, - path_search_tsv TSVECTOR GENERATED ALWAYS AS - (to_tsvector('simple', COALESCE(path_search_text, ''))) STORED, - term_search_text TEXT -); -CREATE INDEX idx_chunk_content_search_tsv ON document_chunks USING GIN (content_search_tsv); -CREATE INDEX idx_chunk_path_search_tsv ON document_chunks USING GIN (path_search_tsv); -""" - -_NOISE_ROWS = 300 - - -@pytest_asyncio.fixture -async def seeded_session( - postgresql_proc: PostgreSQLProcess, -) -> AsyncGenerator[AsyncSession, None]: - dsn = ( - f"postgresql+asyncpg://{postgresql_proc.user}@" - f"{postgresql_proc.host}:{postgresql_proc.port}/postgres" - ) - engine = create_async_engine(dsn, isolation_level="AUTOCOMMIT") - async with engine.begin() as conn: - await conn.execute(text("DROP SCHEMA IF EXISTS bm25_fts CASCADE")) - await conn.execute(text("CREATE SCHEMA bm25_fts")) - await conn.execute(text("SET search_path TO bm25_fts")) - for statement in filter(None, (s.strip() for s in _SCHEMA.split(";"))): - await conn.execute(text(statement)) - await conn.execute(text("INSERT INTO job_results VALUES (1, 'job1')")) - await conn.execute( - text( - "INSERT INTO documents VALUES " - "('d1', 'u1', 'ns1', 'active', 1, 'sample.pdf')" - ) - ) - await conn.execute(text("INSERT INTO document_sections VALUES ('s1', '/root')")) - await conn.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 " - "('hit-en', 'd1', 's1', 'text', 'body', 1, 1, " - " 'alpha beta gamma', 'invoices alpha'), " - "('hit-cjk', 'd1', 's1', 'text', 'body', 1, 2, " - " '合同 条款 甲方', '合同 目录')" - ) - ) - await conn.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) " - "SELECT 'noise-' || i, 'd1', 's1', 'text', 'body', 1, i + 10, " - " 'filler unrelated wording ' || i, 'misc path ' || i " - "FROM generate_series(1, :noise) AS i" - ), - {"noise": _NOISE_ROWS}, - ) - - session_factory = async_sessionmaker(engine, expire_on_commit=False) - async with session_factory() as session: - await session.execute(text("SET search_path TO bm25_fts")) - yield session - await engine.dispose() - - -async def _content_hits(session: AsyncSession, query: str) -> list[str]: - rows = await content_channel( - session, - user_id="u1", - namespace="ns1", - query=query, - top_k=50, - exclude_document_ids=[], - exclude_sections=[], - ) - return [str(row["chunk_id"]) for row in rows] - - -@pytest.mark.asyncio -async def test_content_channel_returns_only_query_matching_chunks( - seeded_session: AsyncSession, -) -> None: - # The corpus holds hundreds of unrelated chunks. Before the prefilter every - # one of them was loaded into Python for BM25 scoring. - assert await _content_hits(seeded_session, "alpha") == ["hit-en"] - - -@pytest.mark.asyncio -async def test_content_channel_matches_cjk_tokens( - seeded_session: AsyncSession, -) -> None: - assert await _content_hits(seeded_session, "合同") == ["hit-cjk"] - - -@pytest.mark.asyncio -async def test_content_channel_uses_or_semantics_across_tokens( - seeded_session: AsyncSession, -) -> None: - # A row matching any single query token must survive, matching how the - # Python BM25 ranker admits rows. - hits = await _content_hits(seeded_session, "alpha 合同") - assert sorted(hits) == ["hit-cjk", "hit-en"] - - -@pytest.mark.asyncio -async def test_tsquery_operators_in_query_do_not_change_filter_shape( - seeded_session: AsyncSession, -) -> None: - # Tokens are lexed by Postgres as data. If operators leaked into tsquery - # syntax, "alpha & zzzz" would AND and drop the row. - assert await _content_hits(seeded_session, "alpha & zzzz") == ["hit-en"] - assert await _content_hits(seeded_session, "!alpha") == ["hit-en"] - - -@pytest.mark.asyncio -async def test_query_matching_nothing_returns_no_rows( - seeded_session: AsyncSession, -) -> None: - # The fallback re-runs the unfiltered scan, and BM25 then scores no row - # above zero, so the channel still yields nothing. - assert await _content_hits(seeded_session, "zzzznomatch") == [] - - -@pytest.mark.asyncio -async def test_path_channel_prefilters_on_path_search_tsv( - seeded_session: AsyncSession, -) -> None: - rows = await path_channel( - seeded_session, - user_id="u1", - namespace="ns1", - query="invoices", - top_k=50, - exclude_document_ids=[], - exclude_sections=[], - ) - assert [str(row["chunk_id"]) for row in rows] == ["hit-en"] - - -@pytest.mark.asyncio -async def test_exclusions_still_apply_under_the_prefilter( - seeded_session: AsyncSession, -) -> None: - rows = await content_channel( - seeded_session, - user_id="u1", - namespace="ns1", - query="alpha", - top_k=50, - exclude_document_ids=["d1"], - 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_retrieval_classic_map_unit_contract.py b/apps/api/tests/contract/test_retrieval_classic_map_unit_contract.py new file mode 100644 index 000000000..e20e647f9 --- /dev/null +++ b/apps/api/tests/contract/test_retrieval_classic_map_unit_contract.py @@ -0,0 +1,282 @@ +from __future__ import annotations + +from collections.abc import Callable +from contextlib import AbstractAsyncContextManager +from typing import Any, cast +from uuid import uuid4 + +from httpx import AsyncClient +from sqlalchemy import select + +from shared.models.database.document import DocumentMapUnit +from shared.services.retrieval.publication_content import ( + replace_document_revision_content, +) +from shared.services.retrieval.publication_models import DocumentPublicationScope +from shared.services.retrieval.serving_generation import lock_namespace_generation +from tests.support.contract_database import ContractDatabase +from tests.support.retrieval_snapshot_support import contract_db_session + +_USER_ID = "local-dev-user" + + +async def test_classic_route_maps_winning_unit_to_one_chunk( + developer_api_client_factory: Callable[ + [], AbstractAsyncContextManager[AsyncClient] + ], +) -> None: + identifier = uuid4().hex[:8] + namespace = f"classic-map-{identifier}" + async with developer_api_client_factory() as api_client: + first = await _publish_document( + namespace=namespace, + source_file_name="first.pdf", + chunks=[ + { + "chunk_id": f"hit-{identifier}", + "type": "text", + "content": "unique alpha ranking marker", + "path": "first.pdf/Root/Hit/body", + "order": 1, + "metadata": {}, + }, + { + "chunk_id": f"other-{identifier}", + "type": "text", + "content": "unrelated filler paragraph", + "path": "first.pdf/Root/Other/body", + "order": 2, + "metadata": {}, + }, + ], + ) + await _publish_document( + namespace=namespace, + source_file_name="second.pdf", + chunks=[ + { + "chunk_id": f"noise-{identifier}", + "type": "text", + "content": "more unrelated filler", + "path": "second.pdf/Root/Noise/body", + "order": 1, + "metadata": {}, + }, + ], + ) + response = await api_client.post( + "/api/v1/retrieval/query", + json={ + "namespace": namespace, + "query": "unique alpha ranking", + "top_k": 1, + "use_agentic": False, + }, + ) + + assert response.status_code == 200 + body = cast(dict[str, object], response.json()) + results = cast(list[dict[str, object]], body["results"]) + assert body["router_used"] == "classic_topk" + assert len(results) == 1 + assert results[0]["chunk_id"] == f"hit-{identifier}" + assert results[0]["chunk_type"] == "text" + assert results[0]["source"] == { + "document_id": first["document_id"], + "source_file_name": "first.pdf", + "section_path": "Root / Hit / body", + } + + +async def test_classic_route_image_filter_scores_only_units_with_images( + developer_api_client_factory: Callable[ + [], AbstractAsyncContextManager[AsyncClient] + ], +) -> None: + identifier = uuid4().hex[:8] + namespace = f"classic-image-{identifier}" + async with developer_api_client_factory() as api_client: + with_image = await _publish_document( + namespace=namespace, + source_file_name="with-image.pdf", + chunks=[ + { + "chunk_id": f"body-{identifier}", + "type": "text", + "content": "shared alpha marker next to a chart", + "path": "with-image.pdf/Root/Section/body", + "order": 1, + "metadata": {"connect_to": [{"target": f"chart-{identifier}"}]}, + }, + { + "chunk_id": f"chart-{identifier}", + "type": "image", + "content": "chart of shared alpha marker", + "path": "images/chart.png", + "order": 2, + "file_path": "images/chart.png", + "metadata": {}, + }, + ], + ) + await _publish_document( + namespace=namespace, + source_file_name="text-only.pdf", + chunks=[ + { + "chunk_id": f"text-only-{identifier}", + "type": "text", + "content": "shared alpha marker with no chart", + "path": "text-only.pdf/Root/Section/body", + "order": 1, + "metadata": {}, + }, + ], + ) + await _publish_document( + namespace=namespace, + source_file_name="other-image.pdf", + chunks=[ + { + "chunk_id": f"other-body-{identifier}", + "type": "text", + "content": "unrelated filler next to a photo", + "path": "other-image.pdf/Root/Section/body", + "order": 1, + "metadata": {"connect_to": [{"target": f"photo-{identifier}"}]}, + }, + { + "chunk_id": f"photo-{identifier}", + "type": "image", + "content": "unrelated landscape photo", + "path": "images/photo.png", + "order": 2, + "file_path": "images/photo.png", + "metadata": {}, + }, + ], + ) + async with contract_db_session() as db: + image_units = list( + ( + await db.execute( + select(DocumentMapUnit).where( + DocumentMapUnit.document_id == with_image["document_id"] + ) + ) + ).scalars() + ) + assert any(unit.has_image for unit in image_units) + + response = await api_client.post( + "/api/v1/retrieval/query", + json={ + "namespace": namespace, + "query": "shared alpha marker", + "top_k": 1, + "use_agentic": False, + "chunk_types": ["image"], + }, + ) + + assert response.status_code == 200 + body = cast(dict[str, object], response.json()) + results = cast(list[dict[str, object]], body["results"]) + assert body["router_used"] == "classic_topk" + assert len(results) == 1 + assert results[0]["chunk_id"] == f"chart-{identifier}" + assert results[0]["chunk_type"] == "image" + + +def _publish_revision_with_generation_lock( + sync_db: Any, + *, + scope: DocumentPublicationScope, + chunks: list[dict[str, Any]], +) -> None: + lock_namespace_generation( + sync_db, user_id=scope.user_id, namespace=scope.namespace + ) + replace_document_revision_content(sync_db, scope=scope, chunks=chunks) + + +async def _publish_document( + *, + namespace: str, + source_file_name: str, + chunks: list[dict[str, Any]], +) -> dict[str, str]: + document_id = f"doc_{uuid4().hex[:12]}" + job_id = f"job_{uuid4().hex[:12]}" + job_result_id = f"result_{uuid4().hex[:12]}" + 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', + :source_file_name, 'chunk', NOW(), NOW() + ) + """, + { + "document_id": document_id, + "user_id": _USER_ID, + "namespace": namespace, + "source_file_name": source_file_name, + }, + ) + 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}, + ) + scope = DocumentPublicationScope( + user_id=_USER_ID, + namespace=namespace, + document_id=document_id, + job_result_id=job_result_id, + source_file_name=source_file_name, + ) + async with contract_db_session() as db: + await db.run_sync( + lambda sync_db: _publish_revision_with_generation_lock( + sync_db, + scope=scope, + chunks=chunks, + ) + ) + await db.commit() + return { + "document_id": document_id, + "job_id": job_id, + "job_result_id": job_result_id, + } 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 6d30731cf..6665eb3e2 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 @@ -356,6 +356,19 @@ async def test_lazy_snapshot_defers_selected_asset_reference_metadata( ) await db.commit() + async with contract_db_session() as db: + persisted_units = list( + ( + await db.execute( + select(DocumentMapUnit).where( + DocumentMapUnit.document_id == document_id + ) + ) + ).scalars() + ) + assert any(unit.has_image for unit in persisted_units) + assert all(not unit.has_table for unit in persisted_units) + calls: list[tuple[str, str]] = [] original = ReadOnlyChunkStore.load_chunk_reference_metadata diff --git a/apps/api/tests/contract/test_retrieval_revision_races_contract.py b/apps/api/tests/contract/test_retrieval_revision_races_contract.py index 3ebb02326..43c47f27d 100644 --- a/apps/api/tests/contract/test_retrieval_revision_races_contract.py +++ b/apps/api/tests/contract/test_retrieval_revision_races_contract.py @@ -1,11 +1,8 @@ -"""Deterministic contracts for revision and channel-session coherence.""" +"""Deterministic contracts for revision-generation coherence.""" from __future__ import annotations -import importlib -from collections.abc import AsyncGenerator -from contextlib import asynccontextmanager -from typing import Any, cast +from typing import cast import pytest from sqlalchemy.ext.asyncio import AsyncSession @@ -14,62 +11,6 @@ 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: diff --git a/apps/api/tests/unit/test_bm25_channel_tsquery.py b/apps/api/tests/unit/test_bm25_channel_tsquery.py deleted file mode 100644 index aea2d9dfa..000000000 --- a/apps/api/tests/unit/test_bm25_channel_tsquery.py +++ /dev/null @@ -1,45 +0,0 @@ -"""Unit tests for BM25 channel Postgres FTS prefilter token preparation.""" - -from __future__ import annotations - -from shared.services.retrieval.search.channels import ( - _MAX_FTS_QUERY_TOKENS, - _prepare_fts_tokens, -) - - -def test_returns_empty_for_no_tokens() -> None: - assert _prepare_fts_tokens([]) == [] - - -def test_keeps_token_order() -> None: - assert _prepare_fts_tokens(["alpha", "beta"]) == ["alpha", "beta"] - - -def test_strips_surrounding_whitespace() -> None: - assert _prepare_fts_tokens([" alpha ", "beta"]) == ["alpha", "beta"] - - -def test_drops_blank_tokens() -> None: - assert _prepare_fts_tokens(["", " ", "alpha"]) == ["alpha"] - - -def test_returns_empty_when_every_token_is_blank() -> None: - assert _prepare_fts_tokens(["", " "]) == [] - - -def test_caps_token_count() -> None: - tokens = [f"tok{index}" for index in range(_MAX_FTS_QUERY_TOKENS + 25)] - assert len(_prepare_fts_tokens(tokens)) == _MAX_FTS_QUERY_TOKENS - - -def test_preserves_cjk_tokens() -> None: - assert _prepare_fts_tokens(["合同", "条款"]) == ["合同", "条款"] - - -def test_passes_tsquery_operators_through_untouched() -> None: - # Tokens travel to Postgres as a text[] parameter and are lexed there, so - # operator characters are data rather than syntax. Nothing is escaped or - # dropped here. - raw = ["alpha' & 'zzzz", "!beta", "a|b"] - assert _prepare_fts_tokens(raw) == raw diff --git a/packages/shared-python/shared/models/database/document.py b/packages/shared-python/shared/models/database/document.py index b85dd84f3..2ebce6cd9 100644 --- a/packages/shared-python/shared/models/database/document.py +++ b/packages/shared-python/shared/models/database/document.py @@ -270,6 +270,12 @@ class DocumentMapUnit(Base): 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) + # Asset presence under this unit's section, after root-asset remount + # (``KnowhereProvider._remount_root_assets``). Lets type-scoped queries + # (e.g. chunk_types=["image"]) narrow map-unit candidates *before* + # scoring instead of scoring everything and discarding after the fact. + has_image: Mapped[bool] = mapped_column(nullable=False, default=False) + has_table: Mapped[bool] = mapped_column(nullable=False, default=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 @@ -284,6 +290,18 @@ class DocumentMapUnit(Base): "unit_id", ), Index("idx_document_map_units_section", "section_id"), + Index( + "idx_document_map_units_has_image", + "document_id", + "job_result_id", + postgresql_where=has_image.is_(True), + ), + Index( + "idx_document_map_units_has_table", + "document_id", + "job_result_id", + postgresql_where=has_table.is_(True), + ), ) diff --git a/packages/shared-python/shared/services/retrieval/execution/routes.py b/packages/shared-python/shared/services/retrieval/execution/routes.py index fe47be194..673e2b6ed 100644 --- a/packages/shared-python/shared/services/retrieval/execution/routes.py +++ b/packages/shared-python/shared/services/retrieval/execution/routes.py @@ -7,7 +7,7 @@ from loguru import logger from sqlalchemy.ext.asyncio import AsyncSession -from shared.services.retrieval.search.discovery import bottom_discovery +from shared.services.retrieval.search.map_unit_discovery import map_unit_discovery from shared.services.retrieval.execution.reference_resolver import ( resolve_workflow_references, ) @@ -118,7 +118,7 @@ async def _try_run_small_corpus_route( async def _run_classic_topk_route( context: RetrievalRouteContext, ) -> RetrievalRouteOutcome: - discovery_result = await bottom_discovery( + discovery_result = await map_unit_discovery( context.db, user_id=context.user_id, namespace=context.namespace, @@ -129,9 +129,6 @@ async def _run_classic_topk_route( chunk_types=context.allowed_chunk_types, signal_paths=context.signal_paths, filter_mode=context.filter_mode, - channels=context.channels, - channel_weights=context.channel_weights, - internal_recall_k=context.internal_recall_k, revision_pins=context.revision_pins, ) diff --git a/packages/shared-python/shared/services/retrieval/map_unit_index.py b/packages/shared-python/shared/services/retrieval/map_unit_index.py index 15454dc9d..4171861c7 100644 --- a/packages/shared-python/shared/services/retrieval/map_unit_index.py +++ b/packages/shared-python/shared/services/retrieval/map_unit_index.py @@ -93,6 +93,11 @@ def replace_document_map_units( 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() + # ``provider.self_units`` already reflects root-asset remount (assets + # referenced via ``connect_to`` are moved onto the text section that + # embeds them), so this is the same ownership the query-time scorer + # sees, not a new computation. + section_chunk_types = {u.chunk_type for u in provider.self_units(section_id)} db.add( DocumentMapUnit( id=map_unit_id, @@ -104,6 +109,8 @@ def replace_document_map_units( path_token_count=len(path_tokens), content_token_count=len(content_tokens), term_search_text_lower=str(unit.get("term_search_text") or "").lower(), + has_image="image" in section_chunk_types, + has_table="table" in section_chunk_types, sort_order=sort_order, ) ) 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 3c237aef7..a937210d6 100644 --- a/packages/shared-python/shared/services/retrieval/nav/knowhere_hybrid.py +++ b/packages/shared-python/shared/services/retrieval/nav/knowhere_hybrid.py @@ -1,7 +1,6 @@ """KnowWhere-style 3-channel hybrid retrieval (path BM25 + content BM25 + term). -Ported from Ontos-AI/knowhere: - packages/shared-python/shared/services/retrieval/search/{scoring,lexical_ranker,channels}.py +Ported from Ontos-AI/knowhere map-unit BM25 (path + content). Reference: https://github.com/Ontos-AI/knowhere """ diff --git a/packages/shared-python/shared/services/retrieval/search/channels.py b/packages/shared-python/shared/services/retrieval/search/channels.py deleted file mode 100644 index cdf67326f..000000000 --- a/packages/shared-python/shared/services/retrieval/search/channels.py +++ /dev/null @@ -1,529 +0,0 @@ -""" -Independent retrieval channels for checkerboard search. - -Each channel queries the full scoped corpus independently and returns -ranked rows. Channels are fused via RRF in the orchestrator. -""" - -from __future__ import annotations - -import time -from collections.abc import Mapping -from typing import Any - -from loguru import logger -from sqlalchemy import text -from sqlalchemy.ext.asyncio import AsyncSession - -from shared.core.config import settings -from shared.services.retrieval.search.lexical_ranker import ( - rank_rows_by_bm25, - tokenize_query_for_ranker, -) -from shared.services.retrieval.search.section_filters import is_excluded_section - -# The generated tsvector columns are built with to_tsvector('simple', ...), so -# queries must use the same configuration or nothing matches. -_FTS_CONFIG = "simple" - -# Guards against pathological queries producing an enormous tsquery. -_MAX_FTS_QUERY_TOKENS = 50 - -_TSV_FIELD_BY_SEARCH_FIELD = { - "content_search_text": "content_search_tsv", - "path_search_text": "path_search_tsv", -} - - -_SCOPED_CORPUS_CTE = """ -WITH scoped_chunks AS ( - SELECT - dc.id, - dc.chunk_id, - dc.document_id, - dc.section_id, - dc.chunk_type, - dc.content, - dc.source_chunk_path, - dc.file_path, - dc.chunk_metadata, - dc.job_result_id, - dc.sort_order, - dc.content_search_text, - dc.content_search_tsv, - dc.path_search_text, - dc.path_search_tsv, - dc.term_search_text, - d.source_file_name, - d.user_id, - d.namespace, - ds.section_path, - jr.job_id - FROM document_chunks dc - JOIN documents d - ON d.document_id = dc.document_id - {revision_join} - LEFT JOIN document_sections ds - ON ds.section_id = dc.section_id - JOIN job_results jr - ON jr.id = dc.job_result_id - 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 "" - # Use PostgreSQL array ANY() to avoid asyncpg tuple-binding pitfalls with - # raw text() + `NOT IN :param` (which asyncpg treats as a record parameter - # and fails with a syntax error). - return "AND d.document_id <> ALL(:excluded_doc_ids)" - - -def _build_base_params( - *, - user_id: str, - namespace: str, - exclude_document_ids: list[str], -) -> dict[str, Any]: - params: dict[str, Any] = { - "user_id": user_id, - "namespace": namespace, - } - if exclude_document_ids: - params["excluded_doc_ids"] = list(exclude_document_ids) - return params - - -def _build_extra_filters( - *, - allowed_chunk_types: set[str] | None, - signal_paths: list[str], - filter_mode: str, -) -> tuple[str, dict[str, Any]]: - """Build additional SQL WHERE clauses for chunk_types and signal_path filtering.""" - clauses: list[str] = [] - params: dict[str, Any] = {} - - if allowed_chunk_types is not None: - placeholders = ", ".join(f":_act_{i}" for i in range(len(allowed_chunk_types))) - clauses.append(f"AND LOWER(dc.chunk_type) IN ({placeholders})") - for i, ct in enumerate(sorted(allowed_chunk_types)): - params[f"_act_{i}"] = ct - - if signal_paths: - # TODO(intent-step): Current implementation uses OR across - # signal_paths keywords. The Intent Step will need hierarchical - # AND (prefix) matching, e.g. signal_paths=["第一章/1.1/（2）"] - # should match only paths containing ALL segments in order. - # Consider adding a `filter_strategy` param: "keyword_or" (current) - # vs "path_prefix" (for Intent Step resolved paths). - ilike_parts = [] - for i, kw in enumerate(signal_paths): - key = f"_sig_{i}" - ilike_parts.append(f"LOWER(COALESCE(ds.section_path, '')) LIKE :{key}") - params[key] = f"%{kw.lower()}%" - combined = " OR ".join(ilike_parts) - if filter_mode == "keep": - clauses.append(f"AND ({combined})") - else: - clauses.append(f"AND NOT ({combined})") - - return "\n ".join(clauses), params - - -def _build_exclude_section_filters( - *, - exclude_sections: list[dict[str, str]], -) -> tuple[str, dict[str, Any]]: - """Exclude an exact section path and its descendants inside the scoped CTE. - - Applied before the FTS candidate LIMIT so excluded sections cannot consume - the bounded candidate budget. Sectionless chunks stay eligible (empty path - does not match), matching ``is_excluded_section``. - """ - clauses: list[str] = [] - params: dict[str, Any] = {} - - for index, item in enumerate(exclude_sections): - if not isinstance(item, dict): - continue - document_id = str(item.get("document_id") or "").strip() - section_path = str(item.get("section_path") or "").strip() - if not document_id or not section_path: - continue - - document_key = f"_exc_section_doc_{index}" - path_key = f"_exc_section_path_{index}" - clauses.append( - f"""AND NOT ( - dc.document_id = :{document_key} - AND ( - COALESCE(ds.section_path, '') = :{path_key} - OR POSITION(:{path_key} || ' / ' IN COALESCE(ds.section_path, '')) = 1 - ) - )""" - ) - params[document_key] = document_id - params[path_key] = section_path - - return "\n ".join(clauses), params - - -def _join_sql_filters(*filters: str) -> str: - return "\n ".join(filter(None, filters)) - - -def _prepare_fts_tokens(tokens: list[str]) -> list[str]: - """Return the ranker tokens to hand to the Postgres FTS prefilter. - - Tokens are passed to SQL as a text[] parameter and lexed by Postgres - itself, so nothing here needs to escape tsquery syntax. Only emptiness and - an upper bound are enforced. - """ - prepared: list[str] = [] - for token in tokens: - cleaned = token.strip() - if cleaned: - prepared.append(cleaned) - if len(prepared) >= _MAX_FTS_QUERY_TOKENS: - break - return prepared - - -def _row_to_dict(row: Any) -> dict[str, Any]: - return dict(row._mapping) - - -def _filter_excluded_sections( - rows: list[dict[str, Any]], - exclude_sections: list[dict[str, str]], -) -> list[dict[str, Any]]: - if not exclude_sections: - return rows - return [ - row - for row in rows - if not is_excluded_section( - document_id=row.get("document_id"), - section_path=row.get("section_path"), - exclude_sections=exclude_sections, - ) - ] - - -async def path_channel( - db: AsyncSession, - *, - user_id: str, - namespace: str, - query: str, - top_k: int, - exclude_document_ids: list[str], - exclude_sections: list[dict[str, str]], - 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. - - This keeps the channel useful when vector search is unavailable. A future - vector score can be fused on top of the returned BM25 score. - """ - return await _bm25_channel( - db, - user_id=user_id, - namespace=namespace, - query=query, - top_k=top_k, - exclude_document_ids=exclude_document_ids, - exclude_sections=exclude_sections, - allowed_chunk_types=allowed_chunk_types, - signal_paths=signal_paths, - filter_mode=filter_mode, - search_field="path_search_text", - revision_pins=revision_pins, - ) - - -async def content_channel( - db: AsyncSession, - *, - user_id: str, - namespace: str, - query: str, - top_k: int, - exclude_document_ids: list[str], - exclude_sections: list[dict[str, str]], - 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( - db, - user_id=user_id, - namespace=namespace, - query=query, - top_k=top_k, - exclude_document_ids=exclude_document_ids, - exclude_sections=exclude_sections, - allowed_chunk_types=allowed_chunk_types, - signal_paths=signal_paths, - filter_mode=filter_mode, - search_field="content_search_text", - revision_pins=revision_pins, - ) - - -async def _bm25_channel( - db: AsyncSession, - *, - user_id: str, - namespace: str, - query: str, - top_k: int, - exclude_document_ids: list[str], - exclude_sections: list[dict[str, str]], - allowed_chunk_types: set[str] | None, - 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}") - - query_tokens = tokenize_query_for_ranker(query) - if not query_tokens: - return [] - - exclude_clause = _build_exclude_clause(exclude_document_ids) - extra_sql, extra_params = _build_extra_filters( - allowed_chunk_types=allowed_chunk_types, - signal_paths=signal_paths or [], - filter_mode=filter_mode, - ) - section_sql, section_params = _build_exclude_section_filters( - exclude_sections=exclude_sections, - ) - extra_sql = _join_sql_filters(extra_sql, section_sql) - params = _build_base_params( - user_id=user_id, - namespace=namespace, - exclude_document_ids=exclude_document_ids, - ) - 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, - ) - full_scan_sql = ( - corpus_cte - + f""" - SELECT sc.* - FROM scoped_chunks sc - WHERE COALESCE(sc.{search_field}, '') <> '' - """ - ) - - started_at = time.perf_counter() - tsv_field = _TSV_FIELD_BY_SEARCH_FIELD[search_field] - fts_tokens = _prepare_fts_tokens(query_tokens) - candidate_limit = int(settings.RETRIEVAL_POSTGRES_FTS_CANDIDATE_LIMIT) - - rows: list[dict[str, Any]] = [] - used_fallback = True - if fts_tokens: - # Postgres lexes the tokens with the same configuration that generated - # the tsvector columns, then ORs the resulting lexemes. Building the - # tsquery server-side keeps the prefilter aligned with the stored - # lexicon and leaves no room for tsquery syntax in user input to - # change the query shape. `fts_query.q` is NULL when no token yields a - # lexeme, which the caller treats as "no usable prefilter". - prefilter_sql = ( - corpus_cte - + f""", - fts_query AS ( - SELECT string_agg(quote_literal(lexeme), ' | ')::tsquery AS q - FROM ( - SELECT DISTINCT - unnest(tsvector_to_array(to_tsvector('{_FTS_CONFIG}', token))) AS lexeme - FROM unnest(CAST(:fts_tokens AS text[])) AS token - ) lexemes - ) - SELECT sc.* - FROM scoped_chunks sc, fts_query fq - WHERE COALESCE(sc.{search_field}, '') <> '' - AND fq.q IS NOT NULL - AND sc.{tsv_field} @@ fq.q - ORDER BY ts_rank_cd(sc.{tsv_field}, fq.q) DESC - LIMIT :fts_candidate_limit - """ - ) - prefilter_params = dict(params) - prefilter_params["fts_tokens"] = fts_tokens - prefilter_params["fts_candidate_limit"] = candidate_limit - result = await db.execute(text(prefilter_sql), prefilter_params) - rows = [_row_to_dict(r) for r in result.all()] - used_fallback = not rows - - # No usable tsquery, or the prefilter matched nothing. Fall back to the - # full scoped scan so recall never regresses against the previous - # behaviour. - if used_fallback: - result = await db.execute(text(full_scan_sql), params) - rows = [_row_to_dict(r) for r in result.all()] - - candidate_count = len(rows) - # Defensive: SQL already owns pre-LIMIT section exclusion. - rows = _filter_excluded_sections(rows, exclude_sections) - - ranked_rows = rank_rows_by_bm25(rows, query_tokens, search_field=search_field) - ranked_rows = ranked_rows[:top_k] - - logger.debug( - "bm25_channel field={} candidates={} limit={} ranked={} fallback={} duration_ms={:.1f}", - search_field, - candidate_count, - candidate_limit, - len(ranked_rows), - used_fallback, - (time.perf_counter() - started_at) * 1000, - ) - return ranked_rows - - -async def term_channel( - db: AsyncSession, - *, - user_id: str, - namespace: str, - query: str, - top_k: int, - exclude_document_ids: list[str], - exclude_sections: list[dict[str, str]], - 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. - - Aligned with KB checkerboard_find() term channel (grep_search()): - exact substring match on content + path, scoring by hit count. - - Note: top_k is already effective_recall_k from app_service. - """ - query_lower = query.lower().strip() - query_tokens = tokenize_query_for_ranker(query) - if not query_lower or not query_tokens: - return [] - - exclude_clause = _build_exclude_clause(exclude_document_ids) - extra_sql, extra_params = _build_extra_filters( - allowed_chunk_types=allowed_chunk_types, - signal_paths=signal_paths or [], - filter_mode=filter_mode, - ) - params = _build_base_params( - user_id=user_id, - namespace=namespace, - exclude_document_ids=exclude_document_ids, - ) - 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}" - ilike_conditions.append(f"LOWER(sc.term_search_text) LIKE :{param_key}") - params[param_key] = f"%{unit}%" - - if not ilike_conditions: - ilike_conditions.append("LOWER(sc.term_search_text) LIKE :full_query") - params["full_query"] = f"%{query_lower}%" - - 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""" - SELECT sc.* - FROM scoped_chunks sc - WHERE sc.term_search_text IS NOT NULL - AND ({where_clause}) - """ - ) - - result = await db.execute(text(sql), params) - rows = [_row_to_dict(r) for r in result.all()] - rows = _filter_excluded_sections(rows, exclude_sections) - - scored: list[dict[str, Any]] = [] - for row in rows: - haystack = (row.get("term_search_text") or "").lower() - if query_lower in haystack: - row["score"] = 100.0 - scored.append(row) - else: - hit_count = sum(1 for u in query_tokens if u in haystack) - if hit_count > 0: - row["score"] = float(hit_count) - scored.append(row) - - scored.sort(key=lambda r: r["score"], reverse=True) - return scored[:top_k] diff --git a/packages/shared-python/shared/services/retrieval/search/discovery.py b/packages/shared-python/shared/services/retrieval/search/discovery.py deleted file mode 100644 index d71ccdd94..000000000 --- a/packages/shared-python/shared/services/retrieval/search/discovery.py +++ /dev/null @@ -1,216 +0,0 @@ -"""Classic 3-channel bottom discovery (moved from agentic/discovery/tools). - -Input (all keyword, via classic route context): - db, user_id, namespace, query, top_k, - exclude_document_ids, exclude_sections, - chunk_types?, signal_paths?, filter_mode, - channels?, channel_weights?, internal_recall_k? - -Output: - DiscoveryResult - status: ``discovery_done`` | ``error`` - payload.fused_rows: RRF-merged scored rows (classic reads this) - payload.top_doc_ids / channel_counts: diagnostics - latency_ms, error? - -Does not call LLM. Does not touch map-nav / wallet / agentic budget. -""" - -from __future__ import annotations - -import asyncio -import time -from collections.abc import Awaitable, Callable, Mapping -from dataclasses import dataclass, field -from typing import Any - -from loguru import logger -from sqlalchemy.ext.asyncio import AsyncSession - -from shared.services.retrieval.search.channels import ( - content_channel, - path_channel, - term_channel, -) -from shared.services.retrieval.search.scoring import ( - merge_channels_rrf, - merge_same_section_rows, - normalize_row_scores, -) -from shared.services.retrieval.settings import ( - CHANNEL_WEIGHT_CONTENT, - CHANNEL_WEIGHT_PATH, - CHANNEL_WEIGHT_TERM, - INTERNAL_RECALL_K_MULTIPLIER, -) - - -@dataclass -class DiscoveryResult: - """Return shape for classic bottom discovery (replaces agentic ToolResult).""" - - status: str - payload: dict[str, Any] = field(default_factory=dict) - latency_ms: int = 0 - error: str | None = None - - -async def bottom_discovery( - db: AsyncSession, - *, - user_id: str, - namespace: str, - query: str, - top_k: int, - exclude_document_ids: list[str], - exclude_sections: list[dict[str, str]], - chunk_types: set[str] | None = None, - signal_paths: list[str] | None = None, - filter_mode: str = "delete", - 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 - if internal_recall_k is not None - else top_k * INTERNAL_RECALL_K_MULTIPLIER - ) - 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, - user_id=user_id, - namespace=namespace, - query=query, - top_k=effective_recall_k, - exclude_document_ids=exclude_document_ids, - exclude_sections=exclude_sections, - 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, - user_id=user_id, - namespace=namespace, - query=query, - top_k=effective_recall_k, - exclude_document_ids=exclude_document_ids, - exclude_sections=exclude_sections, - 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, - user_id=user_id, - namespace=namespace, - query=query, - top_k=effective_recall_k, - exclude_document_ids=exclude_document_ids, - exclude_sections=exclude_sections, - allowed_chunk_types=allowed_chunk_types, - signal_paths=signal_paths, - filter_mode=filter_mode, - revision_pins=revision_pins, - ), - ) - - default_weights = { - "path": CHANNEL_WEIGHT_PATH, - "content": CHANNEL_WEIGHT_CONTENT, - "term": CHANNEL_WEIGHT_TERM, - } - effective_weights = {**default_weights, **(channel_weights or {})} - - channel_lists: list[list[dict[str, Any]]] = [] - weight_list: list[float] = [] - if path_rows: - channel_lists.append(path_rows) - weight_list.append(effective_weights.get("path", CHANNEL_WEIGHT_PATH)) - if content_rows: - channel_lists.append(content_rows) - weight_list.append(effective_weights.get("content", CHANNEL_WEIGHT_CONTENT)) - if term_rows: - channel_lists.append(term_rows) - weight_list.append(effective_weights.get("term", CHANNEL_WEIGHT_TERM)) - - fused_rows = ( - merge_channels_rrf(channel_lists, weight_list, effective_recall_k) - if channel_lists - else [] - ) - fused_rows = merge_same_section_rows(fused_rows) - - if fused_rows: - normalize_row_scores( - fused_rows, - source_field="score", - target_field="discovery_score", - default=0.5, - ) - - doc_id_counts: dict[str, int] = {} - for row in fused_rows: - did = row.get("document_id", "") - if did: - doc_id_counts[did] = doc_id_counts.get(did, 0) + 1 - top_doc_ids = sorted( - doc_id_counts, - key=lambda document_id: doc_id_counts[document_id], - reverse=True, - )[:5] - - latency = int((time.monotonic() - t0) * 1000) - logger.info( - f" search.bottom_discovery: {len(fused_rows)} fused rows, " - f"top_doc_ids={top_doc_ids}, {latency}ms" - ) - return DiscoveryResult( - status="discovery_done", - payload={ - "fused_rows": fused_rows, - "top_doc_ids": top_doc_ids, - "channel_counts": { - "path": len(path_rows), - "content": len(content_rows), - "term": len(term_rows), - }, - }, - latency_ms=latency, - ) - except Exception as exc: - 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/lexical_ranker.py b/packages/shared-python/shared/services/retrieval/search/lexical_ranker.py deleted file mode 100644 index 8b50cd404..000000000 --- a/packages/shared-python/shared/services/retrieval/search/lexical_ranker.py +++ /dev/null @@ -1,74 +0,0 @@ -from __future__ import annotations - -from typing import Any - -from loguru import logger - -from shared.utils.text_utils import tokenize_for_retrieval - - -def tokenize_query_for_ranker(query: str) -> list[str]: - return tokenize_for_retrieval(query, dedupe=True) - - -def rank_rows_by_bm25( - rows: list[dict[str, Any]], - query_tokens: list[str], - *, - search_field: str, -) -> list[dict[str, Any]]: - """Rank matching rows with BM25 over pre-tokenized search text.""" - try: - from rank_bm25 import BM25Okapi - except ImportError: - return _rank_rows_by_token_overlap( - rows, - query_tokens, - search_field=search_field, - ) - - corpus: list[list[str]] = [] - ranked_rows: list[dict[str, Any]] = [] - query_token_set = set(query_tokens) - for row in rows: - tokens = _get_search_tokens(row, search_field=search_field) - if not tokens or not query_token_set.intersection(tokens): - continue - corpus.append(tokens) - ranked_rows.append(row) - - if not corpus or not query_tokens: - return [] - - bm25 = BM25Okapi(corpus) - scores = bm25.get_scores(query_tokens) - - for index, row in enumerate(ranked_rows): - row["score"] = float(scores[index]) - - ranked_rows.sort(key=lambda row: row["score"], reverse=True) - return ranked_rows - - -def _rank_rows_by_token_overlap( - rows: list[dict[str, Any]], - query_tokens: list[str], - *, - search_field: str, -) -> list[dict[str, Any]]: - logger.warning("rank_bm25 not installed, skipping BM25 re-rank") - ranked_rows: list[dict[str, Any]] = [] - query_token_set = set(query_tokens) - for row in rows: - tokens = _get_search_tokens(row, search_field=search_field) - overlap = len(query_token_set.intersection(tokens)) - if overlap <= 0: - continue - row["score"] = float(overlap) - ranked_rows.append(row) - ranked_rows.sort(key=lambda row: row["score"], reverse=True) - return ranked_rows - - -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] diff --git a/packages/shared-python/shared/services/retrieval/search/map_unit_discovery.py b/packages/shared-python/shared/services/retrieval/search/map_unit_discovery.py new file mode 100644 index 000000000..b75649fb1 --- /dev/null +++ b/packages/shared-python/shared/services/retrieval/search/map_unit_discovery.py @@ -0,0 +1,500 @@ +"""Classic-route discovery via the persisted map-unit BM25 scorer. + +Replaces the retired chunk-level 3-channel SQL scan. Scoring uses the same +``score_persisted_corpus_many`` formula as map-nav (path + content only). + +``chunk_types`` is optional: omitted means score every in-scope unit. When +the request is image/table only, ``has_image`` / ``has_table`` (written at +index time) narrow candidates before scoring. + +Hydration returns one primary chunk per winning unit. Asset-only requests +then follow ``connect_to`` to that unit's image/table row. Downstream +``assemble_retrieval_results`` still inlines connected assets for text hits. +""" + +from __future__ import annotations + +import json +import time +from collections.abc import Mapping +from dataclasses import dataclass, field +from typing import Any + +from loguru import logger +from sqlalchemy import text +from sqlalchemy.ext.asyncio import AsyncSession + +from shared.services.retrieval.hydration.connected import hydrate_connected_target_rows +from shared.services.retrieval.hydration.row_utils import ( + iter_connected_target_ids, + normalize_chunk_type, +) +from shared.services.retrieval.nav.knowhere_hybrid import ( + PersistedBm25Stats, + PersistedScoreCorpus, + PersistedScoreUnit, + score_persisted_corpus_many, + tokenize_query_for_ranker, +) +from shared.services.retrieval.search.scoring import normalize_row_scores +from shared.services.retrieval.search.section_filters import is_excluded_section +from shared.services.retrieval.settings import ASSET_CHUNK_TYPES + +_MAP_SCORE_CHANNELS = ("path", "content") + +_SCOPED_UNITS_CTE = """ +WITH scoped_units AS ( + SELECT + dmu.id AS map_unit_id, + dmu.document_id, + dmu.job_result_id, + dmu.section_id, + dmu.path_token_count, + dmu.content_token_count, + ds.section_path + FROM document_map_units dmu + JOIN documents d + ON d.document_id = dmu.document_id + {revision_join} + JOIN document_sections ds + ON ds.section_id = dmu.section_id + WHERE d.user_id = :user_id + AND d.namespace = :namespace + AND d.status = 'active' + {revision_clause} + {exclude_clause} + {type_clause} + {signal_clause} +) +""" + + +@dataclass +class DiscoveryResult: + status: str + payload: dict[str, Any] = field(default_factory=dict) + latency_ms: int = 0 + error: str | None = None + + +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 = dmu.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 (dmu.document_id, dmu.job_result_id) IN ({', '.join(placeholders)})", + params, + ) + + +def _build_type_clause( + allowed_chunk_types: set[str] | None, +) -> tuple[str, dict[str, Any]]: + if allowed_chunk_types is None or not allowed_chunk_types.issubset( + ASSET_CHUNK_TYPES + ): + return "", {} + clauses = [] + if "image" in allowed_chunk_types: + clauses.append("dmu.has_image") + if "table" in allowed_chunk_types: + clauses.append("dmu.has_table") + if not clauses: + return "AND FALSE", {} + return f"AND ({' OR '.join(clauses)})", {} + + +def _build_exclude_clause(exclude_document_ids: list[str]) -> tuple[str, dict[str, Any]]: + if not exclude_document_ids: + return "", {} + return "AND d.document_id <> ALL(:excluded_doc_ids)", { + "excluded_doc_ids": list(exclude_document_ids) + } + + +def _build_signal_clause( + signal_paths: list[str], filter_mode: str +) -> tuple[str, dict[str, Any]]: + if not signal_paths: + return "", {} + ilike_parts = [] + params: dict[str, Any] = {} + for index, keyword in enumerate(signal_paths): + key = f"_sig_{index}" + ilike_parts.append(f"LOWER(COALESCE(ds.section_path, '')) LIKE :{key}") + params[key] = f"%{keyword.lower()}%" + combined = " OR ".join(ilike_parts) + clause = f"AND ({combined})" if filter_mode == "keep" else f"AND NOT ({combined})" + return clause, params + + +async def map_unit_discovery( + db: AsyncSession | None, + *, + user_id: str, + namespace: str, + query: str, + top_k: int, + exclude_document_ids: list[str], + exclude_sections: list[dict[str, str]], + chunk_types: set[str] | None = None, + signal_paths: list[str] | None = None, + filter_mode: str = "delete", + revision_pins: Mapping[str, str] | None = None, + **_kwargs: Any, +) -> DiscoveryResult: + """Score the whole in-scope corpus via the persisted map-unit BM25 scorer.""" + t0 = time.monotonic() + try: + if db is None: + return DiscoveryResult( + status="error", + error="database session required", + latency_ms=0, + ) + query_tokens = tokenize_query_for_ranker(query) + if not query_tokens: + return DiscoveryResult(status="discovery_done", payload={"fused_rows": []}) + + revision_join, revision_clause, revision_params = _build_revision_scope( + revision_pins + ) + exclude_clause, exclude_params = _build_exclude_clause(exclude_document_ids) + type_clause, type_params = _build_type_clause(chunk_types) + signal_clause, signal_params = _build_signal_clause( + signal_paths or [], filter_mode + ) + params: dict[str, Any] = {"user_id": user_id, "namespace": namespace} + params.update(revision_params) + params.update(exclude_params) + params.update(type_params) + params.update(signal_params) + + cte = _SCOPED_UNITS_CTE.format( + revision_join=revision_join, + revision_clause=revision_clause, + exclude_clause=exclude_clause, + type_clause=type_clause, + signal_clause=signal_clause, + ) + unit_result = await db.execute(text(cte + "SELECT * FROM scoped_units"), params) + unit_rows = [dict(row._mapping) for row in unit_result.all()] + unit_rows = [ + row + for row in unit_rows + if not is_excluded_section( + document_id=row.get("document_id"), + section_path=row.get("section_path"), + exclude_sections=exclude_sections, + ) + ] + + if not unit_rows: + return DiscoveryResult(status="discovery_done", payload={"fused_rows": []}) + + map_unit_ids = [row["map_unit_id"] for row in unit_rows] + frequency_result = await db.execute( + text( + "SELECT map_unit_id, channel, token, frequency " + "FROM document_map_unit_tokens " + "WHERE map_unit_id = ANY(:unit_ids) " + "AND channel = ANY(:channels) " + "AND token = ANY(:tokens)" + ), + { + "unit_ids": map_unit_ids, + "channels": list(_MAP_SCORE_CHANNELS), + "tokens": query_tokens, + }, + ) + frequencies: dict[tuple[str, str], dict[str, int]] = {} + for map_unit_id, channel, token, frequency in frequency_result.all(): + frequencies.setdefault((str(map_unit_id), str(channel)), {})[str(token)] = ( + int(frequency) + ) + + path_stats = await _build_bm25_stats( + db, + unit_rows=unit_rows, + map_unit_ids=map_unit_ids, + channel="path", + query_tokens=query_tokens, + frequencies=frequencies, + length_field="path_token_count", + ) + content_stats = await _build_bm25_stats( + db, + unit_rows=unit_rows, + map_unit_ids=map_unit_ids, + channel="content", + query_tokens=query_tokens, + frequencies=frequencies, + length_field="content_token_count", + ) + + corpus = PersistedScoreCorpus( + units=[ + PersistedScoreUnit( + unit_id=row["map_unit_id"], + path_length=int(row["path_token_count"]), + content_length=int(row["content_token_count"]), + path_frequencies=frequencies.get((row["map_unit_id"], "path"), {}), + content_frequencies=frequencies.get( + (row["map_unit_id"], "content"), {} + ), + ) + for row in unit_rows + ], + path_stats=path_stats, + content_stats=content_stats, + ) + scores_by_unit = score_persisted_corpus_many(corpus, [query]).get(query, {}) + + rows_by_unit_id = {row["map_unit_id"]: row for row in unit_rows} + + def _unit_has_query_hit(unit_id: str) -> bool: + # BM25 fused score can be 0 when IDF is 0 (tiny corpus). A unit + # that actually holds a query token is still a hit. + return any( + frequencies.get((unit_id, channel), {}).get(token, 0) > 0 + for channel in _MAP_SCORE_CHANNELS + for token in query_tokens + ) + + ranked_unit_ids = sorted( + (unit_id for unit_id in scores_by_unit if _unit_has_query_hit(unit_id)), + key=lambda unit_id: scores_by_unit[unit_id], + reverse=True, + )[:top_k] + + fused_rows = await _hydrate_winning_units( + db, + ranked_unit_ids=ranked_unit_ids, + rows_by_unit_id=rows_by_unit_id, + scores_by_unit=scores_by_unit, + chunk_types=chunk_types, + exclude_document_ids=exclude_document_ids, + exclude_sections=exclude_sections, + revision_pins=revision_pins, + ) + if fused_rows: + normalize_row_scores( + fused_rows, + source_field="score", + target_field="discovery_score", + default=0.5, + ) + + latency = int((time.monotonic() - t0) * 1000) + logger.info( + " search.map_unit_discovery: {} units scored, {} fused rows, {}ms", + len(unit_rows), + len(fused_rows), + latency, + ) + return DiscoveryResult( + status="discovery_done", + payload={"fused_rows": fused_rows}, + latency_ms=latency, + ) + except Exception as exc: + latency = int((time.monotonic() - t0) * 1000) + logger.error(f" search.map_unit_discovery failed: {exc}") + return DiscoveryResult(status="error", error=str(exc), latency_ms=latency) + + +async def _build_bm25_stats( + session: AsyncSession, + *, + unit_rows: list[dict[str, Any]], + map_unit_ids: list[str], + channel: str, + query_tokens: list[str], + frequencies: dict[tuple[str, str], dict[str, int]], + length_field: str, +) -> PersistedBm25Stats: + lengths = [ + int(row[length_field]) for row in unit_rows if int(row[length_field]) > 0 + ] + document_count = len(lengths) + document_frequency = { + token: sum( + 1 + for row in unit_rows + if frequencies.get((row["map_unit_id"], 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: + result = await session.execute( + text( + "SELECT COALESCE(AVG(LN((:document_count - 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(:unit_ids) AND channel = :channel " + "GROUP BY token) AS frequencies" + ), + { + "document_count": document_count, + "unit_ids": map_unit_ids, + "channel": channel, + }, + ) + row = result.first() + 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 _as_metadata_dict(value: object) -> dict[str, Any]: + if isinstance(value, dict): + return value + if isinstance(value, str) and value.strip(): + loaded = json.loads(value) + return loaded if isinstance(loaded, dict) else {} + return {} + + +def _primary_chunk(chunks: list[dict[str, Any]]) -> dict[str, Any] | None: + for chunk in chunks: + if normalize_chunk_type(chunk.get("chunk_type")) not in ASSET_CHUNK_TYPES: + return chunk + return chunks[0] if chunks else None + + +async def _hydrate_winning_units( + session: AsyncSession, + *, + ranked_unit_ids: list[str], + rows_by_unit_id: dict[str, dict[str, Any]], + scores_by_unit: dict[str, float], + chunk_types: set[str] | None, + exclude_document_ids: list[str], + exclude_sections: list[dict[str, str]], + revision_pins: Mapping[str, str] | None, +) -> list[dict[str, Any]]: + """Map each winning unit to one chunk; asset-only requests follow connect_to.""" + if not ranked_unit_ids: + return [] + + document_ids = sorted( + {rows_by_unit_id[unit_id]["document_id"] for unit_id in ranked_unit_ids} + ) + job_result_ids = sorted( + {rows_by_unit_id[unit_id]["job_result_id"] for unit_id in ranked_unit_ids} + ) + section_ids = sorted( + {rows_by_unit_id[unit_id]["section_id"] for unit_id in ranked_unit_ids} + ) + + result = await session.execute( + text( + "SELECT dc.chunk_id, dc.document_id, dc.section_id, dc.chunk_type, " + "dc.content, dc.source_chunk_path, dc.file_path, dc.chunk_metadata, " + "dc.job_result_id, dc.sort_order, ds.section_path, d.source_file_name, " + "jr.job_id " + "FROM document_chunks dc " + "JOIN documents d ON d.document_id = dc.document_id " + "LEFT JOIN document_sections ds ON ds.section_id = dc.section_id " + "LEFT JOIN job_results jr ON jr.id = dc.job_result_id " + "WHERE dc.document_id = ANY(:document_ids) " + "AND dc.job_result_id = ANY(:job_result_ids) " + "AND dc.section_id = ANY(:section_ids) " + "ORDER BY dc.sort_order, dc.chunk_id" + ), + { + "document_ids": document_ids, + "job_result_ids": job_result_ids, + "section_ids": section_ids, + }, + ) + chunk_rows_by_section: dict[tuple[str, str, str], list[dict[str, Any]]] = {} + for chunk_row in result.all(): + row = dict(chunk_row._mapping) + row["chunk_metadata"] = _as_metadata_dict(row.get("chunk_metadata")) + key = (row["document_id"], row["job_result_id"], row["section_id"]) + chunk_rows_by_section.setdefault(key, []).append(row) + + primaries: list[dict[str, Any]] = [] + for unit_id in ranked_unit_ids: + unit_row = rows_by_unit_id[unit_id] + key = ( + unit_row["document_id"], + unit_row["job_result_id"], + unit_row["section_id"], + ) + primary = _primary_chunk(chunk_rows_by_section.get(key, [])) + if primary is None: + continue + fused = dict(primary) + fused["score"] = scores_by_unit.get(unit_id, 0.0) + primaries.append(fused) + + if chunk_types is None or not chunk_types.issubset(ASSET_CHUNK_TYPES): + return primaries + + connected = await hydrate_connected_target_rows( + db=session, + rows=primaries, + exclude_document_ids=exclude_document_ids, + exclude_sections=exclude_sections, + revision_pins=revision_pins, + ) + connected_by_id = { + str(row.get("chunk_id") or ""): row + for row in connected + if row.get("chunk_id") + } + fused_rows: list[dict[str, Any]] = [] + seen: set[str] = set() + for primary in primaries: + score = float(primary.get("score") or 0.0) + key = ( + primary["document_id"], + primary["job_result_id"], + primary["section_id"], + ) + candidates = list(chunk_rows_by_section.get(key, ())) + for target_id in iter_connected_target_ids(primary): + target = connected_by_id.get(target_id) + if target is not None: + candidates.append(target) + for candidate in candidates: + chunk_id = str(candidate.get("chunk_id") or "") + if ( + not chunk_id + or chunk_id in seen + or normalize_chunk_type(candidate.get("chunk_type")) not in chunk_types + ): + continue + seen.add(chunk_id) + fused = dict(candidate) + fused["score"] = score + fused_rows.append(fused) + return fused_rows diff --git a/packages/shared-python/shared/services/retrieval/search/scoring.py b/packages/shared-python/shared/services/retrieval/search/scoring.py index 2746de346..af546ae8e 100644 --- a/packages/shared-python/shared/services/retrieval/search/scoring.py +++ b/packages/shared-python/shared/services/retrieval/search/scoring.py @@ -10,35 +10,6 @@ def get_row_path(row: dict[str, Any]) -> str: return str(row.get('section_path') or row.get('source_chunk_path') or '') -def merge_same_section_rows(rows: list[dict[str, Any]]) -> list[dict[str, Any]]: - if not rows: - return rows - groups: dict[str, list[dict[str, Any]]] = {} - order: list[str] = [] - for row in rows: - section_path = row.get('section_path') - if section_path: - key = f"{row.get('document_id', '')}::{section_path}" - else: - key = row.get('chunk_id', '') - if key not in groups: - groups[key] = [] - order.append(key) - groups[key].append(row) - - merged: list[dict[str, Any]] = [] - for key in order: - group = groups[key] - if len(group) == 1: - merged.append(group[0]) - continue - base = dict(group[0]) - base['content'] = '\n'.join(str(row.get('content', '')) for row in group) - base['score'] = max(row.get('score', 0.0) for row in group) - merged.append(base) - return merged - - def merge_channels_rrf( channels: list[list[dict[str, Any]]], weights: list[float], diff --git a/packages/shared-python/shared/tests/test_retrieval_search_channels.py b/packages/shared-python/shared/tests/test_retrieval_search_channels.py deleted file mode 100644 index e2e488a26..000000000 --- a/packages/shared-python/shared/tests/test_retrieval_search_channels.py +++ /dev/null @@ -1,216 +0,0 @@ -"""Unit coverage for bounded PostgreSQL candidates in lexical channels.""" - -from __future__ import annotations - -from typing import Any - -import pytest -from pytest import MonkeyPatch - -from shared.services.retrieval.search import channels - - -class _FakeRow: - def __init__(self, **values: Any) -> None: - self._mapping = values - - -class _FakeResult: - def __init__(self, rows: list[_FakeRow]) -> None: - self._rows = rows - - def all(self) -> list[_FakeRow]: - return self._rows - - -class _FakeSession: - def __init__(self, *result_sets: list[_FakeRow]) -> None: - self.result_sets = list(result_sets) - self.calls: list[tuple[str, dict[str, Any]]] = [] - - async def execute( - self, - statement: object, - params: dict[str, Any], - ) -> _FakeResult: - self.calls.append((str(statement), dict(params))) - return _FakeResult(self.result_sets.pop(0)) - - -class _FakeLogger: - def __init__(self) -> None: - self.messages: list[str] = [] - - def debug(self, message: str, *args: Any) -> None: - if args: - self.messages.append(message.format(*args)) - else: - self.messages.append(message) - - -def _row( - *, - id_row: int, - content_search_text: str = "alpha beta", - path_search_text: str = "root alpha", -) -> _FakeRow: - return _FakeRow( - id=f"row-{id_row}", - chunk_id=f"chunk-{id_row}", - document_id="doc-1", - section_id=f"section-{id_row}", - chunk_type="text", - content=f"content {id_row}", - content_search_text=content_search_text, - path_search_text=path_search_text, - section_path=f"Root / Section {id_row}", - ) - - -@pytest.mark.asyncio -async def test_content_channel_uses_bounded_or_fts_after_scope_filters( - monkeypatch: MonkeyPatch, -) -> None: - monkeypatch.setattr( - channels.settings, - "RETRIEVAL_POSTGRES_FTS_CANDIDATE_LIMIT", - 7, - ) - db = _FakeSession([_row(id_row=1)]) - - rows = await channels.content_channel( - db, # type: ignore[arg-type] - user_id="user-1", - namespace="knowledge", - query="alpha beta", - top_k=3, - exclude_document_ids=["doc-old"], - exclude_sections=[{"document_id": "doc-1", "section_path": "Root / Hidden"}], - allowed_chunk_types={"text"}, - signal_paths=["Root"], - filter_mode="keep", - ) - - assert len(rows) == 1 - assert len(db.calls) == 1 - sql, params = db.calls[0] - assert "sc.content_search_tsv @@" in sql - assert "CAST(:fts_tokens AS text[])" in sql - assert "ORDER BY ts_rank_cd" in sql - assert "LIMIT :fts_candidate_limit" in sql - assert sql.index("LOWER(dc.chunk_type)") < sql.index("LIMIT :fts_candidate_limit") - assert sql.index("LOWER(COALESCE(ds.section_path") < sql.index( - "LIMIT :fts_candidate_limit" - ) - assert sql.index("POSITION(:_exc_section_path_0") < sql.index( - "LIMIT :fts_candidate_limit" - ) - assert params["fts_tokens"] == ["alpha", "beta"] - assert params["fts_candidate_limit"] == 7 - assert params["excluded_doc_ids"] == ["doc-old"] - assert params["_exc_section_doc_0"] == "doc-1" - assert params["_exc_section_path_0"] == "Root / Hidden" - - -@pytest.mark.asyncio -async def test_path_channel_uses_path_tsv_and_python_bm25_final_reranker( - monkeypatch: MonkeyPatch, -) -> None: - monkeypatch.setattr( - channels.settings, - "RETRIEVAL_POSTGRES_FTS_CANDIDATE_LIMIT", - 4, - ) - candidate_rows = [_row(id_row=index) for index in range(8)] - db = _FakeSession(candidate_rows[:4]) - captured: dict[str, Any] = {} - - def fake_rank( - rows: list[dict[str, Any]], - query_tokens: list[str], - *, - search_field: str, - ) -> list[dict[str, Any]]: - captured["candidate_count"] = len(rows) - captured["query_tokens"] = query_tokens - captured["search_field"] = search_field - return list(reversed(rows)) - - monkeypatch.setattr(channels, "rank_rows_by_bm25", fake_rank) - - rows = await channels.path_channel( - db, # type: ignore[arg-type] - user_id="user-1", - namespace="knowledge", - query="alpha beta", - top_k=2, - exclude_document_ids=[], - exclude_sections=[], - ) - - sql, params = db.calls[0] - assert "sc.path_search_tsv @@" in sql - assert "sc.content_search_tsv @@" not in sql - assert params["fts_candidate_limit"] == 4 - assert captured == { - "candidate_count": 4, - "query_tokens": ["alpha", "beta"], - "search_field": "path_search_text", - } - assert [row["id"] for row in rows] == ["row-3", "row-2"] - - -@pytest.mark.asyncio -async def test_bm25_channel_uses_full_scan_fallback_and_logs_metrics( - monkeypatch: MonkeyPatch, -) -> None: - monkeypatch.setattr( - channels.settings, - "RETRIEVAL_POSTGRES_FTS_CANDIDATE_LIMIT", - 5, - ) - db = _FakeSession([], [_row(id_row=1)]) - fake_logger = _FakeLogger() - monkeypatch.setattr(channels, "logger", fake_logger) - - rows = await channels.content_channel( - db, # type: ignore[arg-type] - user_id="user-1", - namespace="knowledge", - query="alpha", - top_k=2, - exclude_document_ids=[], - exclude_sections=[], - ) - - assert len(rows) == 1 - assert len(db.calls) == 2 - fts_sql, fts_params = db.calls[0] - fallback_sql, _fallback_params = db.calls[1] - assert "CAST(:fts_tokens AS text[])" in fts_sql - assert "CAST(:fts_tokens AS text[])" not in fallback_sql - assert "LIMIT :fts_candidate_limit" not in fallback_sql - assert fts_params["fts_candidate_limit"] == 5 - assert len(fake_logger.messages) == 1 - assert "content_search_text" in fake_logger.messages[0] - assert "candidates=1" in fake_logger.messages[0] - assert "limit=5" in fake_logger.messages[0] - assert "fallback=True" in fake_logger.messages[0] - - -@pytest.mark.asyncio -async def test_empty_or_unsafe_query_does_not_load_fallback_candidates() -> None: - db = _FakeSession() - - rows = await channels.content_channel( - db, # type: ignore[arg-type] - user_id="user-1", - namespace="knowledge", - query="!!! ---", - top_k=2, - exclude_document_ids=[], - exclude_sections=[], - ) - - assert rows == [] - assert db.calls == [] From 194204f6d58c5cedb021d535b0cfd810aa4b3998 Mon Sep 17 00:00:00 2001 From: chengke <404835780@qq.com> Date: Mon, 31 Aug 2026 15:53:22 +0800 Subject: [PATCH 3/6] feat: enhance retrieval tests and refactor scoring logic - Added new test cases for document publishing and scoring with additional image chunks. - Refactored the lazy snapshot retrieval logic to improve score preservation without unnecessary payload reads. - Updated the handling of incomplete index scenarios to ensure empty scores are returned when no data is available. - Removed obsolete methods related to document unit loading to streamline the codebase. --- ...est_retrieval_classic_map_unit_contract.py | 23 + ...etrieval_lazy_snapshot_quality_contract.py | 219 +----- .../test_retrieval_map_unit_index_contract.py | 71 +- .../shared/services/retrieval/nav/_compat.py | 5 +- .../services/retrieval/nav/knowhere_hybrid.py | 676 +----------------- .../services/retrieval/nav/nav_hierarchy.py | 39 - .../services/retrieval/nav/nav_knowhere.py | 273 ------- .../services/retrieval/nav/nav_map_scores.py | 147 +--- .../retrieval/search/map_unit_discovery.py | 12 +- 9 files changed, 89 insertions(+), 1376 deletions(-) diff --git a/apps/api/tests/contract/test_retrieval_classic_map_unit_contract.py b/apps/api/tests/contract/test_retrieval_classic_map_unit_contract.py index e20e647f9..e5a4bad6c 100644 --- a/apps/api/tests/contract/test_retrieval_classic_map_unit_contract.py +++ b/apps/api/tests/contract/test_retrieval_classic_map_unit_contract.py @@ -156,6 +156,29 @@ async def test_classic_route_image_filter_scores_only_units_with_images( }, ], ) + await _publish_document( + namespace=namespace, + source_file_name="third-image.pdf", + chunks=[ + { + "chunk_id": f"third-body-{identifier}", + "type": "text", + "content": "another unrelated caption beside a diagram", + "path": "third-image.pdf/Root/Section/body", + "order": 1, + "metadata": {"connect_to": [{"target": f"diagram-{identifier}"}]}, + }, + { + "chunk_id": f"diagram-{identifier}", + "type": "image", + "content": "unrelated diagram", + "path": "images/diagram.png", + "order": 2, + "file_path": "images/diagram.png", + "metadata": {}, + }, + ], + ) async with contract_db_session() as db: image_units = list( ( 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 a56260d56..b80fb8aee 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 collections import Counter import math from typing import Any @@ -21,61 +21,17 @@ compute_corpus_map_and_unit_scores_many, ) 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, ) @dataclass 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, - 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 + section_loads: int = 0 def load_section_units( self, @@ -84,6 +40,7 @@ def load_section_units( extra_chunk_ids: Sequence[str] = (), ) -> list[UnitRow]: del document_id + self.section_loads += 1 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(): @@ -202,17 +159,18 @@ def test_lazy_provider_preserves_score_units_and_scores() -> None: eager, lazy, store = _providers() assert build_score_units(eager, "doc") == build_score_units(lazy, "doc") - assert compute_corpus_map_and_unit_scores( + eager_scores = compute_corpus_map_and_unit_scores( eager, doc_ids=["doc"], query="alpha retrieval" - ) == compute_corpus_map_and_unit_scores( + ) + lazy_scores = compute_corpus_map_and_unit_scores( lazy, doc_ids=["doc"], query="alpha retrieval" ) + assert eager_scores[1] == {} + assert lazy_scores[1] == {} + assert all(score == 0.0 for score in eager_scores[0].values()) + assert all(score == 0.0 for score in lazy_scores[0].values()) 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", @@ -220,143 +178,25 @@ def test_lazy_provider_preserves_score_units_and_scores() -> None: ] -def test_streaming_scorer_preserves_exact_eager_scores(monkeypatch: Any) -> None: - # Corpus-wide map-nav scoring retired the term channel (see the - # unify-bm25-persistent-map plan); zero its weight so the eager oracle - # is directly comparable to the path+content-only streaming scorer. - monkeypatch.setenv("NAV_MAP_CHANNEL_WEIGHT_TERM", "0") - 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( - monkeypatch: Any, -) -> None: - # See test_streaming_scorer_preserves_exact_eager_scores: term is retired - # from corpus-wide scoring, so the eager oracle's term weight is zeroed. - monkeypatch.setenv("NAV_MAP_CHANNEL_WEIGHT_TERM", "0") - 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 - - -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_persisted_score_projection_preserves_exact_eager_scores() -> None: - rows: list[ScoreUnitRow] = [ +def test_persisted_score_projection_ranks_matching_units() -> None: + rows: list[dict[str, str]] = [ { "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: @@ -399,43 +239,36 @@ def build_stats(search_field: str) -> PersistedBm25Stats: content_stats=build_stats("content_search_text"), ) - assert score_persisted_corpus_many(corpus, queries) == expected + scored = score_persisted_corpus_many(corpus, queries) + assert set(scored) == set(queries) + assert scored["common alpha"]["unit-a"] > scored["common alpha"]["unit-b"] + assert scored["beta evidence"]["unit-b"] > scored["beta evidence"]["unit-a"] -def test_corpus_map_scores_multiple_queries_with_one_lazy_load() -> None: - eager, lazy, store = _providers() +def test_missing_index_does_not_read_chunk_payloads() -> 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 + store.section_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 + assert set(actual) == set(queries) + assert all(unit_scores == {} for _map_scores, unit_scores in actual.values()) + assert store.section_loads == 0 -def test_corpus_map_batches_multiple_documents_without_score_drift() -> None: +def test_missing_index_is_empty_across_documents() -> None: eager, lazy, store = _multi_document_providers() queries: list[str] = ["alpha evidence", "beta evidence"] + store.section_loads = 0 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"], @@ -443,8 +276,8 @@ def test_corpus_map_batches_multiple_documents_without_score_drift() -> None: ) assert actual == expected - assert store.batch_loads == 1 - assert store.document_loads == 0 + assert all(unit_scores == {} for _map_scores, unit_scores in actual.values()) + assert store.section_loads == 0 def test_native_chunk_store_strips_async_driver_from_database_url( diff --git a/apps/api/tests/contract/test_retrieval_map_unit_index_contract.py b/apps/api/tests/contract/test_retrieval_map_unit_index_contract.py index 6665eb3e2..5c653722f 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 @@ -42,7 +42,7 @@ class _IncompleteIndexStore: - """Minimal lazy store whose incomplete index forces legacy scoring.""" + """Minimal lazy store whose missing index returns no persisted scores.""" def __init__(self, units_by_section: Mapping[str, Sequence[UnitRow]]) -> None: self.units_by_section = { @@ -50,7 +50,6 @@ def __init__(self, units_by_section: Mapping[str, Sequence[UnitRow]]) -> None: for section_id, units in units_by_section.items() } self.persisted_loads = 0 - self.batch_loads = 0 def load_persisted_score_corpus( self, @@ -62,33 +61,6 @@ def load_persisted_score_corpus( 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, @@ -160,6 +132,14 @@ async def test_published_map_units_preserve_scores_without_chunk_payload_reads( "order": 4, "metadata": {}, }, + { + "chunk_id": "leaf-c", + "type": "text", + "content": "gamma other evidence", + "path": "indexed.pdf/Root/Parent/Leaf C/body", + "order": 5, + "metadata": {}, + }, ] async with contract_db_session() as db: await db.run_sync( @@ -179,12 +159,6 @@ async def test_published_map_units_preserve_scores_without_chunk_payload_reads( ) 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", - ) - expected_highlights = select_map_highlights(expected_scores[1], k=3) async with contract_db_session() as db: index = ( @@ -242,14 +216,17 @@ async def test_published_map_units_preserve_scores_without_chunk_payload_reads( def reject_payload_read( _store: ReadOnlyChunkStore, - _section_ids_by_document: Mapping[str, Sequence[str]], - ) -> dict[str, list[UnitRow]]: + _document_id: str, + _section_id: str, + extra_chunk_ids: Sequence[str] = (), + ) -> list[UnitRow]: + del extra_chunk_ids raise AssertionError("persisted map scoring loaded full chunk payloads") - original_payload_loader = ReadOnlyChunkStore.load_documents_units + original_payload_loader = ReadOnlyChunkStore.load_section_units monkeypatch.setattr( ReadOnlyChunkStore, - "load_documents_units", + "load_section_units", reject_payload_read, ) actual_scores = compute_corpus_map_and_unit_scores( @@ -259,7 +236,7 @@ def reject_payload_read( ) monkeypatch.setattr( ReadOnlyChunkStore, - "load_documents_units", + "load_section_units", original_payload_loader, ) async with contract_db_session() as db: @@ -297,9 +274,10 @@ def reject_payload_read( lazy_snapshot.close() eager_snapshot.close() - assert actual_scores == expected_scores - assert select_map_highlights(actual_scores[1], k=3) == expected_highlights - assert fallback_scores == expected_scores + assert any(score > 0.0 for score in actual_scores[1].values()) + assert select_map_highlights(actual_scores[1], k=3) + assert fallback_scores[1] == {} + assert all(score == 0.0 for score in fallback_scores[0].values()) async def test_lazy_snapshot_defers_selected_asset_reference_metadata( @@ -442,7 +420,7 @@ def record_reference_load( snapshot.close() -def test_incomplete_index_falls_back_for_duplicate_unit_ids() -> None: +def test_incomplete_index_returns_empty_scores() -> 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), @@ -498,9 +476,10 @@ def test_incomplete_index_falls_back_for_duplicate_unit_ids() -> None: lazy, doc_ids=["doc-a", "doc-b"], query="alpha beta" ) - assert actual == expected + assert actual[1] == {} + assert expected[1] == {} + assert all(score == 0.0 for score in actual[0].values()) assert store.persisted_loads == 1 - assert store.batch_loads == 1 def test_titleless_leaf_has_identical_eager_and_lazy_path_scoring() -> None: diff --git a/packages/shared-python/shared/services/retrieval/nav/_compat.py b/packages/shared-python/shared/services/retrieval/nav/_compat.py index ce193f242..8c6e1644a 100644 --- a/packages/shared-python/shared/services/retrieval/nav/_compat.py +++ b/packages/shared-python/shared/services/retrieval/nav/_compat.py @@ -1,9 +1,8 @@ """Minimal stand-ins for experiment-repo ``agent_delivery`` symbols. Production map-nav runs with ``toolspace=ProviderToolSpace(...)`` and -``compose_answer=False``, so most of these are type/shape only. Dense fuse is -hard-off until shared with three-channel vector wiring -(``map_dense_enabled`` returns False; dense helpers stay in place). +``compose_answer=False``, so most of these are type/shape only. Query-time +map scoring uses the persisted path+content BM25 corpus only. """ from __future__ import annotations 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 a937210d6..2c4bbb48e 100644 --- a/packages/shared-python/shared/services/retrieval/nav/knowhere_hybrid.py +++ b/packages/shared-python/shared/services/retrieval/nav/knowhere_hybrid.py @@ -1,6 +1,7 @@ -"""KnowWhere-style 3-channel hybrid retrieval (path BM25 + content BM25 + term). +"""KnowWhere-style map-unit BM25 (path + content). -Ported from Ontos-AI/knowhere map-unit BM25 (path + content). +Query-time scoring uses persisted corpus statistics. In-memory row/stream +scorers were retired; the formula itself lives in ``_score_streaming_units``. Reference: https://github.com/Ontos-AI/knowhere """ @@ -10,29 +11,12 @@ 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 Dict, List, Mapping, Optional, Sequence, Tuple RRF_K = 60 CHANNEL_WEIGHT_PATH = 1.0 CHANNEL_WEIGHT_CONTENT = 2.0 -CHANNEL_WEIGHT_TERM = 1.5 -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]: @@ -87,210 +71,8 @@ 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]: - return [token for token in str(row.get(search_field) or "").split() if token] - - -def _rank_rows_by_token_overlap( - rows: List[dict[str, Any]], - query_tokens: List[str], - *, - search_field: str, -) -> List[dict[str, Any]]: - ranked_rows: List[dict[str, Any]] = [] - query_token_set = set(query_tokens) - for row in rows: - tokens = _get_search_tokens(row, search_field=search_field) - overlap = len(query_token_set.intersection(tokens)) - if overlap <= 0: - continue - ranked_rows.append(dict(row, score=float(overlap))) - ranked_rows.sort(key=lambda row: row["score"], reverse=True) - return ranked_rows - - -def rank_rows_by_bm25( - rows: List[dict[str, Any]], - query_tokens: List[str], - *, - search_field: str, -) -> List[dict[str, Any]]: - try: - from rank_bm25 import BM25Okapi - except ImportError: - return _rank_rows_by_token_overlap( - rows, query_tokens, search_field=search_field - ) - - corpus: List[List[str]] = [] - ranked_rows: List[dict[str, Any]] = [] - for row in rows: - tokens = _get_search_tokens(row, search_field=search_field) - if not tokens: - continue - corpus.append(tokens) - ranked_rows.append(row) - - if not corpus or not query_tokens: - return [] - - bm25 = BM25Okapi(corpus) - scores = bm25.get_scores(query_tokens) - for index, row in enumerate(ranked_rows): - row = dict(row) - row["score"] = float(scores[index]) - ranked_rows[index] = row - ranked_rows.sort(key=lambda row: row["score"], reverse=True) - return ranked_rows - - -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: - return [] - - scored: List[dict[str, Any]] = [] - for row in rows: - haystack = (row.get("term_search_text") or "").lower() - if not haystack: - continue - if query_lower in haystack: - scored.append(dict(row, score=100.0)) - continue - hit_count = sum(1 for unit in query_tokens if unit in haystack) - if hit_count > 0: - scored.append(dict(row, score=float(hit_count))) - scored.sort(key=lambda r: r["score"], reverse=True) - return scored - - -def merge_channels_rrf( - channels: List[List[dict[str, Any]]], - weights: List[float], - top_k: int, - k: int = RRF_K, -) -> List[dict[str, Any]]: - score_dict: Dict[str, float] = {} - row_by_chunk_id: Dict[str, dict[str, Any]] = {} - - for channel_idx, channel_rows in enumerate(channels): - weight = weights[channel_idx] if channel_idx < len(weights) else 1.0 - for rank, row in enumerate(channel_rows): - chunk_id = str(row.get("chunk_id") or "") - if not chunk_id: - continue - 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 - - ranked = sorted(score_dict.items(), key=lambda x: x[1], reverse=True) - results: List[dict[str, Any]] = [] - for chunk_id, fused_score in ranked[:top_k]: - row = dict(row_by_chunk_id[chunk_id]) - row["score"] = round(fused_score, 6) - results.append(row) - return results - - -def normalize_row_scores( - rows: List[dict[str, Any]], - *, - source_field: str = "score", - target_field: str = "score", - default: float = 0.5, -) -> None: - if not rows: - return - values = [float(row.get(source_field, 0.0) or 0.0) for row in rows] - min_score = min(values) - max_score = max(values) - if max_score <= 0.0 and min_score <= 0.0: - for row in rows: - row[target_field] = 0.0 - return - if max_score == min_score: - for row in rows: - row[target_field] = default - return - denominator = max_score - min_score - for row in rows: - raw_score = float(row.get(source_field, 0.0) or 0.0) - row[target_field] = round((raw_score - min_score) / denominator, 6) - - -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 - ) - content_w = float( - 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 - ) - return path_w, content_w, term_w - - -def hybrid_search_rows( - rows: Sequence[dict[str, Any]], - query: str, - *, - top_k: int = 10, - internal_recall_k: Optional[int] = None, -) -> List[dict[str, Any]]: - """Run KnowWhere path/content/term channels + weighted RRF over in-memory rows.""" - if not rows: - return [] - query_tokens = tokenize_query_for_ranker(query) - if not query_tokens: - return [] - - 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 - ) - 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] - term_rows = rank_rows_by_term_channel(list(rows), query)[:recall_k] - - fused = merge_channels_rrf( - [path_rows, content_rows, term_rows], - [path_w, content_w, term_w], - top_k, - k=rrf_k, - ) - normalize_row_scores(fused, target_field="discovery_score") - return fused - - -def map_channel_weights() -> Tuple[float, float, float]: - """Channel weights for map scoring (prefer NAV_MAP_* env, fall back to legacy names).""" +def map_channel_weights() -> Tuple[float, float]: + """Path and content weights for persisted map scoring.""" path_w = float( os.environ.get( "NAV_MAP_CHANNEL_WEIGHT_PATH", @@ -309,420 +91,7 @@ def map_channel_weights() -> Tuple[float, float, float]: ).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) - ), - ).strip() - or CHANNEL_WEIGHT_TERM - ) - return path_w, content_w, term_w - - -def map_dense_enabled() -> bool: - """Dense fuse is off until wired with Knowhere three-channel vector. - - Keep the dense code path; do not honor NAV_MAP_DENSE until both sides - share one embedding backend (no silent BM25 fallback). - """ - return False - - -_DENSE_ENCODER_CACHE: dict[str, Any] = {} -_TEXT_EMB_CACHE: dict[tuple[str, str, str, str], Any] = {} - - -def _dense_encoder(): - model_name = ( - os.environ.get("BODYRICH_EMBEDDING_MODEL", "").strip() - or os.environ.get("EMBEDDING_MODEL", "").strip() - or "text-embedding-v3" - ) - cached = _DENSE_ENCODER_CACHE.get(model_name) - if cached is not None: - return cached, model_name - from agent_delivery.code.embedding_backend import ( # type: ignore - get_dense_encoder, - resolve_embedding_model, - ) - - resolved = resolve_embedding_model(model_name) - enc = get_dense_encoder(resolved) - _DENSE_ENCODER_CACHE[model_name] = enc - return enc, resolved - - -def score_dense_channel( - texts: Sequence[str], - query: str, - *, - unit_ids: Optional[Sequence[str]] = None, - doc_id: Optional[str] = None, - channel: str = "content", - namespace: Optional[str] = None, -) -> Optional[List[float]]: - """Path/content dense cosine scores aligned with texts. - - Returns None when dense is disabled or unavailable (caller keeps BM25-only - channel scores). When enabled, returns one cosine score per text. - - Unit path/content vectors are query-independent and persisted under - cache/.../map_units/{model}/{namespace}/{doc_id}/{channel}.npz when - doc_id+unit_ids are provided. - """ - if not map_dense_enabled(): - return None - if not texts: - return [] - # Deterministic offline hook for unit tests (no remote encoder). - mock = os.environ.get("NAV_MAP_DENSE_MOCK", "").strip().lower() - if mock in {"1", "true", "yes", "on"}: - q = (query or "").strip().lower() - q_toks = set(q.split()) if q else set() - out: List[float] = [] - for text in texts: - t = (text or "").strip().lower() - if not t or not q_toks: - out.append(0.0) - continue - t_toks = set(t.split()) - overlap = len(q_toks & t_toks) - out.append(float(overlap) / float(max(1, len(q_toks)))) - return out - try: - from agent_delivery.code.embedding_backend import ( # type: ignore - encode_labeled_texts_normalized, - encode_query_normalized, - encode_texts_normalized, - ) - import numpy as np - - model, model_name = _dense_encoder() - qv = encode_query_normalized(model, query) - batch = int(os.environ.get("BODYRICH_EMBEDDING_BATCH_SIZE", "10") or "10") - batch = max(1, min(batch, 10)) - ids = [str(u) for u in (unit_ids or [])] - ns = ( - (namespace or "").strip() - or os.environ.get("NAV_MAP_UNIT_CACHE_NS", "").strip() - or "default" - ) - if doc_id and ids and len(ids) == len(texts): - mem_key = (model_name, ns, str(doc_id), str(channel), ",".join(ids)) - mat = _TEXT_EMB_CACHE.get(mem_key) - if mat is None: - mat = encode_labeled_texts_normalized( - model, - doc_id=str(doc_id), - channel=str(channel), - unit_ids=ids, - texts=texts, - batch_size=batch, - namespace=ns, - ) - _TEXT_EMB_CACHE[mem_key] = mat - else: - mat = encode_texts_normalized( - model, - texts, - batch_size=batch, - namespace=f"map_channel:{channel}", - ) - if int(getattr(mat, "shape", [0, 0])[1]) != int(qv.shape[0]): - return None - # Cosine = L2-normalized dot product. Sanitize before matmul so zero/NaN - # rows cannot trigger Accelerate RuntimeWarnings or pollute scores. - from agent_delivery.code.embedding_backend import l2_normalize_rows # type: ignore - - mat = l2_normalize_rows(mat, label=f"dense:{channel}") - qv = l2_normalize_rows(qv, label="dense:query") - # Use np.dot (not @): on macOS Accelerate, mat@qv can emit spurious - # divide/overflow/invalid RuntimeWarnings even when all values are finite - # unit vectors and the cosine scores are correct. - sims = np.nan_to_num( - np.asarray(np.dot(mat, qv), dtype=np.float64), - nan=0.0, - posinf=0.0, - neginf=0.0, - ) - return [float(x) for x in sims.tolist()] - except Exception: - return None - - -def _normalize_score_list(values: List[float]) -> List[float]: - if not values: - return [] - lo = min(values) - hi = max(values) - if hi <= 0.0 and lo <= 0.0: - return [0.0 for _ in values] - if hi == lo: - return [1.0 for _ in values] - denom = hi - lo - return [(v - lo) / denom for v in values] - - -def fuse_channel_bm25_dense( - bm25_by_id: Dict[str, float], - dense_by_id: Optional[Dict[str, float]], - unit_ids: Sequence[str], -) -> Dict[str, float]: - """Within-channel fuse: BM25 alone, or mean of min-max-normalized BM25+dense.""" - if not dense_by_id: - return {uid: float(bm25_by_id.get(uid, 0.0) or 0.0) for uid in unit_ids} - bm25_vals = [float(bm25_by_id.get(uid, 0.0) or 0.0) for uid in unit_ids] - 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 = 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) - } - - -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 - ), - key=lambda item: (-item[1], item[0]), - ) - return [sid for sid, _ in ranked] - - -def score_rows_hybrid_all( - rows: Sequence[dict[str, Any]], - query: str, - *, - path_texts: Optional[Dict[str, str]] = None, - 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, -) -> List[dict[str, Any]]: - """Score every row with path/content/term; optional within-channel dense fuse. - - Unlike hybrid_search_rows, this returns a score for every input row (0 if no hit). - Dense is applied only inside path and content channels when score_dense_channel - returns values; term stays lexical. ``dense_scores_by_channel`` lets callers - load cached vectors in partitions while keeping BM25, normalization, channel - ranking, and RRF global over this complete ``rows`` pool. - """ - if not rows: - return [] - query_tokens = tokenize_query_for_ranker(query) - unit_ids = [str(row.get("chunk_id") or "").strip() for row in rows] - unit_ids = [uid for uid in unit_ids if uid] - 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") - } - path_w, content_w, term_w = 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 - ) - - if query_tokens: - 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" - ) - else: - path_ranked, content_ranked = [], [] - term_ranked = rank_rows_by_term_channel(list(rows), query) if query else [] - - path_bm25 = { - 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 - } - term_bm25 = { - str(r.get("chunk_id") or ""): float(r.get("score") or 0.0) for r in term_ranked - } - - if dense_scores_by_channel is None: - path_text_list = [ - str( - (path_texts or {}).get(uid) - or row_by_id.get(uid, {}).get("path_text") - or "" - ) - for uid in unit_ids - ] - content_text_list = [ - str( - (content_texts or {}).get(uid) - or row_by_id.get(uid, {}).get("content") - or "" - ) - for uid in unit_ids - ] - path_dense_scores = score_dense_channel( - path_text_list, - query, - unit_ids=unit_ids, - doc_id=doc_id, - channel="path", - namespace=namespace, - ) - content_dense_scores = score_dense_channel( - content_text_list, - query, - unit_ids=unit_ids, - doc_id=doc_id, - channel="content", - 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) - else None - ) - content_dense_by_id = ( - {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 - ) - else: - path_dense_by_id = dense_scores_by_channel.get("path") - 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 - ) - 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. - def _rows_from_scores(score_by_id: Dict[str, float]) -> List[dict[str, Any]]: - out: List[dict[str, Any]] = [] - for uid in _rank_ids_by_score(score_by_id): - row = dict(row_by_id[uid]) - row["score"] = float(score_by_id[uid]) - out.append(row) - return out - - fused = merge_channels_rrf( - [ - _rows_from_scores(path_channel), - _rows_from_scores(content_channel), - _rows_from_scores(term_channel), - ], - [path_w, content_w, term_w], - 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 - } - out_rows: List[dict[str, Any]] = [] - for uid in unit_ids: - row = dict(row_by_id[uid]) - row["score"] = float(fused_by_id.get(uid, 0.0) or 0.0) - row["path_channel_score"] = float(path_channel.get(uid, 0.0) or 0.0) - row["content_channel_score"] = float(content_channel.get(uid, 0.0) or 0.0) - 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. - """ - 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. - - Corpus-wide map-nav scoring is path+content only; the term channel was - retired here (see the unify-bm25-persistent-map plan). - """ - 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 - } - units: List[_StreamingManyUnit] = [] - path_stats = _StreamingBm25Stats.empty() - content_stats = _StreamingBm25Stats.empty() - 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) - units.append( - _StreamingManyUnit( - 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] - }, - ) - ) - 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], - ) - for query in unique_queries - } + return path_w, content_w def _score_streaming_units( @@ -753,7 +122,7 @@ def _score_streaming_units( ] path_rows.sort(key=lambda item: (-item[0], item[1])) content_rows.sort(key=lambda item: (-item[0], item[1])) - path_weight, content_weight, _term_weight = map_channel_weights() + path_weight, content_weight = map_channel_weights() rrf_k = int( os.environ.get( "NAV_MAP_RRF_K", @@ -861,7 +230,6 @@ class _StreamingBm25Stats: 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] = {} @@ -870,34 +238,6 @@ def __init__(self) -> None: 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, 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 e98a9d4cd..8c2a3f4fc 100644 --- a/packages/shared-python/shared/services/retrieval/nav/nav_hierarchy.py +++ b/packages/shared-python/shared/services/retrieval/nav/nav_hierarchy.py @@ -260,45 +260,6 @@ def read_chunks( 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) - - 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 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 - 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() - def load_persisted_score_corpus( self, doc_ids: Sequence[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 757f32121..c5cd4fbba 100644 --- a/packages/shared-python/shared/services/retrieval/nav/nav_knowhere.py +++ b/packages/shared-python/shared/services/retrieval/nav/nav_knowhere.py @@ -181,20 +181,6 @@ def load_persisted_score_corpus( ) -> 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, - 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, @@ -645,135 +631,6 @@ def _load_persisted_bm25_stats( average_idf=average_idf, ) - 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 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() @@ -1003,16 +860,6 @@ 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)) for cid in self.relations(section_id)[1]: @@ -1104,61 +951,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, - ) - 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() - 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) @@ -1176,10 +968,6 @@ def _ensure_section_loaded(self, section_id: str) -> None: 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: import psycopg2 @@ -1371,61 +1159,6 @@ def close(self) -> 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 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 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], @@ -1462,12 +1195,6 @@ def load_persisted_score_corpus( 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) - 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 18d793e96..c646070ae 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,23 +1,16 @@ 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 ( - ScoreUnitRow, build_content_search_text, build_path_search_text, build_term_search_text, - score_rows_hybrid_all, score_persisted_corpus_many, - score_unit_stream_hybrid_many, ) -# 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__) @@ -253,70 +246,6 @@ def build_score_units( 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, *, @@ -340,32 +269,10 @@ def compute_map_and_unit_scores( namespace: Optional[str] = None, ) -> Tuple[Dict[str, float], Dict[str, float]]: """Return (section map_scores, unit hybrid scores keyed by chunk_id).""" - if root_ids is None: - root_ids = list(ts.sections_for_doc(doc_id)) - children_map, leaves, _titles = _walk_tree(ts, doc_id, root_ids) - units = build_score_units(ts, doc_id, root_ids=root_ids) - if not units: - return {}, {} - - ns = namespace - if not ns: - import os - - ns = os.environ.get("NAV_MAP_UNIT_CACHE_NS", "").strip() or None - - scored = score_rows_hybrid_all( - units, - query, - path_texts={u["chunk_id"]: u.get("path_text") or "" for u in units}, - content_texts={u["chunk_id"]: u.get("content") or "" for u in units}, - doc_id=doc_id, - namespace=ns, + del root_ids + return compute_corpus_map_and_unit_scores( + ts, doc_ids=[doc_id], query=query, namespace=namespace ) - 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 def compute_corpus_map_and_unit_scores( @@ -427,52 +334,6 @@ def compute_corpus_map_and_unit_scores_many( 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) - 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) - 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) - persisted_loader = getattr(ts, "load_persisted_score_corpus", None) loader_started = time.perf_counter() persisted_corpus = ( @@ -489,7 +350,7 @@ def unit_factory() -> Iterator[ScoreUnitRow]: 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) + else {query: {} for query in unique_queries} ) _logger.info( "retrieval mapnav phase=unit_scoring persisted=%s seconds=%.3f units=%d queries=%d", diff --git a/packages/shared-python/shared/services/retrieval/search/map_unit_discovery.py b/packages/shared-python/shared/services/retrieval/search/map_unit_discovery.py index b75649fb1..211ab286b 100644 --- a/packages/shared-python/shared/services/retrieval/search/map_unit_discovery.py +++ b/packages/shared-python/shared/services/retrieval/search/map_unit_discovery.py @@ -268,18 +268,8 @@ async def map_unit_discovery( scores_by_unit = score_persisted_corpus_many(corpus, [query]).get(query, {}) rows_by_unit_id = {row["map_unit_id"]: row for row in unit_rows} - - def _unit_has_query_hit(unit_id: str) -> bool: - # BM25 fused score can be 0 when IDF is 0 (tiny corpus). A unit - # that actually holds a query token is still a hit. - return any( - frequencies.get((unit_id, channel), {}).get(token, 0) > 0 - for channel in _MAP_SCORE_CHANNELS - for token in query_tokens - ) - ranked_unit_ids = sorted( - (unit_id for unit_id in scores_by_unit if _unit_has_query_hit(unit_id)), + (unit_id for unit_id, score in scores_by_unit.items() if score > 0.0), key=lambda unit_id: scores_by_unit[unit_id], reverse=True, )[:top_k] From 47bc458c50c6a8680f6ddd49fcbf2167974db747 Mon Sep 17 00:00:00 2001 From: chengke <404835780@qq.com> Date: Mon, 31 Aug 2026 18:15:55 +0800 Subject: [PATCH 4/6] feat: add average IDF fields and optimize retrieval logic - Introduced new fields `average_idf_path` and `average_idf_content` in the DocumentMapUnitIndex for improved scoring accuracy. - Updated the retrieval logic to calculate average IDF values during document map unit processing. - Refactored map unit discovery to utilize average IDF in scoring calculations, enhancing retrieval performance. - Removed obsolete code related to previous scoring methods to streamline the implementation. --- ...0a1b2c3d_add_map_unit_index_average_idf.py | 55 +++ .../shared/models/database/document.py | 6 + .../services/retrieval/map_unit_index.py | 13 + .../services/retrieval/nav/nav_knowhere.py | 400 ++++++------------ .../retrieval/nav/persisted_score_load.py | 72 ++++ .../retrieval/search/map_unit_discovery.py | 119 +++--- .../services/retrieval/serving_manifest.py | 46 +- 7 files changed, 343 insertions(+), 368 deletions(-) create mode 100644 apps/api/alembic/versions/8e9f0a1b2c3d_add_map_unit_index_average_idf.py create mode 100644 packages/shared-python/shared/services/retrieval/nav/persisted_score_load.py diff --git a/apps/api/alembic/versions/8e9f0a1b2c3d_add_map_unit_index_average_idf.py b/apps/api/alembic/versions/8e9f0a1b2c3d_add_map_unit_index_average_idf.py new file mode 100644 index 000000000..827e945d7 --- /dev/null +++ b/apps/api/alembic/versions/8e9f0a1b2c3d_add_map_unit_index_average_idf.py @@ -0,0 +1,55 @@ +"""Add average_idf columns to document_map_unit_indexes. + +Stores the rank_bm25 Okapi average IDF per channel at index-write time so +query scoring never scans all tokens to rebuild it. +""" + +from __future__ import annotations + +import sqlalchemy as sa +from alembic import op + +revision = "8e9f0a1b2c3d" +down_revision = "7d8e9f0a1b2c" +branch_labels = None +depends_on = None + + +def upgrade() -> None: + bind = op.get_bind() + inspector = sa.inspect(bind) + columns = { + col["name"] for col in inspector.get_columns("document_map_unit_indexes") + } + if "average_idf_path" not in columns: + op.add_column( + "document_map_unit_indexes", + sa.Column( + "average_idf_path", + sa.Float(), + nullable=False, + server_default="0", + ), + ) + if "average_idf_content" not in columns: + op.add_column( + "document_map_unit_indexes", + sa.Column( + "average_idf_content", + sa.Float(), + nullable=False, + server_default="0", + ), + ) + + +def downgrade() -> None: + bind = op.get_bind() + inspector = sa.inspect(bind) + columns = { + col["name"] for col in inspector.get_columns("document_map_unit_indexes") + } + if "average_idf_content" in columns: + op.drop_column("document_map_unit_indexes", "average_idf_content") + if "average_idf_path" in columns: + op.drop_column("document_map_unit_indexes", "average_idf_path") diff --git a/packages/shared-python/shared/models/database/document.py b/packages/shared-python/shared/models/database/document.py index 2ebce6cd9..739221c48 100644 --- a/packages/shared-python/shared/models/database/document.py +++ b/packages/shared-python/shared/models/database/document.py @@ -356,6 +356,12 @@ class DocumentMapUnitIndex(Base): 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) + # rank_bm25 Okapi average IDF for the revision's units (path/content). + # Written at index time so query scoring never rescans all tokens. + average_idf_path: Mapped[float] = mapped_column(Float, nullable=False, default=0.0) + average_idf_content: Mapped[float] = mapped_column( + Float, nullable=False, default=0.0 + ) created_at: Mapped[datetime] = mapped_column( DateTime, default=utc_now_naive, nullable=False ) diff --git a/packages/shared-python/shared/services/retrieval/map_unit_index.py b/packages/shared-python/shared/services/retrieval/map_unit_index.py index 4171861c7..4fff8206a 100644 --- a/packages/shared-python/shared/services/retrieval/map_unit_index.py +++ b/packages/shared-python/shared/services/retrieval/map_unit_index.py @@ -23,6 +23,7 @@ UnitRow, ) from shared.services.retrieval.nav.nav_map_scores import build_score_units +from shared.services.retrieval.nav.persisted_score_load import average_idf_from_unit_dfs from shared.services.retrieval.publication_models import DocumentPublicationScope @@ -85,6 +86,8 @@ def replace_document_map_units( ) persisted_count = 0 token_count = 0 + path_unit_df: Counter[str] = Counter() + content_unit_df: Counter[str] = Counter() 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() @@ -93,6 +96,8 @@ def replace_document_map_units( 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() + path_unit_df.update(set(path_tokens)) + content_unit_df.update(set(content_tokens)) # ``provider.self_units`` already reflects root-asset remount (assets # referenced via ``connect_to`` are moved onto the text section that # embeds them), so this is the same ownership the query-time scorer @@ -139,6 +144,14 @@ def replace_document_map_units( format_version=MAP_UNIT_INDEX_FORMAT_VERSION, unit_count=persisted_count, token_count=token_count, + average_idf_path=average_idf_from_unit_dfs( + unit_count=persisted_count, + token_document_frequency=path_unit_df, + ), + average_idf_content=average_idf_from_unit_dfs( + unit_count=persisted_count, + token_document_frequency=content_unit_df, + ), ) ) 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 c5cd4fbba..577b1b95b 100644 --- a/packages/shared-python/shared/services/retrieval/nav/nav_knowhere.py +++ b/packages/shared-python/shared/services/retrieval/nav/nav_knowhere.py @@ -23,7 +23,6 @@ import os import logging import time -from hashlib import sha256 from dataclasses import dataclass, field from typing import ( Any, @@ -42,7 +41,6 @@ from .nav_address import NavLevel from .nav_hierarchy import NodeMeta from .knowhere_hybrid import ( - PersistedBm25Stats, PersistedScoreCorpus, PersistedScoreUnit, tokenize_query_for_ranker, @@ -207,15 +205,11 @@ 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]] + tuple[tuple[str, str], ...], list[dict[str, object]] ] = {} - self._score_frequency_cache: dict[ - tuple[tuple[str, str], ...], - dict[tuple[str, str], dict[str, int]], + self._score_average_idf_cache: dict[ + tuple[tuple[str, str], ...], tuple[float, float] ] = {} def _connection(self) -> "_SyncConnection": @@ -300,7 +294,17 @@ def load_persisted_score_corpus( allowed_section_ids_by_document: Mapping[str, Sequence[str]], queries: Sequence[str], ) -> Optional[PersistedScoreCorpus]: - """Load query-relevant score inputs when every revision is indexed.""" + """Load query-token score inputs from map-unit tables when indexed. + + Units and lengths come from ``document_map_units``. Frequencies come from + ``document_map_unit_tokens`` filtered to the query tokens. Average IDF + comes from ``document_map_unit_indexes`` (written at index time). + """ + from shared.services.retrieval.nav.persisted_score_load import ( + build_channel_bm25_stats, + combine_average_idf, + ) + revisions = [ (document_id, self._revisions[document_id]) for raw_document_id in document_ids @@ -312,229 +316,141 @@ def load_persisted_score_corpus( revision_params: List[object] = [ value for revision in revisions for value in revision ] + unique_queries = list(dict.fromkeys(str(query) for query in queries)) + query_tokens = list( + dict.fromkeys( + token + for query in unique_queries + for token in tokenize_query_for_ranker(query) + ) + ) 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), + stage_started = time.perf_counter() + try: + cur.execute( + "SELECT indexes.document_id, indexes.job_result_id, " + "indexes.format_version, indexes.unit_count, " + "indexes.average_idf_path, indexes.average_idf_content " + "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, ) - 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) + except Exception as exc: + if getattr(exc, "pgcode", None) in {"42P01", "42703"}: + return None + raise + index_rows = list(cur.fetchall()) _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), + "retrieval map-index load stage=indexes seconds=%.3f rows=%d", + time.perf_counter() - stage_started, + len(index_rows), ) - 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 + if len(index_rows) != len(revisions) or any( + len(row) < 6 or int(row[2]) != _MAP_UNIT_INDEX_FORMAT_VERSION + for row in index_rows ): 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. - # 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. + if revision_key in self._score_average_idf_cache: + average_idf_path, average_idf_content = self._score_average_idf_cache[ + revision_key + ] + else: + average_idf_path = combine_average_idf( + [(float(row[4] or 0.0), int(row[3] or 0)) for row in index_rows] + ) + average_idf_content = combine_average_idf( + [(float(row[5] or 0.0), int(row[3] or 0)) for row in index_rows] + ) + self._score_average_idf_cache[revision_key] = ( + average_idf_path, + average_idf_content, + ) + 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 = [ + allowed_pairs_set = { (document_id, section_id) 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 + } + all_unit_rows = self._score_unit_rows_cache.get(revision_key) + if all_unit_rows is None: + stage_started = time.perf_counter() + cur.execute( + "SELECT units.id, units.document_id, units.unit_id, " + "units.section_id, units.path_token_count, units.content_token_count " + "FROM document_map_units AS units " + f"JOIN (VALUES {values_sql}) AS revisions(document_id, job_result_id) " + "ON units.document_id = revisions.document_id " + "AND units.job_result_id = revisions.job_result_id " + "ORDER BY units.document_id, units.sort_order, units.unit_id", + revision_params, + ) + all_unit_rows = [ + { + "map_unit_id": str(row[0]), + "document_id": str(row[1]), + "unit_id": str(row[2]), + "section_id": str(row[3] or ""), + "path_token_count": int(row[4] or 0), + "content_token_count": int(row[5] or 0), + } + for row in cur.fetchall() ] + expected_unit_count = sum(int(row[3] or 0) for row in index_rows) + if expected_unit_count and len(all_unit_rows) != expected_unit_count: + return None + self._score_unit_rows_cache[revision_key] = all_unit_rows _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, + time.perf_counter() - stage_started, + len(all_unit_rows), + False, ) - 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] + else: + _logger.info( + "retrieval map-index load stage=units seconds=0.000 rows=%d cache_hit=%s", + len(all_unit_rows), + True, ) - ) + + unit_rows = [ + row + for row in all_unit_rows + if (str(row["document_id"]), str(row["section_id"])) in allowed_pairs_set + ] 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 - ] + if unit_rows and 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), - ), + "SELECT units.id, tokens.channel, tokens.token, tokens.frequency " + "FROM document_map_unit_tokens AS tokens " + "JOIN document_map_units AS units ON units.id = tokens.map_unit_id " + f"JOIN (VALUES {values_sql}) AS revisions(document_id, job_result_id) " + "ON units.document_id = revisions.document_id " + "AND units.job_result_id = revisions.job_result_id " + "WHERE tokens.token = ANY(%s) AND tokens.channel = ANY(%s)", + [*revision_params, list(query_tokens), list(_MAP_SCORE_CHANNELS)], ) + allowed_map_unit_ids = {str(row["map_unit_id"]) for row in unit_rows} for map_unit_id, channel, token, frequency in cur.fetchall(): + if str(map_unit_id) not in allowed_map_unit_ids: + continue frequencies.setdefault((str(map_unit_id), str(channel)), {})[ str(token) ] = int(frequency) _logger.info( - "retrieval map-index load stage=frequencies seconds=%.3f units=%d", + "retrieval map-index load stage=frequencies seconds=%.3f units=%d tokens=%d", time.perf_counter() - stage_started, - len(map_unit_ids), + len(unit_rows), + len(query_tokens), ) _logger.info( @@ -542,33 +458,35 @@ def load_persisted_score_corpus( len(unit_rows), len(unique_queries), ) - path_stats = self._load_persisted_bm25_stats( - cur, + path_stats = build_channel_bm25_stats( unit_rows=unit_rows, - map_unit_ids=map_unit_ids, + map_unit_id_field="map_unit_id", + length_field="path_token_count", channel="path", query_tokens=query_tokens, frequencies=frequencies, - length_index=4, + average_idf=average_idf_path, ) - content_stats = self._load_persisted_bm25_stats( - cur, + content_stats = build_channel_bm25_stats( unit_rows=unit_rows, - map_unit_ids=map_unit_ids, + map_unit_id_field="map_unit_id", + length_field="content_token_count", channel="content", query_tokens=query_tokens, frequencies=frequencies, - length_index=5, + average_idf=average_idf_content, ) 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"), {}), + unit_id=str(row["unit_id"]), + path_length=int(row["path_token_count"]), + content_length=int(row["content_token_count"]), + path_frequencies=frequencies.get( + (str(row["map_unit_id"]), "path"), {} + ), content_frequencies=frequencies.get( - (str(row[0]), "content"), {} + (str(row["map_unit_id"]), "content"), {} ), ) for row in unit_rows @@ -579,58 +497,6 @@ def load_persisted_score_corpus( finally: cur.close() - 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: - stage_started = time.perf_counter() - 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 - _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), - document_frequency=document_frequency, - average_idf=average_idf, - ) - def close(self) -> None: if self._conn is not None: self._conn.close() diff --git a/packages/shared-python/shared/services/retrieval/nav/persisted_score_load.py b/packages/shared-python/shared/services/retrieval/nav/persisted_score_load.py new file mode 100644 index 000000000..9ce10c2e1 --- /dev/null +++ b/packages/shared-python/shared/services/retrieval/nav/persisted_score_load.py @@ -0,0 +1,72 @@ +"""Shared helpers for building persisted BM25 corpora from DB rows.""" + +from __future__ import annotations + +import math +from collections.abc import Mapping, Sequence +from typing import Any + +from shared.services.retrieval.nav.knowhere_hybrid import PersistedBm25Stats + + +def average_idf_from_unit_dfs( + *, + unit_count: int, + token_document_frequency: Mapping[str, int], +) -> float: + """rank_bm25 Okapi average IDF over every token that appears in the corpus.""" + if unit_count <= 0 or not token_document_frequency: + return 0.0 + idfs = [ + math.log(unit_count - frequency + 0.5) - math.log(frequency + 0.5) + for frequency in token_document_frequency.values() + if frequency > 0 + ] + if not idfs: + return 0.0 + return sum(idfs) / len(idfs) + + +def combine_average_idf(parts: Sequence[tuple[float, int]]) -> float: + """Unit-count-weighted mean of per-revision average IDF values.""" + total_units = sum(int(unit_count) for _average, unit_count in parts) + if total_units <= 0: + return 0.0 + return ( + sum(float(average) * int(unit_count) for average, unit_count in parts) + / total_units + ) + + +def build_channel_bm25_stats( + *, + unit_rows: Sequence[Mapping[str, Any]], + map_unit_id_field: str, + length_field: str, + channel: str, + query_tokens: Sequence[str], + frequencies: Mapping[tuple[str, str], Mapping[str, int]], + average_idf: float, +) -> PersistedBm25Stats: + """Build channel stats from already-fetched unit rows and query-token freqs.""" + lengths = [ + int(row[length_field]) for row in unit_rows if int(row[length_field]) > 0 + ] + document_count = len(lengths) + document_frequency = { + token: sum( + 1 + for row in unit_rows + if frequencies.get((str(row[map_unit_id_field]), channel), {}).get( + token, 0 + ) + > 0 + ) + for token in query_tokens + } + return PersistedBm25Stats( + document_count=document_count, + total_length=sum(lengths), + document_frequency=document_frequency, + average_idf=float(average_idf), + ) diff --git a/packages/shared-python/shared/services/retrieval/search/map_unit_discovery.py b/packages/shared-python/shared/services/retrieval/search/map_unit_discovery.py index 211ab286b..38b9762d1 100644 --- a/packages/shared-python/shared/services/retrieval/search/map_unit_discovery.py +++ b/packages/shared-python/shared/services/retrieval/search/map_unit_discovery.py @@ -30,12 +30,15 @@ normalize_chunk_type, ) from shared.services.retrieval.nav.knowhere_hybrid import ( - PersistedBm25Stats, PersistedScoreCorpus, PersistedScoreUnit, score_persisted_corpus_many, tokenize_query_for_ranker, ) +from shared.services.retrieval.nav.persisted_score_load import ( + build_channel_bm25_stats, + combine_average_idf, +) from shared.services.retrieval.search.scoring import normalize_row_scores from shared.services.retrieval.search.section_filters import is_excluded_section from shared.services.retrieval.settings import ASSET_CHUNK_TYPES @@ -209,17 +212,19 @@ async def map_unit_discovery( if not unit_rows: return DiscoveryResult(status="discovery_done", payload={"fused_rows": []}) - map_unit_ids = [row["map_unit_id"] for row in unit_rows] frequency_result = await db.execute( text( - "SELECT map_unit_id, channel, token, frequency " - "FROM document_map_unit_tokens " - "WHERE map_unit_id = ANY(:unit_ids) " - "AND channel = ANY(:channels) " - "AND token = ANY(:tokens)" + cte + + """ + SELECT tokens.map_unit_id, tokens.channel, tokens.token, tokens.frequency + FROM document_map_unit_tokens AS tokens + JOIN scoped_units ON scoped_units.map_unit_id = tokens.map_unit_id + WHERE tokens.channel = ANY(:channels) + AND tokens.token = ANY(:tokens) + """ ), { - "unit_ids": map_unit_ids, + **params, "channels": list(_MAP_SCORE_CHANNELS), "tokens": query_tokens, }, @@ -230,23 +235,53 @@ async def map_unit_discovery( int(frequency) ) - path_stats = await _build_bm25_stats( - db, + index_result = await db.execute( + text( + cte + + """ + SELECT indexes.average_idf_path, indexes.average_idf_content, + indexes.unit_count + FROM document_map_unit_indexes AS indexes + JOIN ( + SELECT DISTINCT document_id, job_result_id FROM scoped_units + ) AS scoped_revisions + ON indexes.document_id = scoped_revisions.document_id + AND indexes.job_result_id = scoped_revisions.job_result_id + """ + ), + params, + ) + index_parts = [ + (float(path_idf or 0.0), float(content_idf or 0.0), int(unit_count or 0)) + for path_idf, content_idf, unit_count in index_result.all() + ] + average_idf_path = combine_average_idf( + [(path_idf, unit_count) for path_idf, _content_idf, unit_count in index_parts] + ) + average_idf_content = combine_average_idf( + [ + (content_idf, unit_count) + for _path_idf, content_idf, unit_count in index_parts + ] + ) + + path_stats = build_channel_bm25_stats( unit_rows=unit_rows, - map_unit_ids=map_unit_ids, + map_unit_id_field="map_unit_id", + length_field="path_token_count", channel="path", query_tokens=query_tokens, frequencies=frequencies, - length_field="path_token_count", + average_idf=average_idf_path, ) - content_stats = await _build_bm25_stats( - db, + content_stats = build_channel_bm25_stats( unit_rows=unit_rows, - map_unit_ids=map_unit_ids, + map_unit_id_field="map_unit_id", + length_field="content_token_count", channel="content", query_tokens=query_tokens, frequencies=frequencies, - length_field="content_token_count", + average_idf=average_idf_content, ) corpus = PersistedScoreCorpus( @@ -310,58 +345,6 @@ async def map_unit_discovery( return DiscoveryResult(status="error", error=str(exc), latency_ms=latency) -async def _build_bm25_stats( - session: AsyncSession, - *, - unit_rows: list[dict[str, Any]], - map_unit_ids: list[str], - channel: str, - query_tokens: list[str], - frequencies: dict[tuple[str, str], dict[str, int]], - length_field: str, -) -> PersistedBm25Stats: - lengths = [ - int(row[length_field]) for row in unit_rows if int(row[length_field]) > 0 - ] - document_count = len(lengths) - document_frequency = { - token: sum( - 1 - for row in unit_rows - if frequencies.get((row["map_unit_id"], 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: - result = await session.execute( - text( - "SELECT COALESCE(AVG(LN((:document_count - 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(:unit_ids) AND channel = :channel " - "GROUP BY token) AS frequencies" - ), - { - "document_count": document_count, - "unit_ids": map_unit_ids, - "channel": channel, - }, - ) - row = result.first() - 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 _as_metadata_dict(value: object) -> dict[str, Any]: if isinstance(value, dict): return value diff --git a/packages/shared-python/shared/services/retrieval/serving_manifest.py b/packages/shared-python/shared/services/retrieval/serving_manifest.py index 5fe979e08..888113202 100644 --- a/packages/shared-python/shared/services/retrieval/serving_manifest.py +++ b/packages/shared-python/shared/services/retrieval/serving_manifest.py @@ -58,14 +58,6 @@ def build_revision_serving_payload( ) ) ) - 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 } @@ -121,18 +113,6 @@ def build_revision_serving_payload( } 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, } @@ -143,7 +123,12 @@ def build_revision_statistics_payload( *, scope: DocumentPublicationScope, ) -> dict[str, Any]: - """Build compressed scoring contributions for one revision.""" + """Build compressed scoring contributions for one revision. + + Stores aggregate token frequencies for namespace statistics rebuild. + Per-unit frequencies stay in ``document_map_unit_tokens`` and are loaded + at query time by the query tokens only. + """ units = list( db.scalars( select(DocumentMapUnit) @@ -153,9 +138,8 @@ def build_revision_statistics_payload( ) 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( + for _map_unit_id, channel, token, frequency in db.execute( select( DocumentMapUnitToken.map_unit_id, DocumentMapUnitToken.channel, @@ -164,15 +148,12 @@ def build_revision_statistics_payload( ).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 + if channel_key not in frequencies: + continue + token_key = str(token) + frequencies[channel_key][token_key] = frequencies[channel_key].get( + token_key, 0 + ) + int(frequency) return { "document_id": scope.document_id, "job_result_id": scope.job_result_id, @@ -182,7 +163,6 @@ def build_revision_statistics_payload( int(unit.content_token_count or 0) for unit in units ), "token_frequencies": frequencies, - "unit_frequencies": unit_frequencies, } From 496501b77d66e5f173c038674ebd4158b092095f Mon Sep 17 00:00:00 2001 From: chengke <404835780@qq.com> Date: Mon, 31 Aug 2026 18:21:20 +0800 Subject: [PATCH 5/6] chore: drop unused serving_manifest import in nav_knowhere Co-authored-by: Cursor --- .../shared-python/shared/services/retrieval/nav/nav_knowhere.py | 1 - 1 file changed, 1 deletion(-) 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 577b1b95b..716ff1ec7 100644 --- a/packages/shared-python/shared/services/retrieval/nav/nav_knowhere.py +++ b/packages/shared-python/shared/services/retrieval/nav/nav_knowhere.py @@ -45,7 +45,6 @@ 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). From 5fee16666a8c30a65d43b3e96ccb0938f415b1a4 Mon Sep 17 00:00:00 2001 From: chengke <404835780@qq.com> Date: Mon, 31 Aug 2026 18:37:47 +0800 Subject: [PATCH 6/6] feat: add read-only --check for serving-index fallback readiness Co-authored-by: Cursor --- apps/api/scripts/backfill_map_unit_indexes.py | 214 +++++++++++++++++- 1 file changed, 213 insertions(+), 1 deletion(-) diff --git a/apps/api/scripts/backfill_map_unit_indexes.py b/apps/api/scripts/backfill_map_unit_indexes.py index be2b4dbbf..6eed78eff 100644 --- a/apps/api/scripts/backfill_map_unit_indexes.py +++ b/apps/api/scripts/backfill_map_unit_indexes.py @@ -6,6 +6,10 @@ derived tables leave them empty intentionally. Run this command after deployment with ``--apply`` so each revision is rebuilt and committed independently; without ``--apply`` it is a read-only inventory. + +Use ``--check`` after backfill to verify whether query-time snapshot +fallbacks (manifest_merge / table_scan) would still fire, and whether +map-unit indexes are complete for scoring. """ # ruff: noqa: E402 @@ -15,6 +19,8 @@ import argparse import os import sys +from collections import defaultdict +from dataclasses import dataclass from pathlib import Path @@ -45,7 +51,12 @@ def _bootstrap_python_path() -> None: from sqlalchemy import select from shared.core.database_sync import get_sync_session_factory -from shared.models.database.document import Document +from shared.models.database.document import ( + Document, + DocumentMapUnitIndex, + RetrievalNamespaceMapSnapshot, + RetrievalServingRevisionManifest, +) from shared.services.retrieval.map_unit_index import replace_document_map_units from shared.services.retrieval.namespace_map_snapshot import ( patch_namespace_map_snapshot, @@ -56,6 +67,7 @@ def _bootstrap_python_path() -> None: lock_namespace_generation, ) from shared.services.retrieval.serving_manifest import ( + decode_serving_manifest, persist_revision_serving_state, rebuild_namespace_serving_statistics, ) @@ -70,6 +82,14 @@ def _build_parser() -> argparse.ArgumentParser: action="store_true", help="Build and commit each current revision index.", ) + parser.add_argument( + "--check", + action="store_true", + help=( + "Read-only: report whether snapshot fallbacks would still fire " + "and whether map-unit indexes are complete. Exit 1 if not ready." + ), + ) parser.add_argument( "--document-id", default="", help="Limit the backfill to one document." ) @@ -90,6 +110,192 @@ def _load_documents(document_id: str) -> list[Document]: return list(db.scalars(statement).all()) +@dataclass(frozen=True) +class NamespaceFallbackReport: + user_id: str + namespace: str + active_docs: int + snapshot_status: str + missing_from_snapshot: int + missing_map_index: int + missing_revision_manifest: int + suspicious_zero_idf: int + would_hit_snapshot_fallback: bool + scoring_incomplete: bool + + @property + def ready(self) -> bool: + return not self.would_hit_snapshot_fallback and not self.scoring_incomplete + + +def check_fallback_readiness(*, document_id: str = "") -> list[NamespaceFallbackReport]: + """Inspect active revisions for snapshot coverage and map-unit indexes.""" + documents = _load_documents(document_id) + by_scope: dict[tuple[str, str], list[Document]] = defaultdict(list) + for document in documents: + by_scope[(document.user_id, document.namespace)].append(document) + + session_factory = get_sync_session_factory() + reports: list[NamespaceFallbackReport] = [] + with session_factory() as db: + for (user_id, namespace), scoped_docs in sorted( + by_scope.items(), key=lambda item: (item[0][0], item[0][1]) + ): + revisions = [ + (doc.document_id, str(doc.current_job_result_id)) + for doc in scoped_docs + if doc.current_job_result_id + ] + if not revisions: + continue + + snapshot = db.execute( + select(RetrievalNamespaceMapSnapshot) + .where(RetrievalNamespaceMapSnapshot.user_id == user_id) + .where(RetrievalNamespaceMapSnapshot.namespace == namespace) + ).scalar_one_or_none() + + snapshot_documents: dict[str, object] | None = None + snapshot_status = "missing" + if snapshot is None: + snapshot_status = "missing" + else: + try: + payload = decode_serving_manifest( + bytes(snapshot.payload_zlib), + checksum=str(snapshot.checksum), + format_version=int(snapshot.format_version), + ) + decoded = payload.get("documents") + if isinstance(decoded, dict): + snapshot_documents = decoded + snapshot_status = "ok" + else: + snapshot_status = "corrupt" + except (TypeError, ValueError): + snapshot_status = "corrupt" + + missing_from_snapshot = 0 + if snapshot_documents is None: + missing_from_snapshot = len(revisions) + else: + for document_id_value, job_result_id in revisions: + entry = snapshot_documents.get(document_id_value) + if ( + not isinstance(entry, dict) + or str(entry.get("job_result_id") or "") != job_result_id + ): + missing_from_snapshot += 1 + if missing_from_snapshot and snapshot_status == "ok": + snapshot_status = "stale" + + index_rows = list( + db.execute( + select( + DocumentMapUnitIndex.document_id, + DocumentMapUnitIndex.job_result_id, + DocumentMapUnitIndex.unit_count, + DocumentMapUnitIndex.average_idf_path, + DocumentMapUnitIndex.average_idf_content, + ).where( + DocumentMapUnitIndex.document_id.in_( + [document_id_value for document_id_value, _ in revisions] + ) + ) + ).all() + ) + index_by_revision = { + (str(document_id_value), str(job_result_id)): ( + int(unit_count or 0), + float(average_idf_path or 0.0), + float(average_idf_content or 0.0), + ) + for document_id_value, job_result_id, unit_count, average_idf_path, average_idf_content in index_rows + } + missing_map_index = 0 + suspicious_zero_idf = 0 + for document_id_value, job_result_id in revisions: + stats = index_by_revision.get((document_id_value, job_result_id)) + if stats is None: + missing_map_index += 1 + continue + unit_count, average_idf_path, average_idf_content = stats + if ( + unit_count > 0 + and average_idf_path == 0.0 + and average_idf_content == 0.0 + ): + suspicious_zero_idf += 1 + + manifest_rows = list( + db.execute( + select( + RetrievalServingRevisionManifest.document_id, + RetrievalServingRevisionManifest.job_result_id, + ).where( + RetrievalServingRevisionManifest.document_id.in_( + [document_id_value for document_id_value, _ in revisions] + ) + ) + ).all() + ) + manifest_keys = { + (str(document_id_value), str(job_result_id)) + for document_id_value, job_result_id in manifest_rows + } + missing_revision_manifest = sum( + 1 + for document_id_value, job_result_id in revisions + if (document_id_value, job_result_id) not in manifest_keys + ) + + would_hit_snapshot_fallback = ( + snapshot_status != "ok" or missing_from_snapshot > 0 + ) + scoring_incomplete = missing_map_index > 0 or suspicious_zero_idf > 0 + reports.append( + NamespaceFallbackReport( + user_id=user_id, + namespace=namespace, + active_docs=len(revisions), + snapshot_status=snapshot_status, + missing_from_snapshot=missing_from_snapshot, + missing_map_index=missing_map_index, + missing_revision_manifest=missing_revision_manifest, + suspicious_zero_idf=suspicious_zero_idf, + would_hit_snapshot_fallback=would_hit_snapshot_fallback, + scoring_incomplete=scoring_incomplete, + ) + ) + return reports + + +def print_fallback_check(reports: list[NamespaceFallbackReport]) -> int: + """Print readiness report. Returns process exit code (0 ready, 1 not).""" + if not reports: + print("check: no active documents found") + return 0 + + failed = 0 + for report in reports: + status = "READY" if report.ready else "NOT_READY" + if not report.ready: + failed += 1 + print( + f"check status={status} user={report.user_id} namespace={report.namespace} " + f"active_docs={report.active_docs} snapshot={report.snapshot_status} " + f"missing_from_snapshot={report.missing_from_snapshot} " + f"missing_map_index={report.missing_map_index} " + f"missing_revision_manifest={report.missing_revision_manifest} " + f"suspicious_zero_idf={report.suspicious_zero_idf} " + f"would_hit_snapshot_fallback={report.would_hit_snapshot_fallback} " + f"scoring_incomplete={report.scoring_incomplete}" + ) + ready_count = len(reports) - failed + print(f"check namespaces_ready={ready_count}/{len(reports)}") + return 1 if failed else 0 + + def backfill_map_unit_indexes(*, apply: bool, document_id: str = "") -> int: documents = _load_documents(document_id) if not apply: @@ -157,6 +363,12 @@ def backfill_map_unit_indexes(*, apply: bool, document_id: str = "") -> int: def main() -> None: arguments = _build_parser().parse_args() + if arguments.check: + if arguments.apply: + raise SystemExit("use either --check or --apply, not both") + reports = check_fallback_readiness(document_id=str(arguments.document_id)) + raise SystemExit(print_fallback_check(reports)) + count = backfill_map_unit_indexes( apply=bool(arguments.apply), document_id=str(arguments.document_id) )