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 00000000..300dc99c --- /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/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 00000000..782f01f9 --- /dev/null +++ b/apps/api/tests/contract/test_retrieval_lazy_snapshot_quality_contract.py @@ -0,0 +1,176 @@ +from __future__ import annotations + +from dataclasses import dataclass +from collections.abc import Sequence +from typing import Any + +from shared.services.retrieval.nav.nav_hierarchy import ProviderToolSpace +from shared.services.retrieval.nav.nav_knowhere import ( + KnowhereProvider, + LazyKnowhereProvider, + NamespaceKnowhereProvider, + SectionRow, + UnitRow, +) +from shared.services.retrieval.nav.nav_map_scores import ( + build_score_units, + compute_corpus_map_and_unit_scores, +) +from shared.services.retrieval.nav.knowhere_hybrid import ( + ScoreUnitRow, + score_rows_hybrid_all, + score_unit_stream_hybrid_all, +) + + +@dataclass +class _FakeChunkStore: + units_by_section: dict[str, list[UnitRow]] + + def load_section_units( + self, + document_id: str, + section_id: str, + extra_chunk_ids: Sequence[str] = (), + ) -> list[UnitRow]: + del document_id + units = list(self.units_by_section.get(section_id, ())) + known = {unit.chunk_id for unit in units} + for units_in_section in self.units_by_section.values(): + for unit in units_in_section: + if unit.chunk_id in extra_chunk_ids and unit.chunk_id not in known: + units.append(unit) + return units + + def close(self) -> None: + return None + + +def _providers() -> tuple[ProviderToolSpace, ProviderToolSpace]: + sections = [ + SectionRow("root", None, "Root", "Root", 0, "", 0), + SectionRow("section", "root", "Root / Section", "Section", 1, "", 1), + SectionRow("leaf", "section", "Root / Section / Leaf", "Leaf", 2, "", 2), + ] + text = UnitRow( + "duplicate-chunk", + "leaf", + "text", + "alpha retrieval evidence", + 1, + metadata={"connect_to": [{"target": "asset-1", "relation": "embeds"}]}, + ) + asset = UnitRow( + "asset-1", + "root", + "image", + "", + 2, + file_path="images/asset.png", + metadata={"summary": "supporting image"}, + ) + eager = NamespaceKnowhereProvider( + [KnowhereProvider(doc_id="doc", sections=sections, units=[text, asset])], + titles={"doc": "document"}, + ) + store = _FakeChunkStore({"root": [asset], "leaf": [text]}) + lazy = NamespaceKnowhereProvider( + [ + LazyKnowhereProvider( + doc_id="doc", + sections=sections, + chunk_store=store, + known_chunk_ids=[text.chunk_id, asset.chunk_id], + root_asset_ids=[asset.chunk_id], + remounted_assets_by_section={"leaf": [asset.chunk_id]}, + ) + ], + titles={"doc": "document"}, + chunk_owner_by_id={"duplicate-chunk": "doc", "asset-1": "doc"}, + ) + return ProviderToolSpace(eager), ProviderToolSpace(lazy) + + +def test_lazy_provider_preserves_score_units_and_scores() -> None: + eager, lazy = _providers() + + assert build_score_units(eager, "doc") == build_score_units(lazy, "doc") + assert compute_corpus_map_and_unit_scores( + eager, doc_ids=["doc"], query="alpha retrieval" + ) == compute_corpus_map_and_unit_scores( + lazy, doc_ids=["doc"], query="alpha retrieval" + ) + + lazy_provider = lazy._provider + self_units = getattr(lazy_provider, "self_units") + assert [unit.chunk_id for unit in self_units("leaf")] == [ + "duplicate-chunk", + "asset-1", + ] + + +def test_streaming_scorer_preserves_exact_eager_scores() -> None: + rows: list[ScoreUnitRow] = [ + { + "chunk_id": "unit-a", + "path_search_text": "root alpha", + "content_search_text": "alpha alpha evidence", + "term_search_text": "alpha alpha evidence root", + }, + { + "chunk_id": "unit-b", + "path_search_text": "root beta", + "content_search_text": "beta evidence", + "term_search_text": "beta evidence root", + }, + { + "chunk_id": "unit-c", + "path_search_text": "root common", + "content_search_text": "common evidence", + "term_search_text": "common evidence root", + }, + ] + eager_rows: list[dict[str, Any]] = [dict(row) for row in rows] + eager_scores = { + str(row["chunk_id"]): float(row["score"]) + for row in score_rows_hybrid_all(eager_rows, "alpha evidence") + } + replay_count: int = 0 + + def unit_factory() -> Sequence[ScoreUnitRow]: + nonlocal replay_count + replay_count += 1 + return rows + + assert score_unit_stream_hybrid_all(unit_factory, "alpha evidence") == eager_scores + assert replay_count == 1 + + +def test_streaming_scorer_preserves_duplicate_id_eager_semantics() -> None: + rows: list[ScoreUnitRow] = [ + { + "chunk_id": "duplicate", + "path_search_text": "alpha", + "content_search_text": "alpha", + "term_search_text": "alpha", + }, + { + "chunk_id": "duplicate", + "path_search_text": "beta", + "content_search_text": "beta", + "term_search_text": "beta", + }, + { + "chunk_id": "other", + "path_search_text": "alpha beta", + "content_search_text": "alpha beta", + "term_search_text": "alpha beta", + }, + ] + eager_rows: list[dict[str, Any]] = [dict(row) for row in rows] + eager_scores = { + str(row["chunk_id"]): float(row["score"]) + for row in score_rows_hybrid_all(eager_rows, "alpha beta") + } + + assert score_unit_stream_hybrid_all(lambda: rows, "alpha beta") == eager_scores diff --git a/apps/api/tests/migrations/test_schema_contract.py b/apps/api/tests/migrations/test_schema_contract.py index 24047dcc..3ead1217 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/packages/shared-python/shared/models/database/document.py b/packages/shared-python/shared/models/database/document.py index f9681e17..b2381996 100644 --- a/packages/shared-python/shared/models/database/document.py +++ b/packages/shared-python/shared/models/database/document.py @@ -210,6 +210,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", diff --git a/packages/shared-python/shared/services/retrieval/execution/routes.py b/packages/shared-python/shared/services/retrieval/execution/routes.py index e67a1ead..ad2268ab 100644 --- a/packages/shared-python/shared/services/retrieval/execution/routes.py +++ b/packages/shared-python/shared/services/retrieval/execution/routes.py @@ -186,6 +186,7 @@ async def _run_mapnav_route( namespace=context.namespace, exclude_document_ids=context.exclude_document_ids, exclude_sections=context.exclude_sections, + lazy=True, ) # Small-corpus count / snapshot reads may leave a checkout; drop it before @@ -196,19 +197,22 @@ 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, - ) + 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) + finally: + snapshot.close() async with open_fresh_database_context() as final_db: resolved = await resolve_workflow_references( diff --git a/packages/shared-python/shared/services/retrieval/nav/knowhere_hybrid.py b/packages/shared-python/shared/services/retrieval/nav/knowhere_hybrid.py index d56c72fb..7d8b195b 100644 --- a/packages/shared-python/shared/services/retrieval/nav/knowhere_hybrid.py +++ b/packages/shared-python/shared/services/retrieval/nav/knowhere_hybrid.py @@ -9,7 +9,11 @@ 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 +22,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: @@ -68,7 +85,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] @@ -590,3 +607,188 @@ def _rows_from_scores(score_by_id: Dict[str, float]) -> List[dict[str, Any]]: row["term_channel_score"] = float(term_channel.get(uid, 0.0) or 0.0) out_rows.append(row) return out_rows + + +def score_unit_stream_hybrid_all( + unit_factory: Callable[[], Iterable[ScoreUnitRow]], + query: str, +) -> Dict[str, float]: + """Score replayable units without retaining their payloads. + + This is the corpus scorer used by map-nav. It mirrors the active BM25 and + weighted-RRF implementation, but keeps only token statistics, identifiers, + and final scores between bounded provider reads. + """ + query_tokens = tokenize_query_for_ranker(query) + path_stats = _StreamingBm25Stats.empty() + content_stats = _StreamingBm25Stats.empty() + units: List[_StreamingUnit] = [] + query_token_set = set(query_tokens) + query_lower = query.lower().strip() + for row in unit_factory(): + unit_id = str(row.get("chunk_id") or "").strip() + if not unit_id: + continue + path_tokens = _get_search_tokens(row, search_field="path_search_text") + content_tokens = _get_search_tokens(row, search_field="content_search_text") + path_stats.observe(path_tokens) + content_stats.observe(content_tokens) + path_frequencies = Counter(path_tokens) + content_frequencies = Counter(content_tokens) + term_score = 0.0 + if query_lower: + haystack = str(row.get("term_search_text") or "").lower() + if query_lower in haystack: + term_score = 100.0 + else: + hit_count = sum(1 for token in query_tokens if token in haystack) + if hit_count > 0: + term_score = float(hit_count) + units.append( + _StreamingUnit( + unit_id=unit_id, + path_length=len(path_tokens), + content_length=len(content_tokens), + path_frequencies={ + token: path_frequencies[token] + for token in query_token_set + if path_frequencies[token] + }, + content_frequencies={ + token: content_frequencies[token] + for token in query_token_set + if content_frequencies[token] + }, + term_score=term_score, + ) + ) + path_stats.finalize() + content_stats.finalize() + path_by_id: Dict[str, float] = {} + content_by_id: Dict[str, float] = {} + term_by_id: Dict[str, float] = {} + unit_ids = list(dict.fromkeys(unit.unit_id for unit in units)) + for unit in units: + path_score = path_stats.score( + unit.path_length, unit.path_frequencies, query_tokens + ) + content_score = content_stats.score( + unit.content_length, unit.content_frequencies, query_tokens + ) + path_by_id[unit.unit_id] = path_score + content_by_id[unit.unit_id] = content_score + term_by_id[unit.unit_id] = unit.term_score + + path_rows = [ + (score, unit_id) + for unit_id, score in path_by_id.items() + if score > 0.0 + ] + content_rows = [ + (score, unit_id) + for unit_id, score in content_by_id.items() + if score > 0.0 + ] + term_rows = [ + (score, unit_id) + for unit_id, score in term_by_id.items() + if score > 0.0 + ] + + path_rows.sort(key=lambda item: (-item[0], item[1])) + content_rows.sort(key=lambda item: (-item[0], item[1])) + term_rows.sort(key=lambda item: (-item[0], item[1])) + path_weight, content_weight, term_weight = map_channel_weights() + rrf_k = int( + os.environ.get( + "NAV_MAP_RRF_K", + os.environ.get("NAV_DISCOVERY_RRF_K", str(RRF_K)), + ).strip() + or RRF_K + ) + fused: Dict[str, float] = {unit_id: 0.0 for unit_id in unit_ids} + for rank, (_score, unit_id) in enumerate(path_rows): + fused[unit_id] = fused.get(unit_id, 0.0) + path_weight / (rrf_k + rank + 1) + for rank, (_score, unit_id) in enumerate(content_rows): + fused[unit_id] = fused.get(unit_id, 0.0) + content_weight / (rrf_k + rank + 1) + for rank, (_score, unit_id) in enumerate(term_rows): + fused[unit_id] = fused.get(unit_id, 0.0) + term_weight / (rrf_k + rank + 1) + return {unit_id: round(score, 6) for unit_id, score in fused.items()} + + +@dataclass(frozen=True) +class _StreamingUnit: + unit_id: str + path_length: int + content_length: int + path_frequencies: Mapping[str, int] + content_frequencies: Mapping[str, int] + term_score: float + + +class _StreamingBm25Stats: + """Exact BM25Okapi corpus statistics collected without row retention.""" + + def __init__(self) -> None: + self.document_count: int = 0 + self.document_frequency: Counter[str] = Counter() + self.total_length: int = 0 + self.average_length: float = 0.0 + self.idf_by_token: Dict[str, float] = {} + + @classmethod + def empty(cls) -> "_StreamingBm25Stats": + return cls() + + def observe(self, tokens: List[str]) -> None: + if not tokens: + return + self.document_count += 1 + self.total_length += len(tokens) + self.document_frequency.update(set(tokens)) + + def finalize(self) -> None: + self.average_length = ( + self.total_length / self.document_count if self.document_count else 0.0 + ) + idf_by_token: Dict[str, float] = {} + idf_sum = 0.0 + negative_tokens: List[str] = [] + for token, frequency in self.document_frequency.items(): + idf = math.log(self.document_count - frequency + 0.5) - math.log( + frequency + 0.5 + ) + idf_by_token[token] = idf + idf_sum += idf + if idf < 0.0: + negative_tokens.append(token) + average_idf = idf_sum / len(idf_by_token) if idf_by_token else 0.0 + epsilon_floor = 0.25 * average_idf + for token in negative_tokens: + idf_by_token[token] = epsilon_floor + self.idf_by_token = idf_by_token + + def score( + self, + document_length: int, + frequencies: Dict[str, int], + query_tokens: List[str], + ) -> float: + if ( + not frequencies + or not query_tokens + or not self.document_count + or self.average_length <= 0.0 + ): + return 0.0 + denominator_base = 1.5 * ( + 1.0 - 0.75 + 0.75 * document_length / self.average_length + ) + score = 0.0 + for token in query_tokens: + frequency = frequencies.get(token, 0) + if not frequency: + continue + idf = self.idf_by_token.get(token, 0.0) + score += idf * (frequency * 2.5 / (frequency + denominator_base)) + return score diff --git a/packages/shared-python/shared/services/retrieval/nav/nav_hierarchy.py b/packages/shared-python/shared/services/retrieval/nav/nav_hierarchy.py index fafa66d2..93b35598 100644 --- a/packages/shared-python/shared/services/retrieval/nav/nav_hierarchy.py +++ b/packages/shared-python/shared/services/retrieval/nav/nav_hierarchy.py @@ -231,6 +231,12 @@ def read_chunks(self, section_id: str, query: str, *, doc_id: str, k: int) -> Li del section_id, query, doc_id, k return [] + def release_section_units(self, section_id: str) -> None: + """Release one lazy section without discarding the hierarchy.""" + release = getattr(self._provider, "release_section_units", None) + if callable(release): + release(section_id) + @dataclass class InMemoryNode: diff --git a/packages/shared-python/shared/services/retrieval/nav/nav_knowhere.py b/packages/shared-python/shared/services/retrieval/nav/nav_knowhere.py index 19caf893..d73eb178 100644 --- a/packages/shared-python/shared/services/retrieval/nav/nav_knowhere.py +++ b/packages/shared-python/shared/services/retrieval/nav/nav_knowhere.py @@ -22,7 +22,7 @@ import os from dataclasses import dataclass, field -from typing import Any, Dict, List, Optional, Sequence, Set, Tuple +from typing import Any, Callable, Dict, Iterable, List, Optional, Protocol, Sequence, Set, Tuple from .nav_address import NavLevel from .nav_hierarchy import NodeMeta @@ -127,7 +127,125 @@ 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: + return configured.replace("postgresql+asyncpg", "postgresql+psycopg2") + return _DEFAULT_DSN + + +class ChunkStore(Protocol): + def load_section_units( + self, + document_id: str, + section_id: str, + extra_chunk_ids: Sequence[str] = (), + ) -> List[UnitRow]: + raise NotImplementedError + + def close(self) -> None: + raise NotImplementedError + + +class ReadOnlyChunkStore: + """Episode-local, revision-pinned loader for lazy map-nav chunks.""" + + def __init__( + self, + *, + dsn: str, + revisions: Dict[str, str], + excluded_sections: Optional[Iterable[Tuple[str, str]]] = None, + ) -> None: + self._dsn = str(dsn) + self._revisions = dict(revisions) + self._excluded_sections = set(excluded_sections or ()) + self._conn: Optional[_SyncConnection] = None + + def _connection(self) -> "_SyncConnection": + if self._conn is None: + self._conn = _connect(self._dsn) + self._conn.set_session(readonly=True, autocommit=True) + return self._conn + + def load_section_units( + self, + document_id: str, + section_id: str, + extra_chunk_ids: Sequence[str] = (), + ) -> List[UnitRow]: + doc_id = str(document_id).strip() + sid = str(section_id).strip() + job_result_id = self._revisions.get(doc_id) + if not doc_id or not sid or not job_result_id or (doc_id, sid) in self._excluded_sections: + return [] + cur = self._connection().cursor() + try: + ids = [str(chunk_id).strip() for chunk_id in extra_chunk_ids if str(chunk_id).strip()] + if ids: + cur.execute( + "SELECT chunk_id, section_id, chunk_type, content, sort_order, " + "source_chunk_path, file_path, chunk_metadata " + "FROM document_chunks " + "WHERE document_id = %s AND job_result_id = %s " + "AND (section_id = %s OR chunk_id = ANY(%s)) " + "ORDER BY sort_order, chunk_id, id", + (doc_id, job_result_id, sid, ids), + ) + else: + cur.execute( + "SELECT chunk_id, section_id, chunk_type, content, sort_order, " + "source_chunk_path, file_path, chunk_metadata " + "FROM document_chunks " + "WHERE document_id = %s AND job_result_id = %s AND section_id = %s " + "ORDER BY sort_order, chunk_id, id", + (doc_id, job_result_id, sid), + ) + return [_unit_from_row(row) for row in cur.fetchall()] + finally: + cur.close() + + def close(self) -> None: + if self._conn is not None: + self._conn.close() + self._conn = None + + +class _SyncCursor(Protocol): + def execute(self, query: str, params: Sequence[object]) -> None: + raise NotImplementedError + + def fetchall(self) -> Sequence[Sequence[object]]: + raise NotImplementedError + + def close(self) -> None: + raise NotImplementedError + + +class _SyncConnection(Protocol): + def set_session(self, *, readonly: bool, autocommit: bool) -> None: + raise NotImplementedError + + def cursor(self) -> _SyncCursor: + raise NotImplementedError + + def close(self) -> None: + raise NotImplementedError + + +def _unit_from_row(row: Sequence[object]) -> UnitRow: + return UnitRow( + chunk_id=str(row[0] or ""), + section_id=str(row[1]) if row[1] else None, + chunk_type=str(row[2] or "text"), + content=str(row[3] or ""), + sort_order=int(row[4] or 0), + source_chunk_path=str(row[5] or ""), + file_path=str(row[6] or ""), + metadata=_as_meta(row[7]), + ) class KnowhereProvider: @@ -139,8 +257,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 +279,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 +294,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 +430,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 +504,59 @@ def summaries(self) -> Dict[str, str]: def all_section_ids(self) -> List[str]: return list(self._sections) + def chunk_count(self) -> int: + return len(self._chunk_ids) + + +class LazyKnowhereProvider(KnowhereProvider): + """Hierarchy provider that loads full chunk rows only on first access.""" + + def __init__( + self, + *, + doc_id: str, + sections: Sequence[SectionRow], + chunk_store: ChunkStore, + known_chunk_ids: Sequence[str], + root_asset_ids: Sequence[str] = (), + remounted_assets_by_section: Optional[Dict[str, Sequence[str]]] = None, + ) -> None: + super().__init__( + doc_id=doc_id, + sections=sections, + units=(), + lazy_loader=lambda section_id: chunk_store.load_section_units( + doc_id, + section_id, + (remounted_assets_by_section or {}).get(section_id, ()), + ), + known_chunk_ids=known_chunk_ids, + ) + self._chunk_store = chunk_store + self._root_asset_ids = {str(chunk_id) for chunk_id in root_asset_ids} + + def _ensure_section_loaded(self, section_id: str) -> None: + super()._ensure_section_loaded(section_id) + if ( + section_id in self._units_by_section + and self._root_asset_ids + and is_root_section_path(self.section_path(section_id)) + ): + self._units_by_section[section_id] = [ + unit + for unit in self._units_by_section[section_id] + if unit.chunk_id not in self._root_asset_ids + ] + + def close(self) -> None: + self._chunk_store.close() -def _connect(dsn: str): + 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) @@ -499,6 +704,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 +720,44 @@ 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 address_level(self, node_id: str) -> Optional[NavLevel]: sid = str(node_id or "").strip() if not sid: @@ -562,7 +798,8 @@ 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( + 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( 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 d1049adf..b88df1fa 100644 --- a/packages/shared-python/shared/services/retrieval/nav/nav_map_scores.py +++ b/packages/shared-python/shared/services/retrieval/nav/nav_map_scores.py @@ -1,13 +1,15 @@ from __future__ import annotations +from collections.abc import Iterator from typing import Any, Dict, List, Optional, Sequence, Set, Tuple from .knowhere_hybrid import ( + ScoreUnitRow, build_content_search_text, build_path_search_text, build_term_search_text, - score_dense_channel, score_rows_hybrid_all, + score_unit_stream_hybrid_all, ) @@ -26,6 +28,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) @@ -173,47 +179,6 @@ def score_node(section_id: str) -> float: return map_scores -def _score_dense_units_by_doc( - units_by_doc: Sequence[Tuple[str, List[dict]]], - query: str, - *, - namespace: Optional[str], -) -> Dict[str, Optional[Dict[str, float]]]: - """Read per-doc vector caches, returning raw cosine scores for global fusion. - - Dense cosine is independently comparable across documents. Partitioning only - preserves the existing per-doc disk cache; no ranking or normalization occurs - here. If any partition fails, that whole channel falls back to global BM25. - """ - dense_by_channel: Dict[str, Optional[Dict[str, float]]] = {} - for channel, text_field in (("path", "path_text"), ("content", "content")): - score_by_id: Dict[str, float] = {} - complete = True - for doc_id, units in units_by_doc: - if not units: - continue - unit_ids = [str(unit["chunk_id"]) for unit in units] - scores = score_dense_channel( - [str(unit.get(text_field) or "") for unit in units], - query, - unit_ids=unit_ids, - doc_id=doc_id, - channel=channel, - namespace=namespace, - ) - if scores is None or len(scores) != len(unit_ids): - complete = False - break - score_by_id.update( - { - unit_id: float(scores[index]) - for index, unit_id in enumerate(unit_ids) - } - ) - dense_by_channel[channel] = score_by_id if complete else None - return dense_by_channel - - def build_score_units(ts: Any, doc_id: str, root_ids: Optional[Sequence[str]] = None) -> List[dict]: """Build leaf (+ interstitial self_only) units for hybrid scoring.""" if root_ids is None: @@ -276,6 +241,66 @@ def build_score_units(ts: Any, doc_id: str, root_ids: Optional[Sequence[str]] = return units +def iter_score_units( + ts: Any, + doc_id: str, + *, + children_map: Dict[str, List[str]], + leaves: Set[str], + titles: Dict[str, str], +) -> Iterator[ScoreUnitRow]: + """Yield scoring units while releasing each lazy section after use.""" + release = getattr(ts, "release_section_units", None) + + for leaf_id in sorted(leaves): + try: + content = _section_body_text(ts, leaf_id, doc_id) or ( + titles.get(leaf_id) or _line_content(ts, leaf_id, doc_id) + ) + if content: + path_text = _ancestor_path_titles(ts, leaf_id, doc_id) + title = titles.get(leaf_id) or _line_content(ts, leaf_id, doc_id) + yield { + "chunk_id": leaf_id, + "section_id": leaf_id, + "kind": "leaf", + "content": content, + "path_text": path_text, + "path_search_text": build_path_search_text( + section_path=path_text, section_title=title or content + ), + "content_search_text": build_content_search_text(content), + "term_search_text": build_term_search_text(content, path_text=path_text), + } + finally: + if callable(release): + release(leaf_id) + + for sid, kids in children_map.items(): + if not kids: + continue + try: + self_text, has_interstitial = _self_only_text(ts, sid, doc_id) + if not has_interstitial or not self_text: + continue + path_text = _ancestor_path_titles(ts, sid, doc_id) + yield { + "chunk_id": f"{sid}__self", + "section_id": sid, + "kind": "self_only", + "content": self_text, + "path_text": path_text, + "path_search_text": build_path_search_text( + section_path=path_text, section_title=titles.get(sid) or "" + ), + "content_search_text": build_content_search_text(self_text), + "term_search_text": build_term_search_text(self_text, path_text=path_text), + } + finally: + if callable(release): + release(sid) + + def compute_map_scores( ts: Any, *, @@ -346,41 +371,33 @@ 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, - ) - scored = score_rows_hybrid_all( - all_units, - query, - dense_scores_by_channel=dense_scores, - ) - unit_scores = { - str(row.get("chunk_id") or ""): float(row.get("score") or 0.0) - for row in scored - } + children_map, leaves, titles = _walk_tree(ts, doc_id, root_ids) + tree_by_doc[doc_id] = (children_map, leaves, titles) + + def unit_factory() -> Iterator[ScoreUnitRow]: + for document_id in valid_doc_ids: + children_map, leaves, titles = tree_by_doc[document_id] + yield from iter_score_units( + ts, + document_id, + children_map=children_map, + leaves=leaves, + titles=titles, + ) + + unit_scores = score_unit_stream_hybrid_all(unit_factory, query) map_scores: Dict[str, float] = {} for doc_id in valid_doc_ids: - children_map, leaves = tree_by_doc[doc_id] + children_map, leaves, _titles = tree_by_doc[doc_id] doc_map_scores = _pool_unit_scores_to_tree( children_map, leaves, unit_scores ) diff --git a/packages/shared-python/shared/services/retrieval/nav_snapshot.py b/packages/shared-python/shared/services/retrieval/nav_snapshot.py index 7dc60c6f..1edfd374 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, *,