diff --git a/apps/api/alembic/versions/0c1d2e3f4a5b_add_chunk_revision_section_order_index.py b/apps/api/alembic/versions/0c1d2e3f4a5b_add_chunk_revision_section_order_index.py new file mode 100644 index 000000000..300dc99ca --- /dev/null +++ b/apps/api/alembic/versions/0c1d2e3f4a5b_add_chunk_revision_section_order_index.py @@ -0,0 +1,43 @@ +"""Add the index used by lazy map-nav section loads.""" + +from __future__ import annotations + +from alembic import op + + +revision = "0c1d2e3f4a5b" +down_revision = "fbf0c1d2e3f4" +branch_labels = None +depends_on = None + +_INDEX_NAME = "idx_document_chunks_revision_section_order" + + +def upgrade() -> None: + external_transaction = bool( + op.get_context().opts.get("knowhere_external_transaction", False) + ) + if external_transaction: + op.execute( + f"CREATE INDEX IF NOT EXISTS {_INDEX_NAME} " + "ON document_chunks " + "(document_id, job_result_id, section_id, sort_order, chunk_id, id)" + ) + return + with op.get_context().autocommit_block(): + op.execute( + f"CREATE INDEX CONCURRENTLY IF NOT EXISTS {_INDEX_NAME} " + "ON document_chunks " + "(document_id, job_result_id, section_id, sort_order, chunk_id, id)" + ) + + +def downgrade() -> None: + external_transaction = bool( + op.get_context().opts.get("knowhere_external_transaction", False) + ) + if external_transaction: + op.execute(f"DROP INDEX IF EXISTS {_INDEX_NAME}") + return + with op.get_context().autocommit_block(): + op.execute(f"DROP INDEX CONCURRENTLY IF EXISTS {_INDEX_NAME}") diff --git a/apps/api/alembic/versions/1d2e3f4a5b6c_add_document_map_units.py b/apps/api/alembic/versions/1d2e3f4a5b6c_add_document_map_units.py new file mode 100644 index 000000000..79aa7583c --- /dev/null +++ b/apps/api/alembic/versions/1d2e3f4a5b6c_add_document_map_units.py @@ -0,0 +1,122 @@ +"""Add revision-pinned map-nav score units.""" + +from __future__ import annotations + +import sqlalchemy as sa +from alembic import op + +revision = "1d2e3f4a5b6c" +down_revision = "0c1d2e3f4a5b" +branch_labels = None +depends_on = None + + +def upgrade() -> None: + bind = op.get_bind() + inspector = sa.inspect(bind) + if not inspector.has_table("document_map_unit_indexes"): + op.create_table( + "document_map_unit_indexes", + sa.Column("id", sa.String(length=100), nullable=False), + sa.Column("document_id", sa.String(length=36), nullable=False), + sa.Column("job_result_id", sa.String(length=36), nullable=False), + sa.Column("format_version", sa.Integer(), nullable=False), + sa.Column("unit_count", sa.Integer(), nullable=False), + sa.Column("token_count", sa.Integer(), nullable=False), + sa.Column("created_at", sa.DateTime(), nullable=False), + sa.ForeignKeyConstraint( + ["document_id"], ["documents.document_id"], ondelete="CASCADE" + ), + sa.ForeignKeyConstraint( + ["job_result_id"], ["job_results.id"], ondelete="CASCADE" + ), + sa.PrimaryKeyConstraint("id"), + sa.UniqueConstraint( + "document_id", + "job_result_id", + name="uq_document_map_unit_indexes_revision", + ), + ) + if not inspector.has_table("document_map_units"): + op.create_table( + "document_map_units", + sa.Column("id", sa.String(length=160), nullable=False), + sa.Column("document_id", sa.String(length=36), nullable=False), + sa.Column("job_result_id", sa.String(length=36), nullable=False), + sa.Column("unit_id", sa.String(length=128), nullable=False), + sa.Column("section_id", sa.String(length=36), nullable=False), + sa.Column("unit_kind", sa.String(length=32), nullable=False), + sa.Column("path_token_count", sa.Integer(), nullable=False), + sa.Column("content_token_count", sa.Integer(), nullable=False), + sa.Column("term_search_text_lower", sa.Text(), nullable=False), + sa.Column("sort_order", sa.Integer(), nullable=False), + sa.Column("created_at", sa.DateTime(), nullable=False), + sa.ForeignKeyConstraint( + ["document_id"], ["documents.document_id"], ondelete="CASCADE" + ), + sa.ForeignKeyConstraint( + ["job_result_id"], ["job_results.id"], ondelete="CASCADE" + ), + sa.PrimaryKeyConstraint("id"), + ) + if not inspector.has_table("document_map_unit_tokens"): + op.create_table( + "document_map_unit_tokens", + sa.Column("id", sa.String(length=36), nullable=False), + sa.Column("map_unit_id", sa.String(length=160), nullable=False), + sa.Column("channel", sa.String(length=16), nullable=False), + sa.Column("token", sa.Text(), nullable=False), + sa.Column("token_hash", sa.String(length=64), nullable=False), + sa.Column("frequency", sa.Integer(), nullable=False), + sa.ForeignKeyConstraint( + ["map_unit_id"], ["document_map_units.id"], ondelete="CASCADE" + ), + sa.PrimaryKeyConstraint("id"), + ) + inspector = sa.inspect(bind) + indexes = { + item["name"] for item in inspector.get_indexes("document_map_unit_indexes") + } + if "idx_document_map_unit_indexes_revision" not in indexes: + op.create_index( + "idx_document_map_unit_indexes_revision", + "document_map_unit_indexes", + ["document_id", "job_result_id"], + ) + indexes = {item["name"] for item in inspector.get_indexes("document_map_units")} + if "idx_document_map_units_revision_order" not in indexes: + op.create_index( + "idx_document_map_units_revision_order", + "document_map_units", + ["document_id", "job_result_id", "sort_order", "unit_id"], + ) + if "idx_document_map_units_section" not in indexes: + op.create_index( + "idx_document_map_units_section", "document_map_units", ["section_id"] + ) + indexes = { + item["name"] for item in inspector.get_indexes("document_map_unit_tokens") + } + if "idx_document_map_unit_tokens_lookup" not in indexes: + op.create_index( + "idx_document_map_unit_tokens_lookup", + "document_map_unit_tokens", + ["channel", "token_hash", "map_unit_id"], + ) + if "idx_document_map_unit_tokens_unit" not in indexes: + op.create_index( + "idx_document_map_unit_tokens_unit", + "document_map_unit_tokens", + ["map_unit_id", "channel"], + ) + + +def downgrade() -> None: + existing_tables = set(sa.inspect(op.get_bind()).get_table_names()) + for table_name in ( + "document_map_unit_tokens", + "document_map_units", + "document_map_unit_indexes", + ): + if table_name in existing_tables: + op.drop_table(table_name) diff --git a/apps/api/scripts/backfill_map_unit_indexes.py b/apps/api/scripts/backfill_map_unit_indexes.py new file mode 100644 index 000000000..78a88f6c0 --- /dev/null +++ b/apps/api/scripts/backfill_map_unit_indexes.py @@ -0,0 +1,97 @@ +"""Backfill persisted MAP-NAV lexical indexes for existing revisions. + +The migration creates empty derived tables intentionally. Run this command +after deployment with ``--apply`` so each revision is rebuilt and committed +independently; without ``--apply`` it is a read-only inventory. +""" + +# ruff: noqa: E402 + +from __future__ import annotations + +import argparse +import os +import sys +from pathlib import Path + + +def _bootstrap_python_path() -> None: + api_root = Path(__file__).resolve().parents[1] + repo_root = api_root.parents[1] + shared_root = repo_root / "packages" / "shared-python" + for path in (api_root, shared_root): + value = os.fspath(path) + if value not in sys.path: + sys.path.insert(0, value) + + +_bootstrap_python_path() + +from sqlalchemy import select + +from shared.core.database_sync import get_sync_session_factory +from shared.models.database.document import Document +from shared.services.retrieval.map_unit_index import replace_document_map_units +from shared.services.retrieval.publication_models import DocumentPublicationScope + + +def _build_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser( + description="Backfill MAP-NAV indexes for current document revisions." + ) + parser.add_argument( + "--apply", + action="store_true", + help="Build and commit each current revision index.", + ) + parser.add_argument("--document-id", default="", help="Limit the backfill to one document.") + return parser + + +def _load_documents(document_id: str) -> list[Document]: + session_factory = get_sync_session_factory() + with session_factory() as db: + statement = select(Document).where(Document.current_job_result_id.is_not(None)) + normalized_document_id = document_id.strip() + if normalized_document_id: + statement = statement.where(Document.document_id == normalized_document_id) + return list(db.scalars(statement).all()) + + +def backfill_map_unit_indexes(*, apply: bool, document_id: str = "") -> int: + documents = _load_documents(document_id) + if not apply: + for document in documents: + print(f"would backfill document={document.document_id} revision={document.current_job_result_id}") + return len(documents) + + session_factory = get_sync_session_factory() + for document in documents: + job_result_id = document.current_job_result_id + if not job_result_id: + continue + scope = DocumentPublicationScope( + user_id=document.user_id, + namespace=document.namespace, + document_id=document.document_id, + job_result_id=job_result_id, + source_file_name=str(document.source_file_name or ""), + ) + with session_factory() as db: + replace_document_map_units(db, scope=scope) + db.commit() + print(f"backfilled document={document.document_id} revision={job_result_id}") + return len(documents) + + +def main() -> None: + arguments = _build_parser().parse_args() + count = backfill_map_unit_indexes( + apply=bool(arguments.apply), document_id=str(arguments.document_id) + ) + action = "backfilled" if arguments.apply else "found" + print(f"{action} revisions={count}") + + +if __name__ == "__main__": + main() diff --git a/apps/api/tests/contract/test_retrieval_lazy_snapshot_quality_contract.py b/apps/api/tests/contract/test_retrieval_lazy_snapshot_quality_contract.py new file mode 100644 index 000000000..ae6d78ed3 --- /dev/null +++ b/apps/api/tests/contract/test_retrieval_lazy_snapshot_quality_contract.py @@ -0,0 +1,378 @@ +from __future__ import annotations + +from dataclasses import dataclass +from collections.abc import Mapping, Sequence +from typing import Any + +from shared.services.retrieval.nav.nav_hierarchy import ProviderToolSpace +from shared.services.retrieval.nav.nav_knowhere import ( + KnowhereProvider, + LazyKnowhereProvider, + NamespaceKnowhereProvider, + SectionRow, + UnitRow, + knowhere_database_url, +) +from shared.services.retrieval.nav.nav_map_scores import ( + build_score_units, + compute_corpus_map_and_unit_scores, + compute_corpus_map_and_unit_scores_many, +) +from shared.services.retrieval.nav.knowhere_hybrid import ( + ScoreUnitRow, + score_rows_hybrid_all, + score_unit_stream_hybrid_all, + score_unit_stream_hybrid_many, +) + + +@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 + + def load_section_units( + self, + document_id: str, + section_id: str, + extra_chunk_ids: Sequence[str] = (), + ) -> list[UnitRow]: + del document_id + units = list(self.units_by_section.get(section_id, ())) + known = {unit.chunk_id for unit in units} + for units_in_section in self.units_by_section.values(): + for unit in units_in_section: + if unit.chunk_id in extra_chunk_ids and unit.chunk_id not in known: + units.append(unit) + return units + + def close(self) -> None: + return None + + +def _providers() -> tuple[ProviderToolSpace, ProviderToolSpace, _FakeChunkStore]: + sections = [ + SectionRow("root", None, "Root", "Root", 0, "", 0), + SectionRow("section", "root", "Root / Section", "Section", 1, "", 1), + SectionRow("leaf", "section", "Root / Section / Leaf", "Leaf", 2, "", 2), + ] + text = UnitRow( + "duplicate-chunk", + "leaf", + "text", + "alpha retrieval evidence", + 1, + metadata={"connect_to": [{"target": "asset-1", "relation": "embeds"}]}, + ) + asset = UnitRow( + "asset-1", + "root", + "image", + "", + 2, + file_path="images/asset.png", + metadata={"summary": "supporting image"}, + ) + eager = NamespaceKnowhereProvider( + [KnowhereProvider(doc_id="doc", sections=sections, units=[text, asset])], + titles={"doc": "document"}, + ) + store = _FakeChunkStore({"root": [asset], "leaf": [text]}) + lazy = NamespaceKnowhereProvider( + [ + LazyKnowhereProvider( + doc_id="doc", + sections=sections, + chunk_store=store, + known_chunk_ids=[text.chunk_id, asset.chunk_id], + root_asset_ids=[asset.chunk_id], + remounted_assets_by_section={"leaf": [asset.chunk_id]}, + ) + ], + titles={"doc": "document"}, + chunk_owner_by_id={"duplicate-chunk": "doc", "asset-1": "doc"}, + ) + return ProviderToolSpace(eager), ProviderToolSpace(lazy), store + + +def _multi_document_providers() -> tuple[ + ProviderToolSpace, + ProviderToolSpace, + _FakeChunkStore, +]: + first_sections = [ + SectionRow("root-a", None, "Root A", "Root A", 0, "", 0), + SectionRow("leaf-a", "root-a", "Root A / Leaf A", "Leaf A", 1, "", 1), + ] + second_sections = [ + SectionRow("root-b", None, "Root B", "Root B", 0, "", 0), + SectionRow("leaf-b", "root-b", "Root B / Leaf B", "Leaf B", 1, "", 1), + ] + first_unit = UnitRow("chunk-a", "leaf-a", "text", "alpha evidence", 1) + second_unit = UnitRow("chunk-b", "leaf-b", "text", "beta evidence", 1) + eager = NamespaceKnowhereProvider( + [ + KnowhereProvider( + doc_id="doc-a", + sections=first_sections, + units=[first_unit], + ), + KnowhereProvider( + doc_id="doc-b", + sections=second_sections, + units=[second_unit], + ), + ], + titles={"doc-a": "Document A", "doc-b": "Document B"}, + ) + store = _FakeChunkStore( + { + "leaf-a": [first_unit], + "leaf-b": [second_unit], + } + ) + lazy = NamespaceKnowhereProvider( + [ + LazyKnowhereProvider( + doc_id="doc-a", + sections=first_sections, + chunk_store=store, + known_chunk_ids=[first_unit.chunk_id], + ), + LazyKnowhereProvider( + doc_id="doc-b", + sections=second_sections, + chunk_store=store, + known_chunk_ids=[second_unit.chunk_id], + ), + ], + titles={"doc-a": "Document A", "doc-b": "Document B"}, + chunk_owner_by_id={"chunk-a": "doc-a", "chunk-b": "doc-b"}, + ) + return ProviderToolSpace(eager), ProviderToolSpace(lazy), store + + +def test_lazy_provider_preserves_score_units_and_scores() -> None: + eager, lazy, store = _providers() + + assert build_score_units(eager, "doc") == build_score_units(lazy, "doc") + assert compute_corpus_map_and_unit_scores( + eager, doc_ids=["doc"], query="alpha retrieval" + ) == compute_corpus_map_and_unit_scores( + lazy, doc_ids=["doc"], query="alpha retrieval" + ) + + lazy_provider = lazy._provider + 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", + "asset-1", + ] + + +def test_streaming_scorer_preserves_exact_eager_scores() -> None: + rows: list[ScoreUnitRow] = [ + { + "chunk_id": "unit-a", + "path_search_text": "root alpha", + "content_search_text": "alpha alpha evidence", + "term_search_text": "alpha alpha evidence root", + }, + { + "chunk_id": "unit-b", + "path_search_text": "root beta", + "content_search_text": "beta evidence", + "term_search_text": "beta evidence root", + }, + { + "chunk_id": "unit-c", + "path_search_text": "root common", + "content_search_text": "common evidence", + "term_search_text": "common evidence root", + }, + ] + eager_rows: list[dict[str, Any]] = [dict(row) for row in rows] + eager_scores = { + str(row["chunk_id"]): float(row["score"]) + for row in score_rows_hybrid_all(eager_rows, "alpha evidence") + } + replay_count: int = 0 + + def unit_factory() -> Sequence[ScoreUnitRow]: + nonlocal replay_count + replay_count += 1 + return rows + + assert score_unit_stream_hybrid_all(unit_factory, "alpha evidence") == eager_scores + assert replay_count == 1 + + +def test_streaming_scorer_preserves_duplicate_id_eager_semantics() -> None: + rows: list[ScoreUnitRow] = [ + { + "chunk_id": "duplicate", + "path_search_text": "alpha", + "content_search_text": "alpha", + "term_search_text": "alpha", + }, + { + "chunk_id": "duplicate", + "path_search_text": "beta", + "content_search_text": "beta", + "term_search_text": "beta", + }, + { + "chunk_id": "other", + "path_search_text": "alpha beta", + "content_search_text": "alpha beta", + "term_search_text": "alpha beta", + }, + ] + eager_rows: list[dict[str, Any]] = [dict(row) for row in rows] + eager_scores = { + str(row["chunk_id"]): float(row["score"]) + for row in score_rows_hybrid_all(eager_rows, "alpha beta") + } + + assert score_unit_stream_hybrid_all(lambda: rows, "alpha beta") == eager_scores + + +def test_streaming_scorer_scores_multiple_queries_with_one_corpus_read() -> None: + rows: list[ScoreUnitRow] = [ + { + "chunk_id": "unit-a", + "path_search_text": "root alpha", + "content_search_text": "alpha alpha evidence", + "term_search_text": "alpha alpha evidence root", + }, + { + "chunk_id": "unit-b", + "path_search_text": "root beta", + "content_search_text": "beta evidence", + "term_search_text": "beta evidence root", + }, + { + "chunk_id": "unit-c", + "path_search_text": "root common", + "content_search_text": "common evidence", + "term_search_text": "common evidence root", + }, + ] + queries: list[str] = ["alpha evidence", "beta evidence"] + expected = { + query: score_unit_stream_hybrid_all(lambda: rows, query) + for query in queries + } + read_count: int = 0 + + def unit_factory() -> Sequence[ScoreUnitRow]: + nonlocal read_count + read_count += 1 + return rows + + assert score_unit_stream_hybrid_many(unit_factory, queries) == expected + assert read_count == 1 + + +def test_corpus_map_scores_multiple_queries_with_one_lazy_load() -> None: + eager, lazy, store = _providers() + queries: list[str] = ["alpha retrieval", "supporting image"] + expected = { + query: compute_corpus_map_and_unit_scores( + eager, + doc_ids=["doc"], + query=query, + ) + for query in queries + } + + store.document_loads = 0 + store.batch_loads = 0 + actual = compute_corpus_map_and_unit_scores_many( + lazy, + doc_ids=["doc"], + queries=queries, + ) + + assert actual == expected + assert store.batch_loads == 1 + assert store.document_loads == 0 + + +def test_corpus_map_batches_multiple_documents_without_score_drift() -> None: + eager, lazy, store = _multi_document_providers() + queries: list[str] = ["alpha evidence", "beta evidence"] + expected = compute_corpus_map_and_unit_scores_many( + eager, + doc_ids=["doc-a", "doc-b"], + queries=queries, + ) + + actual = compute_corpus_map_and_unit_scores_many( + lazy, + doc_ids=["doc-a", "doc-b"], + queries=queries, + ) + + assert actual == expected + assert store.batch_loads == 1 + assert store.document_loads == 0 + + +def test_native_chunk_store_strips_async_driver_from_database_url( + monkeypatch: Any, +) -> None: + monkeypatch.setenv( + "DATABASE_URL", + "postgresql+asyncpg://prod-user:prod-password@db.example/knowhere", + ) + monkeypatch.delenv("KNOWHERE_DATABASE_URL", raising=False) + + assert ( + knowhere_database_url() + == "postgresql://prod-user:prod-password@db.example/knowhere" + ) diff --git a/apps/api/tests/contract/test_retrieval_map_unit_index_contract.py b/apps/api/tests/contract/test_retrieval_map_unit_index_contract.py new file mode 100644 index 000000000..2f2027b01 --- /dev/null +++ b/apps/api/tests/contract/test_retrieval_map_unit_index_contract.py @@ -0,0 +1,423 @@ +from __future__ import annotations + +from collections.abc import Callable, Mapping, Sequence +from contextlib import AbstractAsyncContextManager +from uuid import uuid4 + +from httpx import AsyncClient +from sqlalchemy import delete, select + +from shared.models.database.document import ( + DocumentMapUnit, + DocumentMapUnitIndex, + DocumentMapUnitToken, +) +from shared.services.retrieval.nav.nav_hierarchy import ProviderToolSpace +from shared.services.retrieval.nav.nav_map_scores import ( + build_score_units, + compute_corpus_map_and_unit_scores, +) +from shared.services.retrieval.nav.nav_knowhere import ( + KnowhereProvider, + LazyKnowhereProvider, + NamespaceKnowhereProvider, + ReadOnlyChunkStore, + SectionRow, + UnitRow, +) +from shared.services.retrieval.nav_snapshot import load_nav_snapshot +from shared.services.retrieval.publication_content import ( + replace_document_revision_content, +) +from shared.services.retrieval.publication_models import DocumentPublicationScope +from tests.support.contract_database import ContractDatabase +from tests.support.retrieval_snapshot_support import contract_db_session + +_USER_ID = "local-dev-user" + + +class _IncompleteIndexStore: + """Minimal lazy store whose incomplete index forces legacy scoring.""" + + def __init__(self, units_by_section: Mapping[str, Sequence[UnitRow]]) -> None: + self.units_by_section = { + str(section_id): list(units) + for section_id, units in units_by_section.items() + } + self.persisted_loads = 0 + self.batch_loads = 0 + + def load_persisted_score_corpus( + self, + document_ids: Sequence[str], + allowed_section_ids_by_document: Mapping[str, Sequence[str]], + queries: Sequence[str], + ) -> None: + del document_ids, allowed_section_ids_by_document, queries + self.persisted_loads += 1 + return None + + def load_documents_units( + self, + section_ids_by_document: Mapping[str, Sequence[str]], + ) -> dict[str, list[UnitRow]]: + self.batch_loads += 1 + return { + str(document_id): [ + unit + for section_id in section_ids + for unit in self.units_by_section.get(str(section_id), ()) + ] + for document_id, section_ids in section_ids_by_document.items() + } + + def load_document_units( + self, + document_id: str, + section_ids: Sequence[str], + extra_chunk_ids_by_section: Mapping[str, Sequence[str]] | None = None, + ) -> list[UnitRow]: + del document_id, extra_chunk_ids_by_section + return [ + unit + for section_id in section_ids + for unit in self.units_by_section.get(str(section_id), ()) + ] + + def load_section_units( + self, + document_id: str, + section_id: str, + extra_chunk_ids: Sequence[str] = (), + ) -> list[UnitRow]: + del document_id, extra_chunk_ids + return list(self.units_by_section.get(str(section_id), ())) + + def close(self) -> None: + return None + + +async def test_published_map_units_preserve_scores_without_chunk_payload_reads( + developer_api_client_factory: Callable[ + [], AbstractAsyncContextManager[AsyncClient] + ], + monkeypatch, +) -> None: + identifier = uuid4().hex[:8] + namespace = f"map-unit-index-{identifier}" + document_id = f"doc_map_{identifier}" + job_id = f"job_map_{identifier}" + job_result_id = f"result_map_{identifier}" + async with developer_api_client_factory(): + await _seed_revision( + namespace=namespace, + document_id=document_id, + job_id=job_id, + job_result_id=job_result_id, + ) + scope = DocumentPublicationScope( + user_id=_USER_ID, + namespace=namespace, + document_id=document_id, + job_result_id=job_result_id, + source_file_name="indexed.pdf", + ) + chunks = [ + { + "chunk_id": "parent-a", + "type": "text", + "content": "common alpha parent evidence", + "path": "indexed.pdf/Root/Parent/intro-a", + "order": 1, + "metadata": {}, + }, + { + "chunk_id": "parent-b", + "type": "text", + "content": "common beta parent evidence", + "path": "indexed.pdf/Root/Parent/intro-b", + "order": 2, + "metadata": {}, + }, + { + "chunk_id": "leaf-a", + "type": "text", + "content": "common alpha leaf evidence", + "path": "indexed.pdf/Root/Parent/Leaf A/body", + "order": 3, + "metadata": {}, + }, + { + "chunk_id": "leaf-b", + "type": "text", + "content": "common beta leaf evidence", + "path": "indexed.pdf/Root/Parent/Leaf B/body", + "order": 4, + "metadata": {}, + }, + ] + async with contract_db_session() as db: + await db.run_sync( + lambda sync_db: replace_document_revision_content( + sync_db, + scope=scope, + chunks=chunks, + ) + ) + await db.commit() + + async with contract_db_session() as db: + eager_snapshot = await load_nav_snapshot( + db, + user_id=_USER_ID, + namespace=namespace, + ) + eager_toolspace = ProviderToolSpace(eager_snapshot.provider) + expected_units = build_score_units(eager_toolspace, document_id) + expected_scores = compute_corpus_map_and_unit_scores( + eager_toolspace, + doc_ids=[document_id], + query="common alpha", + ) + + async with contract_db_session() as db: + index = ( + await db.execute( + select(DocumentMapUnitIndex).where( + DocumentMapUnitIndex.document_id == document_id + ) + ) + ).scalar_one() + persisted_units = list( + ( + await db.execute( + select(DocumentMapUnit) + .where(DocumentMapUnit.document_id == document_id) + .order_by(DocumentMapUnit.sort_order) + ) + ).scalars() + ) + persisted_tokens = list( + ( + await db.execute( + select(DocumentMapUnitToken).where( + DocumentMapUnitToken.map_unit_id.in_( + [unit.id for unit in persisted_units] + ) + ) + ) + ).scalars() + ) + lazy_snapshot = await load_nav_snapshot( + db, + user_id=_USER_ID, + namespace=namespace, + lazy=True, + ) + + assert index.unit_count == len(expected_units) + assert [unit.unit_id for unit in persisted_units] == [ + str(unit["chunk_id"]) for unit in expected_units + ] + assert persisted_tokens + + def reject_payload_read( + _store: ReadOnlyChunkStore, + _section_ids_by_document: Mapping[str, Sequence[str]], + ) -> dict[str, list[UnitRow]]: + raise AssertionError("persisted map scoring loaded full chunk payloads") + + original_payload_loader = ReadOnlyChunkStore.load_documents_units + monkeypatch.setattr( + ReadOnlyChunkStore, + "load_documents_units", + reject_payload_read, + ) + actual_scores = compute_corpus_map_and_unit_scores( + ProviderToolSpace(lazy_snapshot.provider), + doc_ids=[document_id], + query="common alpha", + ) + monkeypatch.setattr( + ReadOnlyChunkStore, + "load_documents_units", + original_payload_loader, + ) + async with contract_db_session() as db: + token_id = ( + select(DocumentMapUnitToken.id) + .join( + DocumentMapUnit, + DocumentMapUnit.id == DocumentMapUnitToken.map_unit_id, + ) + .where(DocumentMapUnit.document_id == document_id) + .limit(1) + .scalar_subquery() + ) + await db.execute( + delete(DocumentMapUnitToken).where(DocumentMapUnitToken.id == token_id) + ) + await db.commit() + incomplete_snapshot = await load_nav_snapshot( + db, + user_id=_USER_ID, + namespace=namespace, + lazy=True, + ) + fallback_scores = compute_corpus_map_and_unit_scores( + ProviderToolSpace(incomplete_snapshot.provider), + doc_ids=[document_id], + query="common alpha", + ) + incomplete_snapshot.close() + lazy_snapshot.close() + eager_snapshot.close() + + assert actual_scores == expected_scores + assert fallback_scores == expected_scores + + +def test_incomplete_index_falls_back_for_duplicate_unit_ids() -> None: + first_sections = [ + SectionRow("root-a", None, "Root A", "Root A", 0, "", 0), + SectionRow("leaf-a", "root-a", "Root A / Leaf A", "Leaf A", 1, "", 1), + ] + second_sections = [ + SectionRow("root-b", None, "Root B", "Root B", 0, "", 0), + SectionRow("leaf-b", "root-b", "Root B / Leaf B", "Leaf B", 1, "", 1), + ] + first_unit = UnitRow("same-chunk", "leaf-a", "text", "alpha evidence", 1) + second_unit = UnitRow("same-chunk", "leaf-b", "text", "beta evidence", 1) + + eager = ProviderToolSpace( + NamespaceKnowhereProvider( + [ + KnowhereProvider( + doc_id="doc-a", sections=first_sections, units=[first_unit] + ), + KnowhereProvider( + doc_id="doc-b", sections=second_sections, units=[second_unit] + ), + ], + titles={"doc-a": "Document A", "doc-b": "Document B"}, + ) + ) + store = _IncompleteIndexStore( + {"leaf-a": [first_unit], "leaf-b": [second_unit]} + ) + lazy = ProviderToolSpace( + NamespaceKnowhereProvider( + [ + LazyKnowhereProvider( + doc_id="doc-a", + sections=first_sections, + chunk_store=store, + known_chunk_ids=[first_unit.chunk_id], + ), + LazyKnowhereProvider( + doc_id="doc-b", + sections=second_sections, + chunk_store=store, + known_chunk_ids=[second_unit.chunk_id], + ), + ], + titles={"doc-a": "Document A", "doc-b": "Document B"}, + chunk_owner_by_id={"same-chunk": "doc-a"}, + ) + ) + + expected = compute_corpus_map_and_unit_scores( + eager, doc_ids=["doc-a", "doc-b"], query="alpha beta" + ) + actual = compute_corpus_map_and_unit_scores( + lazy, doc_ids=["doc-a", "doc-b"], query="alpha beta" + ) + + assert actual == expected + assert store.persisted_loads == 1 + assert store.batch_loads == 1 + + +def test_titleless_leaf_has_identical_eager_and_lazy_path_scoring() -> None: + sections = [ + SectionRow("root", None, "Root", "Root", 0, "", 0), + SectionRow("leaf", "root", "Root / Leaf", "", 1, "", 1), + ] + unit = UnitRow("titleless-chunk", "leaf", "text", "alpha evidence", 1) + eager = ProviderToolSpace( + KnowhereProvider(doc_id="doc", sections=sections, units=[unit]) + ) + store = _IncompleteIndexStore({"leaf": [unit]}) + lazy = ProviderToolSpace( + LazyKnowhereProvider( + doc_id="doc", + sections=sections, + chunk_store=store, + known_chunk_ids=[unit.chunk_id], + ) + ) + + assert build_score_units(eager, "doc") == build_score_units(lazy, "doc") + assert compute_corpus_map_and_unit_scores( + eager, doc_ids=["doc"], query="alpha" + ) == compute_corpus_map_and_unit_scores( + lazy, doc_ids=["doc"], query="alpha" + ) + + +async def _seed_revision( + *, + namespace: str, + document_id: str, + job_id: str, + job_result_id: str, +) -> None: + await ContractDatabase.execute( + """ + INSERT INTO jobs ( + job_id, user_id, job_type, status, source_type, version, + webhook_enabled, created_at, updated_at, credits_charged, billing_status + ) VALUES ( + :job_id, :user_id, 'document_ingestion', 'done', 'file', 0, + false, NOW(), NOW(), 0, 'skipped' + ) + """, + {"job_id": job_id, "user_id": _USER_ID}, + ) + await ContractDatabase.execute( + """ + INSERT INTO documents ( + document_id, user_id, namespace, status, source_file_name, + parse_track, created_at, updated_at + ) VALUES ( + :document_id, :user_id, :namespace, 'active', + 'indexed.pdf', 'chunk', NOW(), NOW() + ) + """, + { + "document_id": document_id, + "user_id": _USER_ID, + "namespace": namespace, + }, + ) + await ContractDatabase.execute( + """ + INSERT INTO job_results ( + id, job_id, document_id, delivery_mode, created_at, updated_at + ) VALUES ( + :job_result_id, :job_id, :document_id, 'inline', NOW(), NOW() + ) + """, + { + "job_result_id": job_result_id, + "job_id": job_id, + "document_id": document_id, + }, + ) + await ContractDatabase.execute( + """ + UPDATE documents SET current_job_result_id = :job_result_id + WHERE document_id = :document_id + """, + {"job_result_id": job_result_id, "document_id": document_id}, + ) diff --git a/apps/api/tests/migrations/test_schema_contract.py b/apps/api/tests/migrations/test_schema_contract.py index 24047dccd..3ead12179 100644 --- a/apps/api/tests/migrations/test_schema_contract.py +++ b/apps/api/tests/migrations/test_schema_contract.py @@ -233,6 +233,28 @@ def test_should_index_document_chunks_in_snapshot_pagination_order( ) +def test_should_index_document_chunks_in_lazy_section_order( + migrated_head_engine: Engine, +) -> None: + with migrated_head_engine.begin() as connection: + index_definition = connection.execute( + text( + """ + SELECT indexdef + FROM pg_indexes + WHERE schemaname = current_schema() + AND tablename = 'document_chunks' + AND indexname = 'idx_document_chunks_revision_section_order' + """ + ) + ).scalar_one() + + assert ( + "(document_id, job_result_id, section_id, sort_order, chunk_id, id)" + in str(index_definition) + ) + + def test_should_upgrade_with_a_caller_owned_connection( alembic_engine: Engine, ) -> None: diff --git a/deploy/ecs/README.md b/deploy/ecs/README.md index e19200afe..dd0742b4d 100644 --- a/deploy/ecs/README.md +++ b/deploy/ecs/README.md @@ -69,6 +69,34 @@ workflow validates these resources, registers immutable image-digest task definitions, runs the production migration first, and then updates the ECS services. It does not create or delete AWS resources. +## Required post-deploy backfill + +The map-nav lexical-index migration creates the derived index tables, but it does +not rebuild indexes for revisions that already exist. Until those revisions are +backfilled, retrieval remains quality-preserving but uses the legacy scoring +path. Every release containing the map-nav index change must include the +following DevOps action in its release notification. + +Run the commands as a one-off container using the newly deployed API image and +the production database secret. Do not run them inside the long-lived API task. + +```bash +# Read-only inventory +python /app/scripts/backfill_map_unit_indexes.py + +# Optional canary: apply one affected document first +python /app/scripts/backfill_map_unit_indexes.py \ + --document-id \ + --apply + +# Apply to all current document revisions +python /app/scripts/backfill_map_unit_indexes.py --apply +``` + +The script commits each document revision independently and is safe to rerun. +Verify the canary retrieval before starting the full apply. New or republished +documents build their index automatically during publication. + ## Manual staging availability `.github/workflows/manage-staging.yml` exposes three manually dispatched diff --git a/packages/shared-python/shared/models/database/__init__.py b/packages/shared-python/shared/models/database/__init__.py index 2d0e7d8a8..5535cd9bf 100644 --- a/packages/shared-python/shared/models/database/__init__.py +++ b/packages/shared-python/shared/models/database/__init__.py @@ -13,6 +13,9 @@ from .document import ( Document, DocumentChunk, + DocumentMapUnit, + DocumentMapUnitIndex, + DocumentMapUnitToken, DocumentSection, GraphEdge, GraphNode, @@ -56,6 +59,9 @@ "Document", "DocumentSection", "DocumentChunk", + "DocumentMapUnit", + "DocumentMapUnitIndex", + "DocumentMapUnitToken", "DocumentPagePlan", "DemoMaterialization", "GraphNode", diff --git a/packages/shared-python/shared/models/database/document.py b/packages/shared-python/shared/models/database/document.py index f9681e178..33dda2f91 100644 --- a/packages/shared-python/shared/models/database/document.py +++ b/packages/shared-python/shared/models/database/document.py @@ -49,7 +49,9 @@ class Document(Base): document_metadata: Mapped[Optional[Dict[str, Any]]] = mapped_column( JSON, nullable=True ) - parse_track: Mapped[str] = mapped_column(String(32), nullable=False, default="chunk") + parse_track: Mapped[str] = mapped_column( + String(32), nullable=False, default="chunk" + ) created_at: Mapped[datetime] = mapped_column( DateTime, default=utc_now_naive, nullable=False ) @@ -139,6 +141,7 @@ class DocumentChunk(Base): id: Mapped[str] = mapped_column( String(36), primary_key=True, default=lambda: f"dchk_{uuid4().hex[:12]}" ) + chunk_id: Mapped[str] = mapped_column(String(64), nullable=False) user_id: Mapped[str] = mapped_column(Text, nullable=False) namespace: Mapped[str] = mapped_column( @@ -210,6 +213,15 @@ class DocumentChunk(Base): "chunk_id", "id", ), + Index( + "idx_document_chunks_revision_section_order", + "document_id", + "job_result_id", + "section_id", + "sort_order", + "chunk_id", + "id", + ), Index("idx_document_chunks_section", "section_id"), Index( "idx_chunk_content_search_tsv", @@ -224,6 +236,110 @@ class DocumentChunk(Base): ) +class DocumentMapUnit(Base): + """Persisted lexical map unit for one document revision. + + These rows are a derived index of the exact leaf and interstitial units + used by map-nav. Full chunk payloads remain in ``document_chunks`` and are + loaded separately for evidence hydration. + """ + + __tablename__ = "document_map_units" + + id: Mapped[str] = mapped_column(String(160), primary_key=True) + document_id: Mapped[str] = mapped_column( + String(36), + ForeignKey("documents.document_id", ondelete="CASCADE"), + nullable=False, + ) + job_result_id: Mapped[str] = mapped_column( + String(36), ForeignKey("job_results.id", ondelete="CASCADE"), nullable=False + ) + unit_id: Mapped[str] = mapped_column(String(128), nullable=False) + section_id: Mapped[str] = mapped_column(String(36), nullable=False) + unit_kind: Mapped[str] = mapped_column(String(32), nullable=False) + path_token_count: Mapped[int] = mapped_column(Integer, nullable=False) + content_token_count: Mapped[int] = mapped_column(Integer, nullable=False) + term_search_text_lower: Mapped[str] = mapped_column(Text, nullable=False) + sort_order: Mapped[int] = mapped_column(Integer, nullable=False, default=0) + created_at: Mapped[datetime] = mapped_column( + DateTime, default=utc_now_naive, nullable=False + ) + + __table_args__ = ( + Index( + "idx_document_map_units_revision_order", + "document_id", + "job_result_id", + "sort_order", + "unit_id", + ), + Index("idx_document_map_units_section", "section_id"), + ) + + +class DocumentMapUnitToken(Base): + """One exact token frequency in a persisted map unit channel.""" + + __tablename__ = "document_map_unit_tokens" + + id: Mapped[str] = mapped_column(String(36), primary_key=True) + map_unit_id: Mapped[str] = mapped_column( + String(160), + ForeignKey("document_map_units.id", ondelete="CASCADE"), + nullable=False, + ) + channel: Mapped[str] = mapped_column(String(16), nullable=False) + token: Mapped[str] = mapped_column(Text, nullable=False) + token_hash: Mapped[str] = mapped_column(String(64), nullable=False) + frequency: Mapped[int] = mapped_column(Integer, nullable=False) + + __table_args__ = ( + Index( + "idx_document_map_unit_tokens_lookup", + "channel", + "token_hash", + "map_unit_id", + ), + Index("idx_document_map_unit_tokens_unit", "map_unit_id", "channel"), + ) + + +class DocumentMapUnitIndex(Base): + """Completeness marker for a revision's materialized map-unit index.""" + + __tablename__ = "document_map_unit_indexes" + + id: Mapped[str] = mapped_column(String(100), primary_key=True) + document_id: Mapped[str] = mapped_column( + String(36), + ForeignKey("documents.document_id", ondelete="CASCADE"), + nullable=False, + ) + job_result_id: Mapped[str] = mapped_column( + String(36), ForeignKey("job_results.id", ondelete="CASCADE"), nullable=False + ) + format_version: Mapped[int] = mapped_column(Integer, nullable=False) + unit_count: Mapped[int] = mapped_column(Integer, nullable=False) + token_count: Mapped[int] = mapped_column(Integer, nullable=False) + created_at: Mapped[datetime] = mapped_column( + DateTime, default=utc_now_naive, nullable=False + ) + + __table_args__ = ( + UniqueConstraint( + "document_id", + "job_result_id", + name="uq_document_map_unit_indexes_revision", + ), + Index( + "idx_document_map_unit_indexes_revision", + "document_id", + "job_result_id", + ), + ) + + class GraphNode(Base): """Persisted derived graph node used for routing and expansion.""" @@ -376,44 +492,58 @@ class RetrievalHitStat(Base): class RetrievalRun(Base): """One row per agentic retrieval query. Append-only analytics.""" - __tablename__ = 'retrieval_runs' + __tablename__ = "retrieval_runs" - run_id: Mapped[str] = mapped_column(String(36), primary_key=True, default=lambda: f'aret_{uuid4().hex[:12]}') + run_id: Mapped[str] = mapped_column( + String(36), primary_key=True, default=lambda: f"aret_{uuid4().hex[:12]}" + ) user_id: Mapped[str] = mapped_column(Text, nullable=False) - namespace: Mapped[str] = mapped_column(String(255), nullable=False, default='default') + namespace: Mapped[str] = mapped_column( + String(255), nullable=False, default="default" + ) query: Mapped[str] = mapped_column(Text, nullable=False) - query_hash: Mapped[str] = mapped_column(String(32), nullable=False, default='') + query_hash: Mapped[str] = mapped_column(String(32), nullable=False, default="") top_k: Mapped[int] = mapped_column(Integer, nullable=False, default=10) chunk_types: Mapped[Optional[List[str]]] = mapped_column(JSON, nullable=True) filters: Mapped[Optional[Dict[str, Any]]] = mapped_column(JSON, nullable=True) - policy_name: Mapped[str] = mapped_column(String(64), nullable=False, default='rule_based_v1') + policy_name: Mapped[str] = mapped_column( + String(64), nullable=False, default="rule_based_v1" + ) agentic_enabled: Mapped[bool] = mapped_column(nullable=False, default=True) cache_hit: Mapped[bool] = mapped_column(nullable=False, default=False) result_count: Mapped[int] = mapped_column(Integer, nullable=False, default=0) final_doc_ids: Mapped[Optional[List[str]]] = mapped_column(JSON, nullable=True) - result_provenance: Mapped[Optional[Dict[str, Any]]] = mapped_column(JSON, nullable=True) - parent_run_id: Mapped[Optional[str]] = mapped_column(String(36), nullable=True, index=True) + result_provenance: Mapped[Optional[Dict[str, Any]]] = mapped_column( + JSON, nullable=True + ) + parent_run_id: Mapped[Optional[str]] = mapped_column( + String(36), nullable=True, index=True + ) workflow_step_id: Mapped[Optional[str]] = mapped_column(String(64), nullable=True) workflow_plan: Mapped[Optional[Dict[str, Any]]] = mapped_column(JSON, nullable=True) latency_ms: Mapped[int] = mapped_column(Integer, nullable=False, default=0) token_count: Mapped[Optional[int]] = mapped_column(Integer, nullable=True) error: Mapped[Optional[str]] = mapped_column(Text, nullable=True) - created_at: Mapped[datetime] = mapped_column(DateTime, default=datetime.utcnow, nullable=False) + created_at: Mapped[datetime] = mapped_column( + DateTime, default=datetime.utcnow, nullable=False + ) completed_at: Mapped[Optional[datetime]] = mapped_column(DateTime, nullable=True) __table_args__ = ( - Index('idx_retrieval_runs_user_namespace', 'user_id', 'namespace'), - Index('idx_retrieval_runs_created', 'created_at'), - Index('idx_retrieval_runs_query_hash', 'query_hash'), + Index("idx_retrieval_runs_user_namespace", "user_id", "namespace"), + Index("idx_retrieval_runs_created", "created_at"), + Index("idx_retrieval_runs_query_hash", "query_hash"), ) class RetrievalStep(Base): """One row per agent step within a retrieval run. Append-only analytics.""" - __tablename__ = 'retrieval_steps' + __tablename__ = "retrieval_steps" - step_id: Mapped[str] = mapped_column(String(36), primary_key=True, default=lambda: f'arst_{uuid4().hex[:12]}') + step_id: Mapped[str] = mapped_column( + String(36), primary_key=True, default=lambda: f"arst_{uuid4().hex[:12]}" + ) run_id: Mapped[str] = mapped_column(String(36), nullable=False, index=True) step_index: Mapped[int] = mapped_column(Integer, nullable=False, default=0) action_type: Mapped[str] = mapped_column(String(64), nullable=False) @@ -425,9 +555,11 @@ class RetrievalStep(Base): token_count: Mapped[Optional[int]] = mapped_column(Integer, nullable=True) model_name: Mapped[Optional[str]] = mapped_column(String(128), nullable=True) error: Mapped[Optional[str]] = mapped_column(Text, nullable=True) - created_at: Mapped[datetime] = mapped_column(DateTime, default=datetime.utcnow, nullable=False) + created_at: Mapped[datetime] = mapped_column( + DateTime, default=datetime.utcnow, nullable=False + ) __table_args__ = ( - Index('idx_retrieval_steps_run', 'run_id', 'step_index'), - Index('idx_retrieval_steps_created', 'created_at'), + Index("idx_retrieval_steps_run", "run_id", "step_index"), + Index("idx_retrieval_steps_created", "created_at"), ) diff --git a/packages/shared-python/shared/services/retrieval/execution/routes.py b/packages/shared-python/shared/services/retrieval/execution/routes.py index e67a1ead8..979610742 100644 --- a/packages/shared-python/shared/services/retrieval/execution/routes.py +++ b/packages/shared-python/shared/services/retrieval/execution/routes.py @@ -1,6 +1,7 @@ from __future__ import annotations import asyncio +import time from contextlib import AbstractAsyncContextManager from loguru import logger @@ -180,12 +181,22 @@ async def _run_mapnav_route( episode_token_count, episode_workflow_plan, ) + snapshot_started = time.perf_counter() snapshot = await load_nav_snapshot( context.db, user_id=context.user_id, namespace=context.namespace, exclude_document_ids=context.exclude_document_ids, exclude_sections=context.exclude_sections, + lazy=True, + ) + snapshot_seconds = time.perf_counter() - snapshot_started + logger.info( + "retrieval mapnav stage=snapshot_load seconds={:.3f} documents={} refs={}".format( + snapshot_seconds, + len(snapshot.document_ids), + len(snapshot.chunk_ref_index), + ) ) # Small-corpus count / snapshot reads may leave a checkout; drop it before @@ -196,20 +207,31 @@ async def _run_mapnav_route( cfg = build_nav_config() toolspace = ProviderToolSpace(snapshot.provider) - episode = await asyncio.to_thread( - run_nav_episode, - None, - context.query, - corpus_doc_ids=list(snapshot.document_ids), - budget_chars=budget, - compose_answer=False, - policy="llm", - config=cfg, - toolspace=toolspace, - ) + episode_started = time.perf_counter() + try: + episode = await asyncio.to_thread( + run_nav_episode, + None, + context.query, + corpus_doc_ids=list(snapshot.document_ids), + budget_chars=budget, + compose_answer=False, + policy="llm", + config=cfg, + toolspace=toolspace, + ) - refs, score_by_chunk_id = build_referenced_chunks(episode, snapshot) + refs, score_by_chunk_id = build_referenced_chunks(episode, snapshot) + logger.info( + "retrieval mapnav stage=episode seconds={:.3f} refs={}".format( + time.perf_counter() - episode_started, + len(refs), + ) + ) + finally: + snapshot.close() + hydration_started = time.perf_counter() async with open_fresh_database_context() as final_db: resolved = await resolve_workflow_references( db=final_db, @@ -256,6 +278,12 @@ async def _run_mapnav_route( selected_paths=selected_paths, selected_doc_ids=selected_docs, ) + logger.info( + "retrieval mapnav stage=hydration seconds={:.3f} results={}".format( + time.perf_counter() - hydration_started, + len(assembled_rows), + ) + ) stop_reason = str(getattr(episode, "stop_reason", "") or "completed") evidence_text = str(getattr(episode, "evidence_text", "") or "") diff --git a/packages/shared-python/shared/services/retrieval/map_unit_index.py b/packages/shared-python/shared/services/retrieval/map_unit_index.py new file mode 100644 index 000000000..15454dc9d --- /dev/null +++ b/packages/shared-python/shared/services/retrieval/map_unit_index.py @@ -0,0 +1,163 @@ +"""Publication-time materialization of exact map-nav lexical units.""" + +from __future__ import annotations + +from collections import Counter +from hashlib import sha256 +from uuid import uuid4 + +from sqlalchemy import delete, select +from sqlalchemy.orm import Session + +from shared.models.database.document import ( + DocumentChunk, + DocumentMapUnit, + DocumentMapUnitIndex, + DocumentMapUnitToken, + DocumentSection, +) +from shared.services.retrieval.nav.nav_hierarchy import ProviderToolSpace +from shared.services.retrieval.nav.nav_knowhere import ( + KnowhereProvider, + SectionRow, + UnitRow, +) +from shared.services.retrieval.nav.nav_map_scores import build_score_units +from shared.services.retrieval.publication_models import DocumentPublicationScope + + +MAP_UNIT_INDEX_FORMAT_VERSION = 1 + + +def replace_document_map_units( + db: Session, + *, + scope: DocumentPublicationScope, +) -> None: + """Build the derived index through the authoritative map-unit constructor.""" + db.execute( + delete(DocumentMapUnitToken).where( + DocumentMapUnitToken.map_unit_id.in_( + select(DocumentMapUnit.id) + .where(DocumentMapUnit.document_id == scope.document_id) + .where(DocumentMapUnit.job_result_id == scope.job_result_id) + ) + ) + ) + db.execute( + delete(DocumentMapUnit) + .where(DocumentMapUnit.document_id == scope.document_id) + .where(DocumentMapUnit.job_result_id == scope.job_result_id) + ) + db.execute( + delete(DocumentMapUnitIndex) + .where(DocumentMapUnitIndex.document_id == scope.document_id) + .where(DocumentMapUnitIndex.job_result_id == scope.job_result_id) + ) + section_models = list( + db.scalars( + select(DocumentSection) + .where(DocumentSection.document_id == scope.document_id) + .where(DocumentSection.job_result_id == scope.job_result_id) + .order_by(DocumentSection.sort_order, DocumentSection.section_id) + ) + ) + chunk_models = list( + db.scalars( + select(DocumentChunk) + .where(DocumentChunk.document_id == scope.document_id) + .where(DocumentChunk.job_result_id == scope.job_result_id) + .order_by( + DocumentChunk.sort_order, + DocumentChunk.chunk_id, + DocumentChunk.id, + ) + ) + ) + provider = KnowhereProvider( + doc_id=scope.document_id, + sections=[_to_section_row(section) for section in section_models], + units=[_to_unit_row(chunk) for chunk in chunk_models], + ) + score_units = build_score_units( + ProviderToolSpace(provider), + scope.document_id, + ) + persisted_count = 0 + token_count = 0 + for sort_order, unit in enumerate(score_units): + unit_id = str(unit.get("chunk_id") or "").strip() + section_id = str(unit.get("section_id") or "").strip() + if not unit_id or not section_id: + continue + map_unit_id = f"dmu_{uuid4().hex}" + path_tokens = str(unit.get("path_search_text") or "").split() + content_tokens = str(unit.get("content_search_text") or "").split() + db.add( + DocumentMapUnit( + id=map_unit_id, + document_id=scope.document_id, + job_result_id=scope.job_result_id, + unit_id=unit_id, + section_id=section_id, + unit_kind=str(unit.get("kind") or "leaf"), + path_token_count=len(path_tokens), + content_token_count=len(content_tokens), + term_search_text_lower=str(unit.get("term_search_text") or "").lower(), + sort_order=sort_order, + ) + ) + for channel, frequencies in ( + ("path", Counter(path_tokens)), + ("content", Counter(content_tokens)), + ): + for token, frequency in frequencies.items(): + db.add( + DocumentMapUnitToken( + id=f"dmut_{uuid4().hex[:31]}", + map_unit_id=map_unit_id, + channel=channel, + token=token, + token_hash=sha256(token.encode("utf-8")).hexdigest(), + frequency=frequency, + ) + ) + token_count += len(frequencies) + persisted_count += 1 + db.add( + DocumentMapUnitIndex( + id=f"dmui_{uuid4().hex}", + document_id=scope.document_id, + job_result_id=scope.job_result_id, + format_version=MAP_UNIT_INDEX_FORMAT_VERSION, + unit_count=persisted_count, + token_count=token_count, + ) + ) + + +def _to_section_row(section: DocumentSection) -> SectionRow: + return SectionRow( + section_id=section.section_id, + parent_section_id=section.parent_section_id, + section_path=section.section_path, + section_title=str(section.section_title or ""), + section_level=section.section_level, + summary=str(section.summary or ""), + sort_order=section.sort_order, + ) + + +def _to_unit_row(chunk: DocumentChunk) -> UnitRow: + raw_metadata = chunk.chunk_metadata + metadata = dict(raw_metadata) if isinstance(raw_metadata, dict) else {} + return UnitRow( + chunk_id=chunk.chunk_id, + section_id=chunk.section_id, + chunk_type=chunk.chunk_type, + content=str(chunk.content or ""), + sort_order=chunk.sort_order, + source_chunk_path=str(chunk.source_chunk_path or ""), + file_path=str(chunk.file_path or ""), + metadata=metadata, + ) diff --git a/packages/shared-python/shared/services/retrieval/nav/knowhere_hybrid.py b/packages/shared-python/shared/services/retrieval/nav/knowhere_hybrid.py index d56c72fb9..369da962c 100644 --- a/packages/shared-python/shared/services/retrieval/nav/knowhere_hybrid.py +++ b/packages/shared-python/shared/services/retrieval/nav/knowhere_hybrid.py @@ -5,11 +5,16 @@ Reference: https://github.com/Ontos-AI/knowhere """ + from __future__ import annotations import os import re -from typing import Any, Dict, List, Optional, Sequence, Tuple +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 RRF_K = 60 CHANNEL_WEIGHT_PATH = 1.0 @@ -18,6 +23,19 @@ INTERNAL_RECALL_K_MULTIPLIER = 2 +class ScoreUnitRow(TypedDict, total=False): + """Compact scoring-unit shape shared by eager and streaming scorers.""" + + chunk_id: str + section_id: str + kind: str + content: str + path_text: str + path_search_text: str + content_search_text: str + term_search_text: str + + def tokenize_for_retrieval(text: str, *, dedupe: bool = True) -> List[str]: tokens = re.findall(r"[a-z0-9_]+|[\u4e00-\u9fff]", str(text or "").lower()) if not dedupe: @@ -39,7 +57,9 @@ def _space_join_tokens(text: str) -> str: return " ".join(tokenize_for_retrieval(text, dedupe=False)) -def build_content_search_text(content: str, *, section_summary: Optional[str] = None) -> str: +def build_content_search_text( + content: str, *, section_summary: Optional[str] = None +) -> str: parts = [str(content or "").strip()] if section_summary and str(section_summary).strip(): parts.append(str(section_summary).strip()) @@ -68,7 +88,7 @@ def build_term_search_text(content: str, *, path_text: Optional[str] = None) -> return combined -def _get_search_tokens(row: dict[str, Any], *, search_field: str) -> List[str]: +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] @@ -99,7 +119,9 @@ def rank_rows_by_bm25( try: from rank_bm25 import BM25Okapi except ImportError: - return _rank_rows_by_token_overlap(rows, query_tokens, search_field=search_field) + return _rank_rows_by_token_overlap( + rows, query_tokens, search_field=search_field + ) corpus: List[List[str]] = [] ranked_rows: List[dict[str, Any]] = [] @@ -123,7 +145,9 @@ def rank_rows_by_bm25( return ranked_rows -def rank_rows_by_term_channel(rows: List[dict[str, Any]], query: str) -> List[dict[str, Any]]: +def rank_rows_by_term_channel( + rows: List[dict[str, Any]], query: str +) -> List[dict[str, Any]]: query_lower = query.lower().strip() query_tokens = tokenize_query_for_ranker(query) if not query_lower or not query_tokens: @@ -200,12 +224,24 @@ def normalize_row_scores( def _channel_weights() -> Tuple[float, float, float]: - path_w = float(os.environ.get("NAV_DISCOVERY_CHANNEL_WEIGHT_PATH", str(CHANNEL_WEIGHT_PATH)).strip() or CHANNEL_WEIGHT_PATH) + path_w = float( + os.environ.get( + "NAV_DISCOVERY_CHANNEL_WEIGHT_PATH", str(CHANNEL_WEIGHT_PATH) + ).strip() + or CHANNEL_WEIGHT_PATH + ) content_w = float( - os.environ.get("NAV_DISCOVERY_CHANNEL_WEIGHT_CONTENT", str(CHANNEL_WEIGHT_CONTENT)).strip() + os.environ.get( + "NAV_DISCOVERY_CHANNEL_WEIGHT_CONTENT", str(CHANNEL_WEIGHT_CONTENT) + ).strip() or CHANNEL_WEIGHT_CONTENT ) - term_w = float(os.environ.get("NAV_DISCOVERY_CHANNEL_WEIGHT_TERM", str(CHANNEL_WEIGHT_TERM)).strip() or CHANNEL_WEIGHT_TERM) + term_w = float( + os.environ.get( + "NAV_DISCOVERY_CHANNEL_WEIGHT_TERM", str(CHANNEL_WEIGHT_TERM) + ).strip() + or CHANNEL_WEIGHT_TERM + ) return path_w, content_w, term_w @@ -225,14 +261,23 @@ def hybrid_search_rows( recall_k = internal_recall_k if recall_k is None: - mult = int(os.environ.get("NAV_DISCOVERY_RECALL_MULT", str(INTERNAL_RECALL_K_MULTIPLIER)).strip() or INTERNAL_RECALL_K_MULTIPLIER) + mult = int( + os.environ.get( + "NAV_DISCOVERY_RECALL_MULT", str(INTERNAL_RECALL_K_MULTIPLIER) + ).strip() + or INTERNAL_RECALL_K_MULTIPLIER + ) recall_k = max(top_k, top_k * max(1, mult)) rrf_k = int(os.environ.get("NAV_DISCOVERY_RRF_K", str(RRF_K)).strip() or RRF_K) path_w, content_w, term_w = _channel_weights() - path_rows = rank_rows_by_bm25(list(rows), query_tokens, search_field="path_search_text")[:recall_k] - content_rows = rank_rows_by_bm25(list(rows), query_tokens, search_field="content_search_text")[:recall_k] + path_rows = rank_rows_by_bm25( + list(rows), query_tokens, search_field="path_search_text" + )[:recall_k] + content_rows = rank_rows_by_bm25( + list(rows), query_tokens, search_field="content_search_text" + )[:recall_k] term_rows = rank_rows_by_term_channel(list(rows), query)[:recall_k] fused = merge_channels_rrf( @@ -245,27 +290,32 @@ def hybrid_search_rows( return fused - def map_channel_weights() -> Tuple[float, float, float]: """Channel weights for map scoring (prefer NAV_MAP_* env, fall back to legacy names).""" path_w = float( os.environ.get( "NAV_MAP_CHANNEL_WEIGHT_PATH", - os.environ.get("NAV_DISCOVERY_CHANNEL_WEIGHT_PATH", str(CHANNEL_WEIGHT_PATH)), + os.environ.get( + "NAV_DISCOVERY_CHANNEL_WEIGHT_PATH", str(CHANNEL_WEIGHT_PATH) + ), ).strip() or CHANNEL_WEIGHT_PATH ) content_w = float( os.environ.get( "NAV_MAP_CHANNEL_WEIGHT_CONTENT", - os.environ.get("NAV_DISCOVERY_CHANNEL_WEIGHT_CONTENT", str(CHANNEL_WEIGHT_CONTENT)), + os.environ.get( + "NAV_DISCOVERY_CHANNEL_WEIGHT_CONTENT", str(CHANNEL_WEIGHT_CONTENT) + ), ).strip() or CHANNEL_WEIGHT_CONTENT ) term_w = float( os.environ.get( "NAV_MAP_CHANNEL_WEIGHT_TERM", - os.environ.get("NAV_DISCOVERY_CHANNEL_WEIGHT_TERM", str(CHANNEL_WEIGHT_TERM)), + os.environ.get( + "NAV_DISCOVERY_CHANNEL_WEIGHT_TERM", str(CHANNEL_WEIGHT_TERM) + ), ).strip() or CHANNEL_WEIGHT_TERM ) @@ -428,18 +478,23 @@ def fuse_channel_bm25_dense( dense_vals = [float(dense_by_id.get(uid, 0.0) or 0.0) for uid in unit_ids] bm25_n = _normalize_score_list(bm25_vals) dense_n = _normalize_score_list(dense_vals) - dense_w = float(os.environ.get("NAV_MAP_CHANNEL_DENSE_WEIGHT", "0.5").strip() or "0.5") + dense_w = float( + os.environ.get("NAV_MAP_CHANNEL_DENSE_WEIGHT", "0.5").strip() or "0.5" + ) dense_w = min(1.0, max(0.0, dense_w)) bm25_w = 1.0 - dense_w return { - uid: bm25_w * bm25_n[i] + dense_w * dense_n[i] - for i, uid in enumerate(unit_ids) + uid: bm25_w * bm25_n[i] + dense_w * dense_n[i] for i, uid in enumerate(unit_ids) } def _rank_ids_by_score(score_by_id: Dict[str, float]) -> List[str]: ranked = sorted( - ((sid, float(score)) for sid, score in score_by_id.items() if float(score) > 0.0), + ( + (sid, float(score)) + for sid, score in score_by_id.items() + if float(score) > 0.0 + ), key=lambda item: (-item[1], item[0]), ) return [sid for sid, _ in ranked] @@ -453,9 +508,7 @@ def score_rows_hybrid_all( content_texts: Optional[Dict[str, str]] = None, doc_id: Optional[str] = None, namespace: Optional[str] = None, - dense_scores_by_channel: Optional[ - Dict[str, Optional[Dict[str, float]]] - ] = None, + dense_scores_by_channel: Optional[Dict[str, Optional[Dict[str, float]]]] = None, ) -> List[dict[str, Any]]: """Score every row with path/content/term; optional within-channel dense fuse. @@ -473,7 +526,9 @@ def score_rows_hybrid_all( if not unit_ids: return [dict(row, score=0.0) for row in rows] - row_by_id = {str(row.get("chunk_id") or ""): dict(row) for row in rows if row.get("chunk_id")} + row_by_id = { + str(row.get("chunk_id") or ""): dict(row) for row in rows if row.get("chunk_id") + } path_w, content_w, term_w = map_channel_weights() rrf_k = int( os.environ.get( @@ -484,7 +539,9 @@ def score_rows_hybrid_all( ) if query_tokens: - path_ranked = rank_rows_by_bm25(list(rows), query_tokens, search_field="path_search_text") + path_ranked = rank_rows_by_bm25( + list(rows), query_tokens, search_field="path_search_text" + ) content_ranked = rank_rows_by_bm25( list(rows), query_tokens, search_field="content_search_text" ) @@ -496,7 +553,8 @@ def score_rows_hybrid_all( str(r.get("chunk_id") or ""): float(r.get("score") or 0.0) for r in path_ranked } content_bm25 = { - str(r.get("chunk_id") or ""): float(r.get("score") or 0.0) for r in content_ranked + str(r.get("chunk_id") or ""): float(r.get("score") or 0.0) + for r in content_ranked } term_bm25 = { str(r.get("chunk_id") or ""): float(r.get("score") or 0.0) for r in term_ranked @@ -536,19 +594,12 @@ def score_rows_hybrid_all( namespace=namespace, ) path_dense_by_id = ( - { - uid: float(path_dense_scores[i]) - for i, uid in enumerate(unit_ids) - } - if path_dense_scores is not None - and len(path_dense_scores) == len(unit_ids) + {uid: float(path_dense_scores[i]) for i, uid in enumerate(unit_ids)} + if path_dense_scores is not None and len(path_dense_scores) == len(unit_ids) else None ) content_dense_by_id = ( - { - uid: float(content_dense_scores[i]) - for i, uid in enumerate(unit_ids) - } + {uid: float(content_dense_scores[i]) for i, uid in enumerate(unit_ids)} if content_dense_scores is not None and len(content_dense_scores) == len(unit_ids) else None @@ -558,7 +609,9 @@ def score_rows_hybrid_all( content_dense_by_id = dense_scores_by_channel.get("content") path_channel = fuse_channel_bm25_dense(path_bm25, path_dense_by_id, unit_ids) - content_channel = fuse_channel_bm25_dense(content_bm25, content_dense_by_id, unit_ids) + content_channel = fuse_channel_bm25_dense( + content_bm25, content_dense_by_id, unit_ids + ) term_channel = {uid: float(term_bm25.get(uid, 0.0) or 0.0) for uid in unit_ids} # Convert channel scores to ranked lists for existing RRF merger. @@ -580,7 +633,9 @@ def _rows_from_scores(score_by_id: Dict[str, float]) -> List[dict[str, Any]]: top_k=len(unit_ids), k=rrf_k, ) - fused_by_id = {str(r.get("chunk_id") or ""): float(r.get("score") or 0.0) for r in fused} + fused_by_id = { + str(r.get("chunk_id") or ""): float(r.get("score") or 0.0) for r in fused + } out_rows: List[dict[str, Any]] = [] for uid in unit_ids: row = dict(row_by_id[uid]) @@ -590,3 +645,311 @@ def _rows_from_scores(score_by_id: Dict[str, float]) -> List[dict[str, Any]]: row["term_channel_score"] = float(term_channel.get(uid, 0.0) or 0.0) out_rows.append(row) return out_rows + + +def score_unit_stream_hybrid_all( + unit_factory: Callable[[], Iterable[ScoreUnitRow]], + query: str, +) -> Dict[str, float]: + """Score replayable units without retaining their payloads. + + This is the corpus scorer used by map-nav. It mirrors the active BM25 and + weighted-RRF implementation, but keeps only token statistics, identifiers, + and final scores between bounded provider reads. + """ + return score_unit_stream_hybrid_many(unit_factory, [query]).get(query, {}) + + +def score_unit_stream_hybrid_many( + unit_factory: Callable[[], Iterable[ScoreUnitRow]], + queries: Sequence[str], +) -> Dict[str, Dict[str, float]]: + """Score several queries exactly while reading the corpus only once.""" + unique_queries = list(dict.fromkeys(str(query) for query in queries)) + if not unique_queries: + return {} + + query_tokens_by_query = { + query: tokenize_query_for_ranker(query) for query in unique_queries + } + query_token_set = { + token + for query_tokens in query_tokens_by_query.values() + for token in query_tokens + } + query_lower_by_query = {query: query.lower().strip() for query in unique_queries} + units: List[_StreamingManyUnit] = [] + 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) + 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, + path_length=len(path_tokens), + content_length=len(content_tokens), + path_frequencies={ + token: path_frequencies[token] + for token in query_token_set + if path_frequencies[token] + }, + content_frequencies={ + token: content_frequencies[token] + for token in query_token_set + if content_frequencies[token] + }, + term_scores=tuple(term_scores), + ) + ) + path_stats.finalize() + content_stats.finalize() + return { + query: _score_streaming_units( + units, + path_stats=path_stats, + content_stats=content_stats, + query_tokens=query_tokens_by_query[query], + query_index=index, + ) + for index, query in enumerate(unique_queries) + } + + +def _score_streaming_units( + units: Sequence["_StreamingManyUnit"], + *, + path_stats: "_StreamingBm25Stats", + content_stats: "_StreamingBm25Stats", + query_tokens: List[str], + query_index: int, +) -> Dict[str, float]: + path_by_id: Dict[str, float] = {} + content_by_id: Dict[str, float] = {} + term_by_id: Dict[str, float] = {} + unit_ids = list(dict.fromkeys(unit.unit_id for unit in units)) + for unit in units: + path_score = path_stats.score( + unit.path_length, unit.path_frequencies, query_tokens + ) + content_score = content_stats.score( + unit.content_length, unit.content_frequencies, query_tokens + ) + path_by_id[unit.unit_id] = path_score + content_by_id[unit.unit_id] = content_score + term_by_id[unit.unit_id] = 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 + ] + content_rows = [ + (score, unit_id) for unit_id, score in content_by_id.items() if score > 0.0 + ] + term_rows = [ + (score, unit_id) for unit_id, score in term_by_id.items() if score > 0.0 + ] + path_rows.sort(key=lambda item: (-item[0], item[1])) + content_rows.sort(key=lambda item: (-item[0], item[1])) + term_rows.sort(key=lambda item: (-item[0], item[1])) + path_weight, content_weight, term_weight = map_channel_weights() + rrf_k = int( + os.environ.get( + "NAV_MAP_RRF_K", + os.environ.get("NAV_DISCOVERY_RRF_K", str(RRF_K)), + ).strip() + or RRF_K + ) + fused: Dict[str, float] = {unit_id: 0.0 for unit_id in unit_ids} + for rank, (_score, unit_id) in enumerate(path_rows): + fused[unit_id] = fused.get(unit_id, 0.0) + path_weight / (rrf_k + rank + 1) + for rank, (_score, unit_id) in enumerate(content_rows): + fused[unit_id] = fused.get(unit_id, 0.0) + content_weight / (rrf_k + rank + 1) + for rank, (_score, unit_id) in enumerate(term_rows): + fused[unit_id] = fused.get(unit_id, 0.0) + term_weight / (rrf_k + rank + 1) + return {unit_id: round(score, 6) for unit_id, score in fused.items()} + + +@dataclass(frozen=True) +class _StreamingManyUnit: + unit_id: str + path_length: int + content_length: int + path_frequencies: Mapping[str, int] + content_frequencies: Mapping[str, int] + term_scores: Tuple[float, ...] + + +@dataclass(frozen=True) +class PersistedBm25Stats: + """Corpus statistics needed to reproduce ``BM25Okapi`` exactly.""" + + document_count: int + total_length: int + document_frequency: Mapping[str, int] + average_idf: float + + +@dataclass(frozen=True) +class PersistedScoreUnit: + """Query-specific frequencies for one persisted map unit.""" + + unit_id: str + path_length: int + content_length: int + path_frequencies: Mapping[str, int] + content_frequencies: Mapping[str, int] + term_scores: Tuple[float, ...] + + +@dataclass(frozen=True) +class PersistedScoreCorpus: + """Compact query projection loaded from the map-unit index.""" + + units: Sequence[PersistedScoreUnit] + path_stats: PersistedBm25Stats + content_stats: PersistedBm25Stats + + +def score_persisted_corpus_many( + corpus: PersistedScoreCorpus, + queries: Sequence[str], +) -> Dict[str, Dict[str, float]]: + """Apply the existing BM25/RRF scorer to persisted query projections.""" + unique_queries = list(dict.fromkeys(str(query) for query in queries)) + if not unique_queries: + return {} + path_stats = _restore_bm25_stats(corpus.path_stats) + content_stats = _restore_bm25_stats(corpus.content_stats) + units = [ + _StreamingManyUnit( + unit_id=unit.unit_id, + path_length=unit.path_length, + content_length=unit.content_length, + path_frequencies=unit.path_frequencies, + content_frequencies=unit.content_frequencies, + term_scores=unit.term_scores, + ) + for unit in corpus.units + ] + return { + query: _score_streaming_units( + units, + path_stats=path_stats, + content_stats=content_stats, + query_tokens=tokenize_query_for_ranker(query), + query_index=query_index, + ) + for query_index, query in enumerate(unique_queries) + } + + +def _restore_bm25_stats(source: PersistedBm25Stats) -> "_StreamingBm25Stats": + stats = _StreamingBm25Stats.empty() + stats.document_count = source.document_count + stats.total_length = source.total_length + stats.average_length = ( + source.total_length / source.document_count if source.document_count else 0.0 + ) + idf_by_token: Dict[str, float] = {} + for token, frequency in source.document_frequency.items(): + idf = math.log(source.document_count - frequency + 0.5) - math.log( + frequency + 0.5 + ) + idf_by_token[token] = 0.25 * source.average_idf if idf < 0.0 else idf + stats.idf_by_token = idf_by_token + return stats + + +class _StreamingBm25Stats: + """Exact BM25Okapi corpus statistics collected without row retention.""" + + def __init__(self) -> None: + self.document_count: int = 0 + self.document_frequency: Counter[str] = Counter() + self.total_length: int = 0 + self.average_length: float = 0.0 + self.idf_by_token: Dict[str, float] = {} + + @classmethod + def empty(cls) -> "_StreamingBm25Stats": + return cls() + + def observe(self, tokens: List[str]) -> None: + if not tokens: + return + self.document_count += 1 + self.total_length += len(tokens) + self.document_frequency.update(set(tokens)) + + def finalize(self) -> None: + self.average_length = ( + self.total_length / self.document_count if self.document_count else 0.0 + ) + idf_by_token: Dict[str, float] = {} + idf_sum = 0.0 + negative_tokens: List[str] = [] + for token, frequency in self.document_frequency.items(): + idf = math.log(self.document_count - frequency + 0.5) - math.log( + frequency + 0.5 + ) + idf_by_token[token] = idf + idf_sum += idf + if idf < 0.0: + negative_tokens.append(token) + average_idf = idf_sum / len(idf_by_token) if idf_by_token else 0.0 + epsilon_floor = 0.25 * average_idf + for token in negative_tokens: + idf_by_token[token] = epsilon_floor + self.idf_by_token = idf_by_token + + def score( + self, + document_length: int, + frequencies: Dict[str, int], + query_tokens: List[str], + ) -> float: + if ( + not frequencies + or not query_tokens + or not self.document_count + or self.average_length <= 0.0 + ): + return 0.0 + denominator_base = 1.5 * ( + 1.0 - 0.75 + 0.75 * document_length / self.average_length + ) + score = 0.0 + for token in query_tokens: + frequency = frequencies.get(token, 0) + if not frequency: + continue + idf = self.idf_by_token.get(token, 0.0) + score += idf * (frequency * 2.5 / (frequency + denominator_base)) + return score diff --git a/packages/shared-python/shared/services/retrieval/nav/nav_hierarchy.py b/packages/shared-python/shared/services/retrieval/nav/nav_hierarchy.py index fafa66d22..d9e42b075 100644 --- a/packages/shared-python/shared/services/retrieval/nav/nav_hierarchy.py +++ b/packages/shared-python/shared/services/retrieval/nav/nav_hierarchy.py @@ -22,7 +22,22 @@ from __future__ import annotations from dataclasses import dataclass, field -from typing import Any, Dict, List, Optional, Protocol, Sequence, Set, Tuple, runtime_checkable +from typing import ( + Any, + Dict, + List, + Optional, + Protocol, + Sequence, + Set, + Tuple, + TYPE_CHECKING, + cast, + runtime_checkable, +) + +if TYPE_CHECKING: + from .knowhere_hybrid import PersistedScoreCorpus @dataclass @@ -170,7 +185,9 @@ def _node_unit_span(self, section_id: str) -> Tuple[str, int, int]: first_order = int(getattr(units[0], "sort_order", 0) or 0) return "\n".join(texts), first_order, len(units) - def _make_chunk(self, node_id: str, doc_id: str, text: str, order: int, section_id: str) -> Any: + def _make_chunk( + self, node_id: str, doc_id: str, text: str, order: int, section_id: str + ) -> Any: from ._compat import Chunk # type: ignore return Chunk( @@ -208,7 +225,9 @@ def _materialize_leaf_path_chunks(self, section_id: str, doc_id: str) -> List[An text = str(self._provider.content(section_id) or "") if not text.strip(): return [] - return [self._make_chunk(f"{section_id}__path", doc_id, text, 0, section_id)] + return [ + self._make_chunk(f"{section_id}__path", doc_id, text, 0, section_id) + ] # One unit per descendant leaf, plus one per interstitial parent, so # node ids line up with the keys nav_map_scores.build_score_units emits. @@ -227,10 +246,62 @@ def _materialize_leaf_path_chunks(self, section_id: str, doc_id: str) -> List[An out.sort(key=lambda c: (min(c.line_ids or (0,)), c.node_id)) return out - def read_chunks(self, section_id: str, query: str, *, doc_id: str, k: int) -> List[Any]: + def read_chunks( + self, section_id: str, query: str, *, doc_id: str, k: int + ) -> List[Any]: del section_id, query, doc_id, k return [] + 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], + queries: Sequence[str], + ) -> Optional["PersistedScoreCorpus"]: + """Forward the optional revision-pinned map-unit index capability.""" + fn = getattr(self._provider, "load_persisted_score_corpus", None) + if not callable(fn): + return None + return cast(Optional["PersistedScoreCorpus"], fn(doc_ids, queries)) + @dataclass class InMemoryNode: diff --git a/packages/shared-python/shared/services/retrieval/nav/nav_knowhere.py b/packages/shared-python/shared/services/retrieval/nav/nav_knowhere.py index 19caf893b..cf1de7790 100644 --- a/packages/shared-python/shared/services/retrieval/nav/nav_knowhere.py +++ b/packages/shared-python/shared/services/retrieval/nav/nav_knowhere.py @@ -21,16 +21,36 @@ from __future__ import annotations import os +from hashlib import sha256 from dataclasses import dataclass, field -from typing import Any, Dict, List, Optional, Sequence, Set, Tuple +from typing import ( + Any, + Callable, + Dict, + Iterable, + List, + Mapping, + Optional, + Protocol, + Sequence, + Set, + Tuple, +) from .nav_address import NavLevel from .nav_hierarchy import NodeMeta +from .knowhere_hybrid import ( + PersistedBm25Stats, + PersistedScoreCorpus, + PersistedScoreUnit, + tokenize_query_for_ranker, +) _ASSET_TYPES = ("table", "image") # Knowhere sentinel path for the virtual document container (not a collectable leaf). ROOT_SECTION_PATH = "Root" _DEFAULT_DSN = "postgresql://root:root123@127.0.0.1:5433/Knowhere" +_MAP_UNIT_INDEX_FORMAT_VERSION = 1 @dataclass(frozen=True) @@ -127,7 +147,545 @@ def _connect_to_targets(metadata: Dict[str, Any]) -> List[str]: def knowhere_database_url() -> str: - return str(os.environ.get("KNOWHERE_DATABASE_URL") or "").strip() or _DEFAULT_DSN + configured = ( + str(os.environ.get("KNOWHERE_DATABASE_URL") or "").strip() + or str(os.environ.get("DATABASE_URL") or "").strip() + ) + if configured: + # ``ReadOnlyChunkStore`` uses psycopg2's native connector, which + # accepts libpq URLs but not SQLAlchemy's ``+driver`` suffix. + return configured.replace("postgresql+asyncpg://", "postgresql://", 1).replace( + "postgresql+psycopg2://", "postgresql://", 1 + ) + return _DEFAULT_DSN + + +class ChunkStore(Protocol): + def load_persisted_score_corpus( + self, + document_ids: Sequence[str], + allowed_section_ids_by_document: Mapping[str, Sequence[str]], + queries: Sequence[str], + ) -> Optional[PersistedScoreCorpus]: + raise NotImplementedError + + def load_documents_units( + self, + section_ids_by_document: Mapping[str, Sequence[str]], + ) -> Dict[str, List[UnitRow]]: + raise NotImplementedError + + def load_document_units( + self, + document_id: str, + 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, + section_id: str, + extra_chunk_ids: Sequence[str] = (), + ) -> List[UnitRow]: + raise NotImplementedError + + def close(self) -> None: + raise NotImplementedError + + +class ReadOnlyChunkStore: + """Episode-local, revision-pinned loader for lazy map-nav chunks.""" + + def __init__( + self, + *, + dsn: str, + revisions: Dict[str, str], + excluded_sections: Optional[Iterable[Tuple[str, str]]] = None, + ) -> None: + self._dsn = str(dsn) + self._revisions = dict(revisions) + self._excluded_sections = set(excluded_sections or ()) + self._conn: Optional[_SyncConnection] = None + + def _connection(self) -> "_SyncConnection": + if self._conn is None: + self._conn = _connect(self._dsn) + self._conn.set_session(readonly=True, autocommit=True) + return self._conn + + def load_section_units( + self, + document_id: str, + section_id: str, + extra_chunk_ids: Sequence[str] = (), + ) -> List[UnitRow]: + doc_id = str(document_id).strip() + sid = str(section_id).strip() + job_result_id = self._revisions.get(doc_id) + if ( + not doc_id + or not sid + or not job_result_id + or (doc_id, sid) in self._excluded_sections + ): + return [] + cur = self._connection().cursor() + try: + ids = [ + str(chunk_id).strip() + for chunk_id in extra_chunk_ids + if str(chunk_id).strip() + ] + if ids: + cur.execute( + "SELECT chunk_id, section_id, chunk_type, content, sort_order, " + "source_chunk_path, file_path, chunk_metadata " + "FROM document_chunks " + "WHERE document_id = %s AND job_result_id = %s " + "AND (section_id = %s OR chunk_id = ANY(%s)) " + "ORDER BY sort_order, chunk_id, id", + (doc_id, job_result_id, sid, ids), + ) + else: + cur.execute( + "SELECT chunk_id, section_id, chunk_type, content, sort_order, " + "source_chunk_path, file_path, chunk_metadata " + "FROM document_chunks " + "WHERE document_id = %s AND job_result_id = %s AND section_id = %s " + "ORDER BY sort_order, chunk_id, id", + (doc_id, job_result_id, sid), + ) + return [_unit_from_row(row) for row in cur.fetchall()] + finally: + cur.close() + + def load_persisted_score_corpus( + self, + document_ids: Sequence[str], + allowed_section_ids_by_document: Mapping[str, Sequence[str]], + queries: Sequence[str], + ) -> Optional[PersistedScoreCorpus]: + """Load query-relevant score inputs when every revision is indexed.""" + revisions = [ + (document_id, self._revisions[document_id]) + for raw_document_id in document_ids + if (document_id := str(raw_document_id).strip()) in self._revisions + ] + if not revisions or len(revisions) != len(document_ids): + return None + values_sql = ", ".join(["(%s, %s)"] * len(revisions)) + revision_params: List[object] = [ + value for revision in revisions for value in revision + ] + cur = self._connection().cursor() + try: + cur.execute( + "SELECT indexes.document_id, indexes.job_result_id, " + "indexes.format_version, indexes.unit_count, indexes.token_count " + "FROM document_map_unit_indexes AS indexes " + f"JOIN (VALUES {values_sql}) AS revisions(document_id, job_result_id) " + "ON indexes.document_id = revisions.document_id " + "AND indexes.job_result_id = revisions.job_result_id", + revision_params, + ) + manifests = list(cur.fetchall()) + if len(manifests) != len(revisions) or any( + int(row[2]) != _MAP_UNIT_INDEX_FORMAT_VERSION for row in manifests + ): + return None + + cur.execute( + "SELECT COUNT(*), COUNT(DISTINCT (units.document_id, units.unit_id)) " + "FROM document_map_units AS units " + f"JOIN (VALUES {values_sql}) AS revisions(document_id, job_result_id) " + "ON units.document_id = revisions.document_id " + "AND units.job_result_id = revisions.job_result_id", + revision_params, + ) + unit_count_row = cur.fetchone() + indexed_unit_count = int(unit_count_row[0]) if unit_count_row else 0 + distinct_unit_count = int(unit_count_row[1]) if unit_count_row else 0 + expected_count = sum(int(row[3]) for row in manifests) + if indexed_unit_count != expected_count or distinct_unit_count != indexed_unit_count: + return None + cur.execute( + "SELECT COUNT(*) FROM document_map_unit_tokens AS tokens " + "JOIN document_map_units AS units ON units.id = tokens.map_unit_id " + f"JOIN (VALUES {values_sql}) AS revisions(document_id, job_result_id) " + "ON units.document_id = revisions.document_id " + "AND units.job_result_id = revisions.job_result_id", + revision_params, + ) + token_row = cur.fetchone() + indexed_token_count = int(token_row[0]) if token_row else 0 + expected_token_count = sum(int(row[4]) for row in manifests) + if indexed_token_count != expected_token_count: + return None + # The public chunk id is content-derived and may repeat within a + # revision. The persisted scorer keys scores by that id, so use + # the legacy payload path whenever ambiguity would change results. + allowed_by_document = { + str(document_id): {str(section_id) for section_id in section_ids} + for document_id, section_ids in allowed_section_ids_by_document.items() + } + allowed_pairs = [ + (document_id, section_id) + for document_id, section_ids in allowed_by_document.items() + for section_id in section_ids + ] + unit_rows: list[Sequence[object]] = [] + if allowed_pairs: + allowed_document_ids = [pair[0] for pair in allowed_pairs] + allowed_section_ids = [pair[1] for pair in allowed_pairs] + cur.execute( + "SELECT units.id, units.document_id, units.unit_id, units.section_id, " + "units.path_token_count, units.content_token_count " + "FROM document_map_units AS units " + f"JOIN (VALUES {values_sql}) AS revisions(document_id, job_result_id) " + "ON units.document_id = revisions.document_id " + "AND units.job_result_id = revisions.job_result_id " + "JOIN UNNEST(%s::text[], %s::text[]) " + "AS allowed(document_id, section_id) " + "ON units.document_id = allowed.document_id " + "AND units.section_id = allowed.section_id " + "ORDER BY units.document_id, units.sort_order, units.unit_id", + [*revision_params, allowed_document_ids, allowed_section_ids], + ) + unit_rows = list(cur.fetchall()) + map_unit_ids = [str(row[0]) for row in unit_rows] + unique_queries = list(dict.fromkeys(str(query) for query in queries)) + query_tokens_by_query = { + query: tokenize_query_for_ranker(query) for query in unique_queries + } + query_tokens = list( + dict.fromkeys( + token + for query in unique_queries + for token in query_tokens_by_query[query] + ) + ) + frequencies: Dict[Tuple[str, str], Dict[str, int]] = {} + if map_unit_ids and query_tokens: + query_token_hashes = [ + sha256(token.encode("utf-8")).hexdigest() for token in query_tokens + ] + cur.execute( + "SELECT map_unit_id, channel, token, frequency " + "FROM document_map_unit_tokens " + "WHERE map_unit_id = ANY(%s) AND token_hash = ANY(%s) " + "AND token = ANY(%s)", + (map_unit_ids, query_token_hashes, query_tokens), + ) + for map_unit_id, channel, token, frequency in cur.fetchall(): + frequencies.setdefault((str(map_unit_id), str(channel)), {})[ + str(token) + ] = int(frequency) + + term_scores = self._load_term_scores( + cur, + map_unit_ids=map_unit_ids, + queries=unique_queries, + query_tokens_by_query=query_tokens_by_query, + ) + path_stats = self._load_persisted_bm25_stats( + cur, + unit_rows=unit_rows, + map_unit_ids=map_unit_ids, + channel="path", + query_tokens=query_tokens, + frequencies=frequencies, + length_index=4, + ) + content_stats = self._load_persisted_bm25_stats( + cur, + unit_rows=unit_rows, + map_unit_ids=map_unit_ids, + channel="content", + query_tokens=query_tokens, + frequencies=frequencies, + length_index=5, + ) + return PersistedScoreCorpus( + units=[ + PersistedScoreUnit( + unit_id=str(row[2]), + path_length=int(row[4]), + content_length=int(row[5]), + path_frequencies=frequencies.get((str(row[0]), "path"), {}), + content_frequencies=frequencies.get( + (str(row[0]), "content"), {} + ), + term_scores=term_scores.get( + str(row[0]), tuple(0.0 for _query in unique_queries) + ), + ) + for row in unit_rows + ], + path_stats=path_stats, + content_stats=content_stats, + ) + finally: + cur.close() + + def _load_term_scores( + self, + cur: "_SyncCursor", + *, + map_unit_ids: Sequence[str], + queries: Sequence[str], + query_tokens_by_query: Mapping[str, Sequence[str]], + ) -> Dict[str, Tuple[float, ...]]: + if not map_unit_ids or not queries: + return {} + expressions: List[str] = [] + params: List[object] = [] + for query in queries: + query_lower = query.lower().strip() + if not query_lower: + expressions.append("0.0") + continue + token_expressions = [ + "CASE WHEN POSITION(%s IN term_search_text_lower) > 0 THEN 1 ELSE 0 END" + for _token in query_tokens_by_query[query] + ] + token_sum = " + ".join(token_expressions) or "0" + expressions.append( + "CASE WHEN POSITION(%s IN term_search_text_lower) > 0 " + f"THEN 100.0 ELSE ({token_sum})::double precision END" + ) + params.append(query_lower) + params.extend(query_tokens_by_query[query]) + params.append(list(map_unit_ids)) + cur.execute( + "SELECT id, " + ", ".join(expressions) + " " + "FROM document_map_units WHERE id = ANY(%s)", + params, + ) + return { + str(row[0]): tuple(float(value) for value in row[1:]) + for row in cur.fetchall() + } + + def _load_persisted_bm25_stats( + self, + cur: "_SyncCursor", + *, + unit_rows: Sequence[Sequence[object]], + map_unit_ids: Sequence[str], + channel: str, + query_tokens: Sequence[str], + frequencies: Mapping[Tuple[str, str], Mapping[str, int]], + length_index: int, + ) -> PersistedBm25Stats: + lengths = [ + int(row[length_index]) for row in unit_rows if int(row[length_index]) > 0 + ] + document_count = len(lengths) + document_frequency = { + token: sum( + 1 + for row in unit_rows + if frequencies.get((str(row[0]), channel), {}).get(token, 0) > 0 + ) + for token in query_tokens + } + needs_average_idf = any( + frequency > document_count / 2 for frequency in document_frequency.values() + ) + average_idf = 0.0 + if needs_average_idf and map_unit_ids and document_count: + cur.execute( + "SELECT COALESCE(AVG(LN((%s - frequencies.document_frequency + 0.5) " + "/ (frequencies.document_frequency + 0.5))), 0.0) " + "FROM (SELECT token, COUNT(*) AS document_frequency " + "FROM document_map_unit_tokens " + "WHERE map_unit_id = ANY(%s) AND channel = %s " + "GROUP BY token) AS frequencies", + (document_count, list(map_unit_ids), channel), + ) + row = cur.fetchone() + average_idf = float(row[0]) if row else 0.0 + return PersistedBm25Stats( + document_count=document_count, + total_length=sum(lengths), + document_frequency=document_frequency, + average_idf=average_idf, + ) + + def load_document_units( + self, + document_id: str, + 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() + self._conn = None + + +class _SyncCursor(Protocol): + def execute(self, query: str, params: Sequence[object]) -> None: + raise NotImplementedError + + def fetchall(self) -> Sequence[Sequence[object]]: + raise NotImplementedError + + def fetchone(self) -> Optional[Sequence[object]]: + raise NotImplementedError + + def close(self) -> None: + raise NotImplementedError + + +class _SyncConnection(Protocol): + def set_session(self, *, readonly: bool, autocommit: bool) -> None: + raise NotImplementedError + + def cursor(self) -> _SyncCursor: + raise NotImplementedError + + def close(self) -> None: + raise NotImplementedError + + +def _unit_from_row(row: Sequence[object]) -> UnitRow: + return UnitRow( + chunk_id=str(row[0] or ""), + section_id=str(row[1]) if row[1] else None, + chunk_type=str(row[2] or "text"), + content=str(row[3] or ""), + sort_order=int(row[4] or 0), + source_chunk_path=str(row[5] or ""), + file_path=str(row[6] or ""), + metadata=_as_meta(row[7]), + ) class KnowhereProvider: @@ -139,8 +697,12 @@ def __init__( doc_id: str, sections: Sequence[SectionRow], units: Sequence[UnitRow], + lazy_loader: Optional[Callable[[str], Sequence[UnitRow]]] = None, + known_chunk_ids: Optional[Sequence[str]] = None, ) -> None: self.doc_id = str(doc_id) + self._lazy_loader = lazy_loader + self._loaded_sections: Set[str] = set() self._sections: Dict[str, SectionRow] = {s.section_id: s for s in sections} self._children: Dict[str, List[str]] = {} self._roots: List[str] = [] @@ -157,6 +719,12 @@ def __init__( self._units_by_section: Dict[str, List[UnitRow]] = {} self._chunk_ids: Set[str] = set() + if known_chunk_ids: + self._chunk_ids.update( + str(chunk_id).strip() + for chunk_id in known_chunk_ids + if str(chunk_id).strip() + ) for unit in sorted(units, key=lambda u: (u.sort_order, u.chunk_id)): sid = unit.section_id if not sid or sid not in self._sections: @@ -166,6 +734,21 @@ def __init__( self._chunk_ids.add(unit.chunk_id) self._remount_root_assets() + def _ensure_section_loaded(self, section_id: str) -> None: + if self._lazy_loader is None or section_id in self._loaded_sections: + return + loaded = list(self._lazy_loader(section_id) or ()) + self._loaded_sections.add(section_id) + if not loaded: + return + current = self._units_by_section.setdefault(section_id, []) + known = {unit.chunk_id for unit in current} + for unit in loaded: + if unit.chunk_id and unit.chunk_id not in known: + current.append(unit) + known.add(unit.chunk_id) + current.sort(key=lambda unit: (unit.sort_order, unit.chunk_id)) + def _remount_root_assets(self) -> None: """Reattach Root-FK image|table units to host sections via ``connect_to``. @@ -287,12 +870,23 @@ def content(self, section_id: str) -> str: return "\n".join(self.unit_text(u) for u in units if self.unit_text(u)) def self_units(self, section_id: str) -> List[UnitRow]: + self._ensure_section_loaded(section_id) return list(self._units_by_section.get(section_id, ())) + def release_section_units(self, section_id: str) -> None: + """Drop one section's loaded payload while keeping its structure.""" + if self._lazy_loader is None: + return + sid = str(section_id or "").strip() + if not sid: + return + self._units_by_section.pop(sid, None) + self._loaded_sections.discard(sid) + def subtree_units(self, section_id: str) -> List[UnitRow]: - out = list(self._units_by_section.get(section_id, ())) + out = list(self.self_units(section_id)) for cid in self.relations(section_id)[1]: - out.extend(self._units_by_section.get(cid, ())) + out.extend(self.self_units(cid)) out.sort(key=lambda u: (u.sort_order, u.chunk_id)) return out @@ -350,8 +944,114 @@ def summaries(self) -> Dict[str, str]: def all_section_ids(self) -> List[str]: return list(self._sections) + def chunk_count(self) -> int: + return len(self._chunk_ids) -def _connect(dsn: str): + +class LazyKnowhereProvider(KnowhereProvider): + """Hierarchy provider that loads full chunk rows only on first access.""" + + def __init__( + self, + *, + doc_id: str, + sections: Sequence[SectionRow], + chunk_store: ChunkStore, + known_chunk_ids: Sequence[str], + root_asset_ids: Sequence[str] = (), + remounted_assets_by_section: Optional[Dict[str, Sequence[str]]] = None, + ) -> None: + super().__init__( + doc_id=doc_id, + sections=sections, + units=(), + lazy_loader=lambda section_id: chunk_store.load_section_units( + doc_id, + section_id, + (remounted_assets_by_section or {}).get(section_id, ()), + ), + known_chunk_ids=known_chunk_ids, + ) + self._chunk_store = chunk_store + self._root_asset_ids = {str(chunk_id) for chunk_id in root_asset_ids} + 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) + if ( + section_id in self._units_by_section + and self._root_asset_ids + and is_root_section_path(self.section_path(section_id)) + ): + self._units_by_section[section_id] = [ + unit + for unit in self._units_by_section[section_id] + if unit.chunk_id not in self._root_asset_ids + ] + + def close(self) -> None: + self._chunk_store.close() + + def release_loaded_units(self) -> None: + self._units_by_section.clear() + self._loaded_sections.clear() + + +def _connect(dsn: str) -> _SyncConnection: import psycopg2 return psycopg2.connect(dsn) @@ -482,7 +1182,9 @@ def load_namespace_from_db( providers = [load_document_from_db(did, dsn=url) for did in wanted] merged_titles = dict(auto_titles) if titles: - merged_titles.update({str(k): str(v) for k, v in titles.items() if str(k).strip()}) + merged_titles.update( + {str(k): str(v) for k, v in titles.items() if str(k).strip()} + ) return NamespaceKnowhereProvider(providers, titles=merged_titles or None) @@ -499,6 +1201,7 @@ def __init__( providers: Sequence[KnowhereProvider], *, titles: Optional[Dict[str, str]] = None, + chunk_owner_by_id: Optional[Dict[str, str]] = None, ) -> None: self._docs: Dict[str, KnowhereProvider] = { p.doc_id: p for p in providers if p.doc_id @@ -514,14 +1217,127 @@ def __init__( for doc_id, provider in self._docs.items(): for sid in provider.all_section_ids(): self._section_owner[sid] = doc_id - for sid in provider.all_section_ids(): - for unit in provider.self_units(sid): - if unit.chunk_id: - self._chunk_owner[unit.chunk_id] = doc_id + if chunk_owner_by_id: + self._chunk_owner.update( + { + str(chunk_id): str(doc_id) + for chunk_id, doc_id in chunk_owner_by_id.items() + if str(chunk_id).strip() and str(doc_id).strip() + } + ) + else: + for doc_id, provider in self._docs.items(): + for sid in provider.all_section_ids(): + for unit in provider.self_units(sid): + if unit.chunk_id: + self._chunk_owner[unit.chunk_id] = doc_id def document_ids(self) -> List[str]: return list(self._docs) + def close(self) -> None: + for provider in self._docs.values(): + close = getattr(provider, "close", None) + if callable(close): + close() + + def release_loaded_units(self) -> None: + for provider in self._docs.values(): + release = getattr(provider, "release_loaded_units", None) + if callable(release): + release() + + def release_section_units(self, section_id: str) -> None: + owner = self._section_owner.get(str(section_id or "").strip()) + if not owner: + return + release = getattr(self._docs[owner], "release_section_units", None) + if callable(release): + release(section_id) + + def 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], + queries: Sequence[str], + ) -> Optional[PersistedScoreCorpus]: + """Return the index projection only when all documents share one store.""" + providers = [ + self._docs[document_id] + for raw_document_id in doc_ids + if (document_id := str(raw_document_id).strip()) in self._docs + ] + if len(providers) != len(doc_ids) or not all( + isinstance(provider, LazyKnowhereProvider) for provider in providers + ): + return None + lazy_providers = [ + provider + for provider in providers + if isinstance(provider, LazyKnowhereProvider) + ] + stores = { + id(provider._chunk_store): provider._chunk_store + for provider in lazy_providers + } + if len(stores) != 1: + return None + store = next(iter(stores.values())) + loader = getattr(store, "load_persisted_score_corpus", None) + if not callable(loader): + return None + return loader( + [provider.doc_id for provider in lazy_providers], + {provider.doc_id: list(provider._sections) for provider in lazy_providers}, + queries, + ) + + def release_document_units(self, doc_id: str) -> None: + provider = self._docs.get(str(doc_id).strip()) + release = getattr(provider, "release_document_units", None) + if callable(release): + release() + def address_level(self, node_id: str) -> Optional[NavLevel]: sid = str(node_id or "").strip() if not sid: @@ -562,8 +1378,13 @@ def node_meta(self, section_id: str) -> NodeMeta: sid = str(section_id or "").strip() if sid in self._docs: provider = self._docs[sid] - n_chunks = sum( - len(provider.self_units(sec)) for sec in provider.all_section_ids() + count_fn = getattr(provider, "chunk_count", None) + n_chunks = ( + int(count_fn()) + if callable(count_fn) + else sum( + len(provider.self_units(sec)) for sec in provider.all_section_ids() + ) ) return NodeMeta( title=self._titles.get(sid, sid), @@ -595,7 +1416,9 @@ def content(self, section_id: str) -> str: if sid in self._docs: provider = self._docs[sid] return "\n".join( - provider.content(root) for root in provider.roots(sid) if provider.content(root) + provider.content(root) + for root in provider.roots(sid) + if provider.content(root) ) owner = self._section_owner.get(sid) if not owner: diff --git a/packages/shared-python/shared/services/retrieval/nav/nav_map_scores.py b/packages/shared-python/shared/services/retrieval/nav/nav_map_scores.py index d1049adf2..6318f4903 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,22 +1,31 @@ from __future__ import annotations +from collections.abc import Iterator from typing import Any, Dict, List, Optional, Sequence, Set, Tuple from .knowhere_hybrid import ( + ScoreUnitRow, build_content_search_text, build_path_search_text, build_term_search_text, - score_dense_channel, score_rows_hybrid_all, + score_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 + def _children_ids(ts: Any, section_id: str, doc_id: str) -> List[str]: children_fn = getattr(ts, "_children_for_section_path", None) if not callable(children_fn): st = ts.get_structure(section_id) rows = st.get("children") or [] - return [str(r.get("section_id") or "").strip() for r in rows if r.get("section_id")] + return [ + str(r.get("section_id") or "").strip() for r in rows if r.get("section_id") + ] rows = children_fn(section_id, doc_id, limit=100000) return [str(r.get("section_id") or "").strip() for r in rows if r.get("section_id")] @@ -26,6 +35,10 @@ def _line_content(ts: Any, section_id: str, doc_id: str) -> str: idx = getattr(ts, "_idx", None) b = getattr(idx, "_bundles", {}).get(doc_id) if idx is not None else None if b is None: + path_fn = getattr(ts, "path_titles", None) + if callable(path_fn): + path = str(path_fn(section_id, doc_id) or "").strip() + return path.rsplit(" / ", 1)[-1] if path else "" st = ts.get_structure(section_id) return str(st.get("preview") or "").strip() loc = getattr(idx, "_node_to_doc_line", {}).get(section_id) @@ -142,8 +155,7 @@ def _pool_unit_scores_to_tree( ) -> Dict[str, float]: """MAX-pool globally comparable unit scores onto one document tree.""" map_scores = { - leaf_id: float(unit_scores.get(leaf_id, 0.0) or 0.0) - for leaf_id in leaves + leaf_id: float(unit_scores.get(leaf_id, 0.0) or 0.0) for leaf_id in leaves } def score_node(section_id: str) -> float: @@ -154,12 +166,9 @@ def score_node(section_id: str) -> float: score = float(unit_scores.get(section_id, 0.0) or 0.0) map_scores[section_id] = score return score - descendant_leaves = _collect_descendant_leaves( - section_id, children_map, leaves - ) + descendant_leaves = _collect_descendant_leaves(section_id, children_map, leaves) parts = [ - float(unit_scores.get(leaf_id, 0.0) or 0.0) - for leaf_id in descendant_leaves + float(unit_scores.get(leaf_id, 0.0) or 0.0) for leaf_id in descendant_leaves ] self_key = f"{section_id}__self" if self_key in unit_scores: @@ -173,48 +182,9 @@ def score_node(section_id: str) -> float: return map_scores -def _score_dense_units_by_doc( - units_by_doc: Sequence[Tuple[str, List[dict]]], - query: str, - *, - namespace: Optional[str], -) -> Dict[str, Optional[Dict[str, float]]]: - """Read per-doc vector caches, returning raw cosine scores for global fusion. - - Dense cosine is independently comparable across documents. Partitioning only - preserves the existing per-doc disk cache; no ranking or normalization occurs - here. If any partition fails, that whole channel falls back to global BM25. - """ - dense_by_channel: Dict[str, Optional[Dict[str, float]]] = {} - for channel, text_field in (("path", "path_text"), ("content", "content")): - score_by_id: Dict[str, float] = {} - complete = True - for doc_id, units in units_by_doc: - if not units: - continue - unit_ids = [str(unit["chunk_id"]) for unit in units] - scores = score_dense_channel( - [str(unit.get(text_field) or "") for unit in units], - query, - unit_ids=unit_ids, - doc_id=doc_id, - channel=channel, - namespace=namespace, - ) - if scores is None or len(scores) != len(unit_ids): - complete = False - break - score_by_id.update( - { - unit_id: float(scores[index]) - for index, unit_id in enumerate(unit_ids) - } - ) - dense_by_channel[channel] = score_by_id if complete else None - return dense_by_channel - - -def build_score_units(ts: Any, doc_id: str, root_ids: Optional[Sequence[str]] = None) -> List[dict]: +def build_score_units( + ts: Any, doc_id: str, root_ids: Optional[Sequence[str]] = None +) -> List[dict]: """Build leaf (+ interstitial self_only) units for hybrid scoring.""" if root_ids is None: root_ids = list(ts.sections_for_doc(doc_id)) @@ -243,7 +213,9 @@ def build_score_units(ts: Any, doc_id: str, root_ids: Optional[Sequence[str]] = section_path=path_text, section_title=title or content ), "content_search_text": build_content_search_text(content), - "term_search_text": build_term_search_text(content, path_text=path_text), + "term_search_text": build_term_search_text( + content, path_text=path_text + ), } ) @@ -270,12 +242,78 @@ def build_score_units(ts: Any, doc_id: str, root_ids: Optional[Sequence[str]] = section_path=path_text, section_title=titles.get(sid) or "" ), "content_search_text": build_content_search_text(self_text), - "term_search_text": build_term_search_text(self_text, path_text=path_text), + "term_search_text": build_term_search_text( + self_text, path_text=path_text + ), } ) return units +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, *, @@ -320,7 +358,9 @@ def compute_map_and_unit_scores( doc_id=doc_id, namespace=ns, ) - unit_score = {str(r.get("chunk_id") or ""): float(r.get("score") or 0.0) for r in scored} + unit_score = { + str(r.get("chunk_id") or ""): float(r.get("score") or 0.0) for r in scored + } map_scores = _pool_unit_scores_to_tree(children_map, leaves, unit_score) return map_scores, unit_score @@ -337,6 +377,26 @@ def compute_corpus_map_and_unit_scores( All documents share one BM25 corpus, path/content normalization, channel ranking, and RRF pass. Document-level scores are keyed by bare ``document_id``. """ + return compute_corpus_map_and_unit_scores_many( + ts, + doc_ids=doc_ids, + queries=[query], + namespace=namespace, + ).get(query, ({}, {})) + + +def compute_corpus_map_and_unit_scores_many( + ts: Any, + *, + doc_ids: Sequence[str], + queries: Sequence[str], + namespace: Optional[str] = None, +) -> Dict[str, Tuple[Dict[str, float], Dict[str, float]]]: + """Globally score several queries with one replay of the corpus units.""" + unique_queries = list(dict.fromkeys(str(query) for query in queries)) + if not unique_queries: + return {} + valid_doc_ids: List[str] = [] seen_doc_ids: Set[str] = set() for raw in doc_ids: @@ -346,51 +406,91 @@ def compute_corpus_map_and_unit_scores( seen_doc_ids.add(doc_id) valid_doc_ids.append(doc_id) - ns = namespace - if not ns: - import os - - ns = os.environ.get("NAV_MAP_UNIT_CACHE_NS", "").strip() or None + del namespace # Dense scoring is intentionally disabled for the corpus path. - tree_by_doc: Dict[str, Tuple[Dict[str, List[str]], Set[str]]] = {} - units_by_doc: List[Tuple[str, List[dict]]] = [] - all_units: List[dict] = [] + tree_by_doc: Dict[ + str, + Tuple[Dict[str, List[str]], Set[str], Dict[str, str]], + ] = {} for doc_id in valid_doc_ids: root_ids = list(ts.sections_for_doc(doc_id)) - children_map, leaves, _titles = _walk_tree(ts, doc_id, root_ids) - units = build_score_units(ts, doc_id, root_ids=root_ids) - tree_by_doc[doc_id] = (children_map, leaves) - units_by_doc.append((doc_id, units)) - all_units.extend(units) - - dense_scores = _score_dense_units_by_doc( - units_by_doc, - query, - namespace=ns, + children_map, leaves, titles = _walk_tree(ts, doc_id, root_ids) + tree_by_doc[doc_id] = (children_map, leaves, titles) + + def unit_factory() -> Iterator[ScoreUnitRow]: + 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) + persisted_corpus = ( + persisted_loader(valid_doc_ids, unique_queries) + if callable(persisted_loader) + else None ) - scored = score_rows_hybrid_all( - all_units, - query, - dense_scores_by_channel=dense_scores, + 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) ) - unit_scores = { - str(row.get("chunk_id") or ""): float(row.get("score") or 0.0) - for row in scored - } - - map_scores: Dict[str, float] = {} - for doc_id in valid_doc_ids: - children_map, leaves = tree_by_doc[doc_id] - doc_map_scores = _pool_unit_scores_to_tree( - children_map, leaves, unit_scores - ) - map_scores.update(doc_map_scores) - doc_max = max( - (float(value) for value in doc_map_scores.values()), - default=0.0, - ) - map_scores[doc_id] = doc_max - return map_scores, unit_scores + results: Dict[str, Tuple[Dict[str, float], Dict[str, float]]] = {} + for query in unique_queries: + unit_scores = unit_scores_by_query.get(query, {}) + map_scores: Dict[str, float] = {} + for doc_id in valid_doc_ids: + children_map, leaves, _titles = tree_by_doc[doc_id] + doc_map_scores = _pool_unit_scores_to_tree( + children_map, leaves, unit_scores + ) + map_scores.update(doc_map_scores) + doc_max = max( + (float(value) for value in doc_map_scores.values()), + default=0.0, + ) + map_scores[doc_id] = doc_max + results[query] = (map_scores, unit_scores) + return results def unit_id_to_section_id(unit_id: str) -> str: @@ -449,3 +549,44 @@ def relight_map_for_query( ts, doc_ids=doc_ids, query=query ) return map_scores, unit_scores, select_map_highlights(unit_scores, k=int(top_k)) + + +def relight_maps_for_queries( + ts: Any, + *, + doc_id: str, + queries: Sequence[str], + top_k: int = 6, +) -> Dict[str, Tuple[Dict[str, float], Dict[str, float], List[str]]]: + """Re-score a shared map for several queries with one corpus replay.""" + unique_queries = list(dict.fromkeys(str(query) for query in queries)) + if not unique_queries: + return {} + doc = str(doc_id or "").strip() + if doc: + return { + query: relight_map_for_query( + ts, + doc_id=doc, + query=query, + top_k=top_k, + ) + for query in unique_queries + } + + doc_ids = [str(value) for value in (ts.document_ids() or ()) if str(value).strip()] + if not doc_ids: + return {} + scored = compute_corpus_map_and_unit_scores_many( + ts, + doc_ids=doc_ids, + queries=unique_queries, + ) + return { + query: ( + map_scores, + unit_scores, + select_map_highlights(unit_scores, k=int(top_k)), + ) + for query, (map_scores, unit_scores) in scored.items() + } diff --git a/packages/shared-python/shared/services/retrieval/nav/nav_orchestrate.py b/packages/shared-python/shared/services/retrieval/nav/nav_orchestrate.py index 1e30b8f4c..91f45e1c5 100644 --- a/packages/shared-python/shared/services/retrieval/nav/nav_orchestrate.py +++ b/packages/shared-python/shared/services/retrieval/nav/nav_orchestrate.py @@ -126,6 +126,16 @@ def _unbound_retrieval_query(subgoal: Subgoal) -> str: return raw or (subgoal.need or "").strip() or subgoal.retrieval_query +def _resolve_subgoal_query(state: NavState, subgoal: Subgoal) -> str: + query = bind_slots(subgoal.retrieval_query, state.slot_bindings) + if unbound_slots(query): + query = _unbound_retrieval_query(subgoal) + refined = str( + (state.subgoal_refined_queries or {}).get(subgoal.id) or "" + ).strip() + return refined or query + + def _run_navigate_for_query( ts: Any, state: NavState, @@ -180,6 +190,7 @@ def _relit_map( config: NavConfig, *, query: str, + prepared: Optional[Tuple[Dict[str, float], Dict[str, float], List[str]]] = None, ) -> Iterator[None]: """Score the shared map against the harvest ``query`` for one call. @@ -189,9 +200,9 @@ def _relit_map( the query the policy is told to pursue. Scoring failures degrade to the episode lighting. """ - relit: Optional[Tuple[Dict[str, float], Dict[str, float], List[str]]] = None + relit = prepared q = (query or "").strip() - if q: + if relit is None and q: try: from .nav_map_scores import relight_map_for_query @@ -224,6 +235,10 @@ def _execute_subgoal_harvest_once( subgoal: Subgoal, *, steps_out: Optional[List[Any]], + retrieval_query: Optional[str] = None, + prepared_relight: Optional[ + Tuple[Dict[str, float], Dict[str, float], List[str]] + ] = None, ) -> Dict[str, Any]: """One harvest() call for this subgoal this wave — no internal retry loop. @@ -232,22 +247,21 @@ def _execute_subgoal_harvest_once( """ from .nav_harvest import harvest - rq = bind_slots(subgoal.retrieval_query, state.slot_bindings) - if unbound_slots(rq): - # F1: deps may be "settled" (satisfied or dropped) without ever - # producing this subgoal's referenced slot — degrade to a query with - # the unresolved {{...}} braces stripped rather than stalling. - rq = _unbound_retrieval_query(subgoal) + rq = retrieval_query or _resolve_subgoal_query(state, subgoal) refined = str((state.subgoal_refined_queries or {}).get(subgoal.id) or "").strip() - if refined: - rq = refined _set_focus(state, subgoal, rq) # Always enter at namespace/document root; prior dead-ends stay hidden via # subgoal_dismissed_section_ids so the next harvest sees siblings instead. before_sections = set(state.collected_section_ids) before_explicit = set(state.explicit_collect_ids) before_len = len(state.collected) - with _relit_map(ts, state, config, query=rq): + with _relit_map( + ts, + state, + config, + query=rq, + prepared=prepared_relight, + ): harvest_result = harvest( ts, state, @@ -433,10 +447,39 @@ def execute_plan( by_id = {s.id: s for s in plan.subgoals} outputs: List[Dict[str, Any]] = [] + query_by_subgoal = { + sid: _resolve_subgoal_query(state, by_id[sid]) for sid in ready + } + prepared_relights: Dict[ + str, + Tuple[Dict[str, float], Dict[str, float], List[str]], + ] = {} + try: + from .nav_map_scores import relight_maps_for_queries + + prepared_relights = relight_maps_for_queries( + ts, + doc_id=state.doc_id, + queries=list(query_by_subgoal.values()), + top_k=int(config.collect_top_k), + ) + except Exception: + prepared_relights = {} def _run_one(sid: str, working_state: NavState, out_steps: Optional[List[Any]]) -> Dict[str, Any]: + query = query_by_subgoal[sid] + prepared = prepared_relights.get(query) + if prepared is not None and not prepared[0]: + prepared = None return _execute_subgoal_harvest_once( - ts, working_state, config, plan, by_id[sid], steps_out=out_steps + ts, + working_state, + config, + plan, + by_id[sid], + steps_out=out_steps, + retrieval_query=query, + prepared_relight=prepared, ) # Serial wave execution (parallel fan-out retired with ThreadPoolExecutor). diff --git a/packages/shared-python/shared/services/retrieval/nav_snapshot.py b/packages/shared-python/shared/services/retrieval/nav_snapshot.py index 7dc60c6f3..1edfd3748 100644 --- a/packages/shared-python/shared/services/retrieval/nav_snapshot.py +++ b/packages/shared-python/shared/services/retrieval/nav_snapshot.py @@ -7,6 +7,7 @@ from __future__ import annotations +import json from dataclasses import dataclass from typing import Any, Protocol @@ -16,10 +17,13 @@ from shared.models.database.document import Document, DocumentChunk, DocumentSection from shared.models.database.job_result import JobResult from shared.services.retrieval.nav.nav_knowhere import ( + LazyKnowhereProvider, KnowhereProvider, NamespaceKnowhereProvider, + ReadOnlyChunkStore, SectionRow, UnitRow, + knowhere_database_url, ) from shared.services.retrieval.search.section_filters import is_excluded_section @@ -48,6 +52,11 @@ class NavSnapshot: document_ids: list[str] document_titles: dict[str, str] + def close(self) -> None: + close = getattr(self.provider, "close", None) + if callable(close): + close() + def build_nav_snapshot( *, @@ -86,6 +95,7 @@ async def load_nav_snapshot( namespace: str, exclude_document_ids: list[str] | None = None, exclude_sections: list[dict[str, str]] | None = None, + lazy: bool = False, ) -> NavSnapshot: """Preload namespace current revision into a sync map-nav snapshot.""" excluded_docs = [str(x).strip() for x in (exclude_document_ids or ()) if str(x).strip()] @@ -136,6 +146,71 @@ async def load_nav_snapshot( document_revisions=document_revisions, exclude_sections=excluded_secs, ) + if lazy: + chunk_ids_by_doc, chunk_ref_index, remounted_assets = await _load_chunk_index( + db, + document_revisions=document_revisions, + exclude_sections=excluded_secs, + section_path_by_id=section_path_by_id, + job_id_by_result_id=job_id_by_result_id, + ) + kept_titles = { + did: title for did, title in document_titles.items() if sections_by_doc.get(did) + } + if not kept_titles: + raise ValueError( + f"nav snapshot empty after excludes for " + f"user_id={user_id!r} namespace={namespace!r}" + ) + revisions = {did: result_id for did, result_id in document_revisions if did in kept_titles} + store = ReadOnlyChunkStore( + dsn=knowhere_database_url(), + revisions=revisions, + excluded_sections={ + ( + str(item.get("document_id") or "").strip(), + section_id, + ) + for section_id, section_path in section_path_by_id.items() + for item in excluded_secs + if isinstance(item, dict) + and str(item.get("document_id") or "").strip() + and str(item.get("section_path") or "").strip() == section_path + }, + ) + try: + providers = [ + LazyKnowhereProvider( + doc_id=did, + sections=sections_by_doc.get(did, ()), + chunk_store=store, + known_chunk_ids=chunk_ids_by_doc.get(did, ()), + root_asset_ids=remounted_assets.get(did, {}).get("root", ()), + remounted_assets_by_section=remounted_assets.get(did, {}).get("owners", {}), + ) + for did in kept_titles + ] + provider = NamespaceKnowhereProvider( + providers, + titles=kept_titles, + chunk_owner_by_id={ + chunk_id: meta["document_id"] + for chunk_id, meta in chunk_ref_index.items() + if ":" not in chunk_id + and isinstance(meta, dict) + and meta.get("document_id") + }, + ) + except Exception: + store.close() + raise + return NavSnapshot( + provider=provider, + chunk_ref_index=dict(chunk_ref_index), + document_ids=list(provider.document_ids()), + document_titles={did: kept_titles.get(did, did) for did in provider.document_ids()}, + ) + units_by_doc, chunk_ref_index = await _load_chunks( db, document_revisions=document_revisions, @@ -164,6 +239,126 @@ async def load_nav_snapshot( ) +async def _load_chunk_index( + db: SnapshotSession, + *, + document_revisions: list[tuple[str, str]], + exclude_sections: list[dict[str, str]], + section_path_by_id: dict[str, str], + job_id_by_result_id: dict[str, str], +) -> tuple[dict[str, list[str]], dict[str, dict[str, Any]], dict[str, dict[str, Any]]]: + """Load only IDs/reference metadata; content remains lazy.""" + ids_by_doc: dict[str, list[str]] = {} + ref_index: dict[str, dict[str, Any]] = {} + root_assets_by_doc: dict[str, set[str]] = {} + text_connections_by_doc: dict[str, list[tuple[str, str]]] = {} + for group_start in range(0, len(document_revisions), _REVISION_GROUP_SIZE): + revision_group = document_revisions[group_start : group_start + _REVISION_GROUP_SIZE] + last_key: tuple[str, str, int, str, str] | None = None + while True: + stmt = ( + select( + DocumentChunk.document_id, + DocumentChunk.job_result_id, + DocumentChunk.chunk_id, + DocumentChunk.section_id, + DocumentChunk.chunk_type, + DocumentChunk.file_path, + DocumentChunk.chunk_metadata["connect_to"].label("connect_to"), + DocumentChunk.sort_order, + DocumentChunk.id, + ) + .where(tuple_(DocumentChunk.document_id, DocumentChunk.job_result_id).in_(revision_group)) + .order_by( + DocumentChunk.document_id, + DocumentChunk.job_result_id, + DocumentChunk.sort_order, + DocumentChunk.chunk_id, + DocumentChunk.id, + ) + .limit(_CHUNK_BATCH_SIZE) + ) + if last_key is not None: + stmt = stmt.where( + tuple_( + DocumentChunk.document_id, + DocumentChunk.job_result_id, + DocumentChunk.sort_order, + DocumentChunk.chunk_id, + DocumentChunk.id, + ) + > tuple_(*[literal(value) for value in last_key]) + ) + rows = (await db.execute(stmt)).all() + if not rows: + break + for row in rows: + document_id = str(row[0]) + job_result_id = str(row[1]) + chunk_id = str(row[2] or "").strip() + section_id = str(row[3]) if row[3] else None + section_path = section_path_by_id.get(section_id) if section_id else None + if is_excluded_section( + document_id=document_id, + section_path=section_path, + exclude_sections=exclude_sections, + ) or (section_id and section_id not in section_path_by_id): + continue + if not chunk_id: + continue + chunk_type = str(row[4] or "text") + meta = { + "document_id": document_id, + "section_path": section_path, + "chunk_type": chunk_type, + "file_path": str(row[5] or "") or None, + "job_id": job_id_by_result_id.get(job_result_id), + } + ids_by_doc.setdefault(document_id, []).append(chunk_id) + ref_index[chunk_id] = meta + ref_index[f"{document_id}:{chunk_id}"] = meta + if ( + chunk_type in {"image", "table"} + and section_id + and section_path == "Root" + ): + root_assets_by_doc.setdefault(document_id, set()).add(chunk_id) + connections = row[6] + if isinstance(connections, str) and connections.strip(): + try: + connections = json.loads(connections) + except json.JSONDecodeError: + connections = None + if chunk_type == "text" and isinstance(connections, list): + for connection in connections: + if not isinstance(connection, dict): + continue + target = str(connection.get("target") or "").strip() + if not target: + continue + text_connections_by_doc.setdefault(document_id, []).append( + (section_id or "", target) + ) + last = rows[-1] + last_key = ( + str(last[0]), + str(last[1]), + int(last[7] or 0), + str(last[2] or ""), + str(last[8]), + ) + if len(rows) < _CHUNK_BATCH_SIZE: + break + remounted: dict[str, dict[str, Any]] = {} + for document_id, asset_ids in root_assets_by_doc.items(): + owners: dict[str, list[str]] = {} + for section_id, target in text_connections_by_doc.get(document_id, ()): + if target in asset_ids: + owners.setdefault(section_id, []).append(target) + remounted[document_id] = {"root": sorted(asset_ids), "owners": owners} + return ids_by_doc, ref_index, remounted + + async def _load_sections( db: SnapshotSession, *, diff --git a/packages/shared-python/shared/services/retrieval/publication_content.py b/packages/shared-python/shared/services/retrieval/publication_content.py index 9ebbdd74a..19a7f2074 100644 --- a/packages/shared-python/shared/services/retrieval/publication_content.py +++ b/packages/shared-python/shared/services/retrieval/publication_content.py @@ -6,7 +6,13 @@ from sqlalchemy import delete from sqlalchemy.orm import Session -from shared.models.database.document import DocumentChunk, DocumentSection +from shared.models.database.document import ( + DocumentChunk, + DocumentMapUnit, + DocumentMapUnitIndex, + DocumentSection, +) +from shared.services.retrieval.map_unit_index import replace_document_map_units from shared.services.retrieval.publication_models import DocumentPublicationScope from shared.services.retrieval.search.lexical_text import ( build_content_lexical_text, @@ -57,7 +63,9 @@ def replace_document_revision_content( """Replace retrieval sections and chunks for one published document revision.""" _delete_existing_revision_content(db, scope=scope) section_publisher = DocumentSectionPublisher( - db=db, scope=scope, section_summaries=section_summaries, + db=db, + scope=scope, + section_summaries=section_summaries, ) for index, chunk in enumerate(chunks): safe_chunk = cast(dict[str, Any], remove_nul_characters(chunk)) @@ -81,6 +89,8 @@ def replace_document_revision_content( fallback_sort_order=index, ) ) + db.flush() + replace_document_map_units(db, scope=scope) class DocumentSectionPublisher: @@ -146,6 +156,16 @@ def _delete_existing_revision_content( *, scope: DocumentPublicationScope, ) -> None: + db.execute( + delete(DocumentMapUnitIndex) + .where(DocumentMapUnitIndex.document_id == scope.document_id) + .where(DocumentMapUnitIndex.job_result_id == scope.job_result_id) + ) + db.execute( + delete(DocumentMapUnit) + .where(DocumentMapUnit.document_id == scope.document_id) + .where(DocumentMapUnit.job_result_id == scope.job_result_id) + ) db.execute( delete(DocumentChunk) .where(DocumentChunk.document_id == scope.document_id)