From 2b8c9367c8528d48032ed89a31c04498ca803f66 Mon Sep 17 00:00:00 2001 From: suguanYang Date: Fri, 28 Aug 2026 23:58:58 +0800 Subject: [PATCH 01/19] perf: instrument mapnav episode timings --- .../services/retrieval/nav/nav_agent.py | 27 ++++++- .../shared/services/retrieval/nav/nav_llm.py | 74 ++++++++++++------- .../services/retrieval/nav/nav_orchestrate.py | 37 ++++++++-- 3 files changed, 104 insertions(+), 34 deletions(-) diff --git a/packages/shared-python/shared/services/retrieval/nav/nav_agent.py b/packages/shared-python/shared/services/retrieval/nav/nav_agent.py index f4d0d12d..197af5af 100644 --- a/packages/shared-python/shared/services/retrieval/nav/nav_agent.py +++ b/packages/shared-python/shared/services/retrieval/nav/nav_agent.py @@ -1,7 +1,8 @@ from __future__ import annotations -import time +import logging import os +import time from typing import Any, List, Optional, Sequence, Tuple from ._compat import AgentStep, EpisodeResult @@ -37,6 +38,7 @@ # Back-compat aliases for tests / callers. _evidence_owner_section_id = evidence_owner_section_id _unit_score_for_evidence_chunk = unit_score_for_evidence_chunk +_logger = logging.getLogger(__name__) def _chunks_to_retrieved_nodes(chunks: List[Chunk]) -> List[str]: @@ -398,6 +400,7 @@ def _run_nav_episode_body( state = NavState(doc_id=episode_doc, query=query, task_type=task_type) steps: List[AgentStep] = [] + map_started = time.perf_counter() if namespace_mode: section_ids = list(ts.sections_for_doc("")) state.map_scores, state.unit_scores = compute_corpus_map_and_unit_scores( @@ -408,6 +411,12 @@ def _run_nav_episode_body( state.map_scores, state.unit_scores = compute_map_and_unit_scores( ts, doc_id=episode_doc, query=query, root_ids=section_ids ) + _logger.info( + "retrieval mapnav phase=map_scoring seconds=%.3f documents=%d sections=%d", + time.perf_counter() - map_started, + len(corpus_ids), + len(section_ids), + ) state.highlight_ids = select_map_highlights( state.unit_scores, k=int(cfg.collect_top_k) ) @@ -434,6 +443,11 @@ def _run_nav_episode_body( }, t0=plan_t0), ) ) + _logger.info( + "retrieval mapnav phase=planner seconds=%.3f subgoals=%d", + time.perf_counter() - plan_t0, + len(retrieval_plan.subgoals), + ) # Checklist: wave orchestration; navigate mode: classic single navigate. if cfg.is_checklist and state.retrieval_plan is not None: @@ -452,6 +466,11 @@ def _run_nav_episode_body( detail=stamp_step_detail(orch_detail, t0=orch_t0), ) ) + _logger.info( + "retrieval mapnav phase=orchestration seconds=%.3f waves=%d", + time.perf_counter() - orch_t0, + len(orch_detail.get("waves", [])), + ) else: navigate( ts, @@ -464,6 +483,7 @@ def _run_nav_episode_body( steps_out=steps, ) + evidence_started = time.perf_counter() fill = pack_nav_evidence( _dedupe_scored(list(state.collected)), ts, @@ -471,6 +491,11 @@ def _run_nav_episode_body( cfg, budget_chars=budget_chars, ) + _logger.info( + "retrieval mapnav phase=evidence_pack seconds=%.3f chunks=%d", + time.perf_counter() - evidence_started, + len(fill.kept_chunks), + ) scored_chunks = list(fill.scored_chunks) retrieval_seconds = time.perf_counter() - retrieval_t0 composed = "" diff --git a/packages/shared-python/shared/services/retrieval/nav/nav_llm.py b/packages/shared-python/shared/services/retrieval/nav/nav_llm.py index f8ba3187..767c1675 100644 --- a/packages/shared-python/shared/services/retrieval/nav/nav_llm.py +++ b/packages/shared-python/shared/services/retrieval/nav/nav_llm.py @@ -20,7 +20,9 @@ from __future__ import annotations +import logging import os +import time from contextlib import contextmanager from contextvars import ContextVar from typing import Any, Callable, Dict, Iterator, Optional, Sequence @@ -30,6 +32,7 @@ NavChatBackend = Callable[..., Dict[str, Any]] _backend: Optional[NavChatBackend] = None +_logger = logging.getLogger(__name__) _DS_DEFAULT_MODEL = "deepseek-v4-flash" _DEFAULT_PLANNER_THINK_MAX = 16384 @@ -190,25 +193,34 @@ def nav_chat( raise NavTokenLimit(used=nav_tokens_used(), limit=nav_token_limit()) merged_extra = _merge_thinking_extra(extra, role=thinking_role, model=model) + call_started = time.perf_counter() if _backend is not None: - result = _backend( - purpose=purpose, - messages=list(messages), - model=model, - temperature=temperature, - max_tokens=max_tokens, - response_format=response_format, - extra=merged_extra, - thinking_role=thinking_role, - context=context, - api_key_env=api_key_env, - base_url_env=base_url_env, - timeout=timeout, - usage_tag=usage_tag, - ) - record_episode_tokens((result or {}).get("usage")) - return result + try: + result = _backend( + purpose=purpose, + messages=list(messages), + model=model, + temperature=temperature, + max_tokens=max_tokens, + response_format=response_format, + extra=merged_extra, + thinking_role=thinking_role, + context=context, + api_key_env=api_key_env, + base_url_env=base_url_env, + timeout=timeout, + usage_tag=usage_tag, + ) + record_episode_tokens((result or {}).get("usage")) + return result + finally: + _logger.info( + "retrieval mapnav llm_call purpose=%s role=%s seconds=%.3f", + purpose, + thinking_role, + time.perf_counter() - call_started, + ) from ._compat import cached_chat_completion # type: ignore from ._compat import ( # type: ignore @@ -230,16 +242,24 @@ def nav_chat( f"(model={model!r}; need DS_KEY for deepseek-* or OPENAI_API_KEY)." ) client = make_openai_client(api_key=key, base_url=base_url, timeout=timeout) - cached = cached_chat_completion( - client, - purpose=purpose, - model=model, - messages=list(messages), - temperature=temperature, - max_tokens=max_tokens, - response_format=response_format, - extra=merged_extra, - ) + try: + cached = cached_chat_completion( + client, + purpose=purpose, + model=model, + messages=list(messages), + temperature=temperature, + max_tokens=max_tokens, + response_format=response_format, + extra=merged_extra, + ) + finally: + _logger.info( + "retrieval mapnav llm_call purpose=%s role=%s seconds=%.3f", + purpose, + thinking_role, + time.perf_counter() - call_started, + ) if usage_tag: record_usage(usage_tag, cached.get("usage")) record_episode_tokens(cached.get("usage")) diff --git a/packages/shared-python/shared/services/retrieval/nav/nav_orchestrate.py b/packages/shared-python/shared/services/retrieval/nav/nav_orchestrate.py index 91f45e1c..2e82be08 100644 --- a/packages/shared-python/shared/services/retrieval/nav/nav_orchestrate.py +++ b/packages/shared-python/shared/services/retrieval/nav/nav_orchestrate.py @@ -8,14 +8,14 @@ from __future__ import annotations -from .nav_token_budget import stamp_step_detail - +import logging import re import time from contextlib import contextmanager from dataclasses import asdict from typing import Any, Dict, Iterator, List, Optional, Sequence, Set, Tuple +from .nav_token_budget import stamp_step_detail from .nav_navigate import navigate from .nav_plan import ( RetrievalPlan, @@ -29,6 +29,7 @@ from .nav_verify import apply_bindings_from_result, build_subgoal_result _SLOT_STRIP_RE = re.compile(r"\{\{\s*[^}]+\s*\}\}") +_logger = logging.getLogger(__name__) def ready_subgoal_ids( @@ -483,8 +484,18 @@ def _run_one(sid: str, working_state: NavState, out_steps: Optional[List[Any]]) ) # Serial wave execution (parallel fan-out retired with ThreadPoolExecutor). + wave_started = time.perf_counter() for sid in ready: - outputs.append(_run_one(sid, state, steps_out)) + subgoal_started = time.perf_counter() + try: + outputs.append(_run_one(sid, state, steps_out)) + finally: + _logger.info( + "retrieval mapnav harvest subgoal=%s wave=%d seconds=%.3f", + sid, + wave_idx, + time.perf_counter() - subgoal_started, + ) # Bookkeeping shared by both decision paths. for item in outputs: @@ -508,9 +519,17 @@ def _run_one(sid: str, working_state: NavState, out_steps: Optional[List[Any]]) state.subgoal_attempt_counts.get(sid, 0) ) + 1 - control_detail = _apply_plan_control( - ts, state, config, plan=plan, outputs=outputs, by_id=by_id, steps_out=steps_out - ) + control_started = time.perf_counter() + try: + control_detail = _apply_plan_control( + ts, state, config, plan=plan, outputs=outputs, by_id=by_id, steps_out=steps_out + ) + finally: + _logger.info( + "retrieval mapnav plan_control wave=%d seconds=%.3f", + wave_idx, + time.perf_counter() - control_started, + ) wave_detail["plan_control"] = control_detail replan_requested = bool(control_detail.get("replan")) if control_detail.get("done"): @@ -525,6 +544,12 @@ def _run_one(sid: str, working_state: NavState, out_steps: Optional[List[Any]]) ) ) summary["waves"].append(wave_detail) + _logger.info( + "retrieval mapnav wave=%d ready=%d seconds=%.3f", + wave_idx, + len(ready), + time.perf_counter() - wave_started, + ) if replan_requested: cap = int(getattr(config, "max_replans", 0) or 0) From 259065fec195db91d6abf6b572e479d2f26aafe5 Mon Sep 17 00:00:00 2001 From: suguanYang Date: Sat, 29 Aug 2026 14:41:02 +0800 Subject: [PATCH 02/19] perf: eliminate lazy mapnav tree n-plus-one loads --- ...etrieval_lazy_snapshot_quality_contract.py | 90 ++++++++++++++++++- .../test_retrieval_lazy_tree_contract.py | 33 +++++++ .../services/retrieval/nav/nav_hierarchy.py | 18 ++-- .../services/retrieval/nav/nav_knowhere.py | 54 ++++++++++- .../services/retrieval/nav/nav_map_scores.py | 31 +++++++ .../shared/services/retrieval/nav_snapshot.py | 34 ++++++- 6 files changed, 250 insertions(+), 10 deletions(-) create mode 100644 apps/api/tests/contract/test_retrieval_lazy_tree_contract.py diff --git a/apps/api/tests/contract/test_retrieval_lazy_snapshot_quality_contract.py b/apps/api/tests/contract/test_retrieval_lazy_snapshot_quality_contract.py index ae6d78ed..0752a3b9 100644 --- a/apps/api/tests/contract/test_retrieval_lazy_snapshot_quality_contract.py +++ b/apps/api/tests/contract/test_retrieval_lazy_snapshot_quality_contract.py @@ -2,6 +2,8 @@ from dataclasses import dataclass from collections.abc import Mapping, Sequence +from collections import Counter +import math from typing import Any from shared.services.retrieval.nav.nav_hierarchy import ProviderToolSpace @@ -20,7 +22,11 @@ ) from shared.services.retrieval.nav.knowhere_hybrid import ( ScoreUnitRow, + PersistedBm25Stats, + PersistedScoreCorpus, + PersistedScoreUnit, score_rows_hybrid_all, + score_persisted_corpus_many, score_unit_stream_hybrid_all, score_unit_stream_hybrid_many, ) @@ -304,8 +310,7 @@ def test_streaming_scorer_scores_multiple_queries_with_one_corpus_read() -> None ] queries: list[str] = ["alpha evidence", "beta evidence"] expected = { - query: score_unit_stream_hybrid_all(lambda: rows, query) - for query in queries + query: score_unit_stream_hybrid_all(lambda: rows, query) for query in queries } read_count: int = 0 @@ -318,6 +323,87 @@ def unit_factory() -> Sequence[ScoreUnitRow]: assert read_count == 1 +def test_persisted_score_projection_preserves_exact_eager_scores() -> None: + rows: list[ScoreUnitRow] = [ + { + "chunk_id": "unit-a", + "path_search_text": "root alpha", + "content_search_text": "common alpha alpha evidence", + "term_search_text": "common alpha alpha evidence root", + }, + { + "chunk_id": "unit-b", + "path_search_text": "root beta", + "content_search_text": "common beta evidence", + "term_search_text": "common beta evidence root", + }, + { + "chunk_id": "unit-c", + "path_search_text": "root common", + "content_search_text": "common evidence", + "term_search_text": "common evidence root", + }, + ] + queries = ["common alpha", "beta evidence"] + expected = { + query: score_unit_stream_hybrid_all(lambda: rows, query) for query in queries + } + query_tokens = {token for query in queries for token in query.split()} + + def build_stats(search_field: str) -> PersistedBm25Stats: + token_rows = [str(row[search_field]).split() for row in rows] + document_frequency = Counter( + token for tokens in token_rows for token in set(tokens) + ) + document_count = len(token_rows) + raw_idfs = [ + math.log(document_count - frequency + 0.5) - math.log(frequency + 0.5) + for frequency in document_frequency.values() + ] + return PersistedBm25Stats( + document_count=document_count, + total_length=sum(len(tokens) for tokens in token_rows), + document_frequency={ + token: document_frequency[token] for token in query_tokens + }, + average_idf=sum(raw_idfs) / len(raw_idfs), + ) + + corpus = PersistedScoreCorpus( + units=[ + PersistedScoreUnit( + unit_id=str(row["chunk_id"]), + path_length=len(str(row["path_search_text"]).split()), + content_length=len(str(row["content_search_text"]).split()), + path_frequencies={ + token: str(row["path_search_text"]).split().count(token) + for token in query_tokens + }, + content_frequencies={ + token: str(row["content_search_text"]).split().count(token) + for token in query_tokens + }, + term_scores=tuple( + 100.0 + if query in str(row["term_search_text"]) + else float( + sum( + token in str(row["term_search_text"]) + for token in query.split() + ) + ) + for query in queries + ), + ) + for row in rows + ], + path_stats=build_stats("path_search_text"), + content_stats=build_stats("content_search_text"), + ) + + assert score_persisted_corpus_many(corpus, queries) == expected + + def test_corpus_map_scores_multiple_queries_with_one_lazy_load() -> None: eager, lazy, store = _providers() queries: list[str] = ["alpha retrieval", "supporting image"] diff --git a/apps/api/tests/contract/test_retrieval_lazy_tree_contract.py b/apps/api/tests/contract/test_retrieval_lazy_tree_contract.py new file mode 100644 index 00000000..30994dad --- /dev/null +++ b/apps/api/tests/contract/test_retrieval_lazy_tree_contract.py @@ -0,0 +1,33 @@ +"""Contract coverage for lazy MAP-NAV tree traversal.""" + +from __future__ import annotations + +from shared.services.retrieval.nav.nav_hierarchy import NodeMeta, ProviderToolSpace +from shared.services.retrieval.nav.nav_knowhere import KnowhereProvider, SectionRow +from shared.services.retrieval.nav.nav_map_scores import _walk_tree + + +class _MetadataForbiddenProvider(KnowhereProvider): + def node_meta(self, section_id: str) -> NodeMeta: + raise AssertionError(f"tree traversal materialized metadata for {section_id}") + + +def test_tree_walk_reads_children_and_titles_without_materializing_metadata() -> None: + provider = _MetadataForbiddenProvider( + doc_id="doc", + sections=[ + SectionRow("root", None, "Root", "Root", 0, "", 0), + SectionRow("child", "root", "Root / Child", "Child", 1, "", 1), + ], + units=(), + ) + + children, leaves, titles = _walk_tree( + ProviderToolSpace(provider), + "doc", + ["root"], + ) + + assert children == {"root": ["child"], "child": []} + assert leaves == {"child"} + assert titles == {"root": "Root", "child": "Child"} diff --git a/packages/shared-python/shared/services/retrieval/nav/nav_hierarchy.py b/packages/shared-python/shared/services/retrieval/nav/nav_hierarchy.py index d9e42b07..e98a9d4c 100644 --- a/packages/shared-python/shared/services/retrieval/nav/nav_hierarchy.py +++ b/packages/shared-python/shared/services/retrieval/nav/nav_hierarchy.py @@ -145,14 +145,22 @@ def get_structure(self, section_id: str) -> dict: def _children_for_section_path( self, section_id: str, doc_id: str, limit: Optional[int] = None ) -> List[dict]: - del doc_id # section_id is globally addressable in this provider model. child_ids = [str(c) for c in self._provider.children(section_id)] if limit is not None: child_ids = child_ids[: max(0, int(limit))] - return [ - {"section_id": cid, "preview": self._provider.node_meta(cid).title} - for cid in child_ids - ] + # ``node_meta`` may materialize a lazy subtree to calculate chunk + # counts. Tree traversal needs only the child id/title; avoid an N+1 + # payload load while building the scoring tree. + out: List[dict] = [] + for cid in child_ids: + path = self.path_titles(cid, doc_id) + out.append( + { + "section_id": cid, + "preview": path.rsplit(" / ", 1)[-1] if path else "", + } + ) + return out def section_relation_ids( self, section_id: str, doc_id: str diff --git a/packages/shared-python/shared/services/retrieval/nav/nav_knowhere.py b/packages/shared-python/shared/services/retrieval/nav/nav_knowhere.py index cf1de779..c056a2d2 100644 --- a/packages/shared-python/shared/services/retrieval/nav/nav_knowhere.py +++ b/packages/shared-python/shared/services/retrieval/nav/nav_knowhere.py @@ -21,6 +21,8 @@ from __future__ import annotations import os +import logging +import time from hashlib import sha256 from dataclasses import dataclass, field from typing import ( @@ -51,6 +53,7 @@ ROOT_SECTION_PATH = "Root" _DEFAULT_DSN = "postgresql://root:root123@127.0.0.1:5433/Knowhere" _MAP_UNIT_INDEX_FORMAT_VERSION = 1 +_logger = logging.getLogger(__name__) @dataclass(frozen=True) @@ -282,6 +285,7 @@ def load_persisted_score_corpus( ] cur = self._connection().cursor() try: + stage_started = time.perf_counter() cur.execute( "SELECT indexes.document_id, indexes.job_result_id, " "indexes.format_version, indexes.unit_count, indexes.token_count " @@ -292,11 +296,17 @@ def load_persisted_score_corpus( revision_params, ) manifests = list(cur.fetchall()) + _logger.info( + "retrieval map-index load stage=manifests seconds=%.3f rows=%d", + time.perf_counter() - stage_started, + len(manifests), + ) if len(manifests) != len(revisions) or any( int(row[2]) != _MAP_UNIT_INDEX_FORMAT_VERSION for row in manifests ): return None + stage_started = time.perf_counter() cur.execute( "SELECT COUNT(*), COUNT(DISTINCT (units.document_id, units.unit_id)) " "FROM document_map_units AS units " @@ -308,9 +318,15 @@ def load_persisted_score_corpus( 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 + _logger.info( + "retrieval map-index load stage=unit_count seconds=%.3f rows=%d", + time.perf_counter() - stage_started, + indexed_unit_count, + ) 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 + stage_started = time.perf_counter() cur.execute( "SELECT COUNT(*) FROM document_map_unit_tokens AS tokens " "JOIN document_map_units AS units ON units.id = tokens.map_unit_id " @@ -321,6 +337,11 @@ def load_persisted_score_corpus( ) token_row = cur.fetchone() indexed_token_count = int(token_row[0]) if token_row else 0 + _logger.info( + "retrieval map-index load stage=token_count seconds=%.3f rows=%d", + time.perf_counter() - stage_started, + indexed_token_count, + ) expected_token_count = sum(int(row[4]) for row in manifests) if indexed_token_count != expected_token_count: return None @@ -340,6 +361,7 @@ def load_persisted_score_corpus( if allowed_pairs: allowed_document_ids = [pair[0] for pair in allowed_pairs] allowed_section_ids = [pair[1] for pair in allowed_pairs] + stage_started = time.perf_counter() cur.execute( "SELECT units.id, units.document_id, units.unit_id, units.section_id, " "units.path_token_count, units.content_token_count " @@ -355,6 +377,11 @@ def load_persisted_score_corpus( [*revision_params, allowed_document_ids, allowed_section_ids], ) unit_rows = list(cur.fetchall()) + _logger.info( + "retrieval map-index load stage=units seconds=%.3f rows=%d", + time.perf_counter() - stage_started, + len(unit_rows), + ) 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 = { @@ -372,6 +399,7 @@ def load_persisted_score_corpus( query_token_hashes = [ sha256(token.encode("utf-8")).hexdigest() for token in query_tokens ] + stage_started = time.perf_counter() cur.execute( "SELECT map_unit_id, channel, token, frequency " "FROM document_map_unit_tokens " @@ -383,6 +411,11 @@ def load_persisted_score_corpus( frequencies.setdefault((str(map_unit_id), str(channel)), {})[ str(token) ] = int(frequency) + _logger.info( + "retrieval map-index load stage=frequencies seconds=%.3f units=%d", + time.perf_counter() - stage_started, + len(map_unit_ids), + ) term_scores = self._load_term_scores( cur, @@ -390,6 +423,11 @@ def load_persisted_score_corpus( queries=unique_queries, query_tokens_by_query=query_tokens_by_query, ) + _logger.info( + "retrieval map-index load stage=complete units=%d queries=%d", + len(unit_rows), + len(unique_queries), + ) path_stats = self._load_persisted_bm25_stats( cur, unit_rows=unit_rows, @@ -459,14 +497,22 @@ def _load_term_scores( params.append(query_lower) params.extend(query_tokens_by_query[query]) params.append(list(map_unit_ids)) + stage_started = time.perf_counter() cur.execute( "SELECT id, " + ", ".join(expressions) + " " "FROM document_map_units WHERE id = ANY(%s)", params, ) + rows = cur.fetchall() + _logger.info( + "retrieval map-index load stage=term_scores units=%d queries=%d seconds=%.3f", + len(map_unit_ids), + len(queries), + time.perf_counter() - stage_started, + ) return { str(row[0]): tuple(float(value) for value in row[1:]) - for row in cur.fetchall() + for row in rows } def _load_persisted_bm25_stats( @@ -497,6 +543,7 @@ def _load_persisted_bm25_stats( ) average_idf = 0.0 if needs_average_idf and map_unit_ids and document_count: + stage_started = time.perf_counter() cur.execute( "SELECT COALESCE(AVG(LN((%s - frequencies.document_frequency + 0.5) " "/ (frequencies.document_frequency + 0.5))), 0.0) " @@ -508,6 +555,11 @@ def _load_persisted_bm25_stats( ) row = cur.fetchone() average_idf = float(row[0]) if row else 0.0 + _logger.info( + "retrieval map-index load stage=average_idf channel=%s seconds=%.3f", + channel, + time.perf_counter() - stage_started, + ) return PersistedBm25Stats( document_count=document_count, total_length=sum(lengths), diff --git a/packages/shared-python/shared/services/retrieval/nav/nav_map_scores.py b/packages/shared-python/shared/services/retrieval/nav/nav_map_scores.py index 6318f490..18d793e9 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,6 +1,8 @@ from __future__ import annotations from collections.abc import Iterator +import logging +import time from typing import Any, Dict, List, Optional, Sequence, Set, Tuple from .knowhere_hybrid import ( @@ -16,6 +18,7 @@ # Keep one bulk read bounded when a namespace contains a very large document, # while still replacing the per-document N+1 access pattern. _CORPUS_PREFETCH_GROUP_SIZE = 8 +_logger = logging.getLogger(__name__) def _children_ids(ts: Any, section_id: str, doc_id: str) -> List[str]: @@ -412,10 +415,17 @@ def compute_corpus_map_and_unit_scores_many( str, Tuple[Dict[str, List[str]], Set[str], Dict[str, str]], ] = {} + tree_started = time.perf_counter() for doc_id in valid_doc_ids: root_ids = list(ts.sections_for_doc(doc_id)) children_map, leaves, titles = _walk_tree(ts, doc_id, root_ids) tree_by_doc[doc_id] = (children_map, leaves, titles) + _logger.info( + "retrieval mapnav phase=tree_build seconds=%.3f documents=%d sections=%d", + time.perf_counter() - tree_started, + len(valid_doc_ids), + sum(len(value[0]) for value in tree_by_doc.values()), + ) def unit_factory() -> Iterator[ScoreUnitRow]: prefetch_batch = getattr(ts, "prefetch_document_units_batch", None) @@ -464,17 +474,32 @@ def unit_factory() -> Iterator[ScoreUnitRow]: release(document_id) persisted_loader = getattr(ts, "load_persisted_score_corpus", None) + loader_started = time.perf_counter() persisted_corpus = ( persisted_loader(valid_doc_ids, unique_queries) if callable(persisted_loader) else None ) + _logger.info( + "retrieval mapnav phase=index_load seconds=%.3f persisted=%s", + time.perf_counter() - loader_started, + persisted_corpus is not None, + ) + score_started = time.perf_counter() unit_scores_by_query = ( score_persisted_corpus_many(persisted_corpus, unique_queries) if persisted_corpus is not None else score_unit_stream_hybrid_many(unit_factory, unique_queries) ) + _logger.info( + "retrieval mapnav phase=unit_scoring persisted=%s seconds=%.3f units=%d queries=%d", + persisted_corpus is not None, + time.perf_counter() - score_started, + sum(len(scores) for scores in unit_scores_by_query.values()), + len(unique_queries), + ) results: Dict[str, Tuple[Dict[str, float], Dict[str, float]]] = {} + pooling_started = time.perf_counter() for query in unique_queries: unit_scores = unit_scores_by_query.get(query, {}) map_scores: Dict[str, float] = {} @@ -490,6 +515,12 @@ def unit_factory() -> Iterator[ScoreUnitRow]: ) map_scores[doc_id] = doc_max results[query] = (map_scores, unit_scores) + _logger.info( + "retrieval mapnav phase=map_pooling seconds=%.3f documents=%d sections=%d", + time.perf_counter() - pooling_started, + len(valid_doc_ids), + sum(len(value[0]) for value in tree_by_doc.values()), + ) return results diff --git a/packages/shared-python/shared/services/retrieval/nav_snapshot.py b/packages/shared-python/shared/services/retrieval/nav_snapshot.py index 1edfd374..17668767 100644 --- a/packages/shared-python/shared/services/retrieval/nav_snapshot.py +++ b/packages/shared-python/shared/services/retrieval/nav_snapshot.py @@ -8,6 +8,8 @@ from __future__ import annotations import json +import logging +import time from dataclasses import dataclass from typing import Any, Protocol @@ -33,7 +35,10 @@ # statement bounded. The contract benchmark verifies this page size against # the full 2 KiB content and metadata payload. _CHUNK_BATCH_SIZE = 10_000 -_REVISION_GROUP_SIZE = 32 +# Keep revision predicates bounded while reducing round trips for large +# namespaces. Keyset paging still caps each payload query at 10,000 rows. +_REVISION_GROUP_SIZE = 64 +_logger = logging.getLogger(__name__) class SnapshotSession(Protocol): @@ -252,6 +257,9 @@ async def _load_chunk_index( ref_index: dict[str, dict[str, Any]] = {} root_assets_by_doc: dict[str, set[str]] = {} text_connections_by_doc: dict[str, list[tuple[str, str]]] = {} + query_seconds = 0.0 + assembly_seconds = 0.0 + query_count = 0 for group_start in range(0, len(document_revisions), _REVISION_GROUP_SIZE): revision_group = document_revisions[group_start : group_start + _REVISION_GROUP_SIZE] last_key: tuple[str, str, int, str, str] | None = None @@ -289,9 +297,13 @@ async def _load_chunk_index( ) > tuple_(*[literal(value) for value in last_key]) ) + query_started = time.perf_counter() rows = (await db.execute(stmt)).all() + query_seconds += time.perf_counter() - query_started + query_count += 1 if not rows: break + assembly_started = time.perf_counter() for row in rows: document_id = str(row[0]) job_result_id = str(row[1]) @@ -339,6 +351,7 @@ async def _load_chunk_index( text_connections_by_doc.setdefault(document_id, []).append( (section_id or "", target) ) + assembly_seconds += time.perf_counter() - assembly_started last = rows[-1] last_key = ( str(last[0]), @@ -356,6 +369,13 @@ async def _load_chunk_index( if target in asset_ids: owners.setdefault(section_id, []).append(target) remounted[document_id] = {"root": sorted(asset_ids), "owners": owners} + _logger.info( + "retrieval snapshot phase=chunk_index rows=%d queries=%d query_seconds=%.3f assembly_seconds=%.3f", + sum(len(chunk_ids) for chunk_ids in ids_by_doc.values()), + query_count, + query_seconds, + assembly_seconds, + ) return ids_by_doc, ref_index, remounted @@ -392,7 +412,11 @@ async def _load_sections( by_doc: dict[str, list[SectionRow]] = {} path_by_id: dict[str, str] = {} - for row in (await db.execute(stmt)).all(): + query_started = time.perf_counter() + rows = (await db.execute(stmt)).all() + query_seconds = time.perf_counter() - query_started + assembly_started = time.perf_counter() + for row in rows: document_id = str(row[0]) section_path = str(row[3] or "") if is_excluded_section( @@ -413,6 +437,12 @@ async def _load_sections( ) by_doc.setdefault(document_id, []).append(section) path_by_id[section_id] = section_path + _logger.info( + "retrieval snapshot phase=sections rows=%d query_seconds=%.3f assembly_seconds=%.3f", + len(rows), + query_seconds, + time.perf_counter() - assembly_started, + ) return by_doc, path_by_id From 0d859b93ca1dc10d68e7598c405084be25b2b5b4 Mon Sep 17 00:00:00 2001 From: suguanYang Date: Sat, 29 Aug 2026 16:08:54 +0800 Subject: [PATCH 03/19] perf: defer lazy reference metadata and tune token lookup --- ...b6c7d_add_map_unit_token_covering_index.py | 40 +++++ .../test_retrieval_map_unit_index_contract.py | 147 +++++++++++++++++- .../shared/models/database/document.py | 7 + .../services/retrieval/nav/nav_knowhere.py | 41 ++++- .../shared/services/retrieval/nav_snapshot.py | 71 +++++++-- 5 files changed, 290 insertions(+), 16 deletions(-) create mode 100644 apps/api/alembic/versions/2e3f4a5b6c7d_add_map_unit_token_covering_index.py diff --git a/apps/api/alembic/versions/2e3f4a5b6c7d_add_map_unit_token_covering_index.py b/apps/api/alembic/versions/2e3f4a5b6c7d_add_map_unit_token_covering_index.py new file mode 100644 index 00000000..61250639 --- /dev/null +++ b/apps/api/alembic/versions/2e3f4a5b6c7d_add_map_unit_token_covering_index.py @@ -0,0 +1,40 @@ +"""Add a covering index for persisted map-nav token lookups.""" + +from __future__ import annotations + +from alembic import op + + +revision = "2e3f4a5b6c7d" +down_revision = "1d2e3f4a5b6c" +branch_labels = None +depends_on = None + +_INDEX_NAME = "idx_document_map_unit_tokens_unit_lookup" + + +def upgrade() -> None: + external_transaction = bool( + op.get_context().opts.get("knowhere_external_transaction", False) + ) + statement = ( + f"CREATE INDEX {{concurrently}}IF NOT EXISTS {_INDEX_NAME} " + "ON document_map_unit_tokens (map_unit_id, channel, token_hash) " + "INCLUDE (token, frequency)" + ) + if external_transaction: + op.execute(statement.format(concurrently="")) + return + with op.get_context().autocommit_block(): + op.execute(statement.format(concurrently="CONCURRENTLY ")) + + +def downgrade() -> None: + external_transaction = bool( + op.get_context().opts.get("knowhere_external_transaction", False) + ) + if external_transaction: + op.execute(f"DROP INDEX IF EXISTS {_INDEX_NAME}") + return + with op.get_context().autocommit_block(): + op.execute(f"DROP INDEX CONCURRENTLY IF EXISTS {_INDEX_NAME}") diff --git a/apps/api/tests/contract/test_retrieval_map_unit_index_contract.py b/apps/api/tests/contract/test_retrieval_map_unit_index_contract.py index 2f2027b0..56a2dced 100644 --- a/apps/api/tests/contract/test_retrieval_map_unit_index_contract.py +++ b/apps/api/tests/contract/test_retrieval_map_unit_index_contract.py @@ -5,7 +5,7 @@ from uuid import uuid4 from httpx import AsyncClient -from sqlalchemy import delete, select +from sqlalchemy import delete, select, text from shared.models.database.document import ( DocumentMapUnit, @@ -13,9 +13,11 @@ DocumentMapUnitToken, ) from shared.services.retrieval.nav.nav_hierarchy import ProviderToolSpace +from shared.services.retrieval.nav._compat import Chunk, EpisodeResult from shared.services.retrieval.nav.nav_map_scores import ( build_score_units, compute_corpus_map_and_unit_scores, + select_map_highlights, ) from shared.services.retrieval.nav.nav_knowhere import ( KnowhereProvider, @@ -29,6 +31,7 @@ from shared.services.retrieval.publication_content import ( replace_document_revision_content, ) +from shared.services.retrieval.nav_bridge import build_referenced_chunks from shared.services.retrieval.publication_models import DocumentPublicationScope from tests.support.contract_database import ContractDatabase from tests.support.retrieval_snapshot_support import contract_db_session @@ -179,6 +182,7 @@ async def test_published_map_units_preserve_scores_without_chunk_payload_reads( doc_ids=[document_id], query="common alpha", ) + expected_highlights = select_map_highlights(expected_scores[1], k=3) async with contract_db_session() as db: index = ( @@ -208,6 +212,18 @@ async def test_published_map_units_preserve_scores_without_chunk_payload_reads( ) ).scalars() ) + index_names = { + str(row[0]) + for row in ( + await db.execute( + text( + "SELECT indexname FROM pg_indexes " + "WHERE schemaname = current_schema() " + "AND tablename = 'document_map_unit_tokens'" + ) + ) + ).all() + } lazy_snapshot = await load_nav_snapshot( db, user_id=_USER_ID, @@ -220,6 +236,7 @@ async def test_published_map_units_preserve_scores_without_chunk_payload_reads( str(unit["chunk_id"]) for unit in expected_units ] assert persisted_tokens + assert "idx_document_map_unit_tokens_lookup" in index_names def reject_payload_read( _store: ReadOnlyChunkStore, @@ -274,9 +291,137 @@ def reject_payload_read( eager_snapshot.close() assert actual_scores == expected_scores + assert select_map_highlights(actual_scores[1], k=3) == expected_highlights assert fallback_scores == expected_scores +async def test_lazy_snapshot_defers_selected_asset_reference_metadata( + developer_api_client_factory: Callable[ + [], AbstractAsyncContextManager[AsyncClient] + ], + monkeypatch, +) -> None: + identifier = uuid4().hex[:8] + namespace = f"lazy-ref-{identifier}" + document_id = f"doc_ref_{identifier}" + job_id = f"job_ref_{identifier}" + job_result_id = f"result_ref_{identifier}" + async with developer_api_client_factory(): + await _seed_revision( + namespace=namespace, + document_id=document_id, + job_id=job_id, + job_result_id=job_result_id, + ) + scope = DocumentPublicationScope( + user_id=_USER_ID, + namespace=namespace, + document_id=document_id, + job_result_id=job_result_id, + source_file_name="refs.pdf", + ) + chunks = [ + { + "chunk_id": "body", + "type": "text", + "content": "body evidence", + "path": "refs.pdf/Root/Section/body", + "order": 1, + "metadata": {"connect_to": [{"target": "asset"}]}, + }, + { + "chunk_id": "asset", + "type": "image", + "content": "image description", + "path": "refs.pdf/Root/image", + "order": 2, + "file_path": "images/asset.png", + "metadata": {}, + }, + ] + async with contract_db_session() as db: + await db.run_sync( + lambda sync_db: replace_document_revision_content( + sync_db, + scope=scope, + chunks=chunks, + ) + ) + await db.commit() + + calls: list[tuple[str, str]] = [] + original = ReadOnlyChunkStore.load_chunk_reference_metadata + + def record_reference_load( + store: ReadOnlyChunkStore, + document: str, + chunk: str, + ) -> Mapping[str, object] | None: + calls.append((document, chunk)) + return original(store, document, chunk) + + monkeypatch.setattr( + ReadOnlyChunkStore, + "load_chunk_reference_metadata", + record_reference_load, + ) + async with contract_db_session() as db: + snapshot = await load_nav_snapshot( + db, + user_id=_USER_ID, + namespace=namespace, + lazy=True, + ) + + assert calls == [] + assert snapshot.chunk_ref_index[f"{document_id}:asset"]["file_path"] == ( + "images/asset.png" + ) + assert calls == [(document_id, "asset")] + episode = EpisodeResult( + representation="", + steps=[], + scored_chunks=[ + ( + Chunk( + node_id="asset", + doc_id=document_id, + text="image description", + line_ids=(2,), + section_id="root", + ), + 0.75, + ) + ], + kept_chunks=[ + Chunk( + node_id="asset", + doc_id=document_id, + text="image description", + line_ids=(2,), + section_id="root", + ) + ], + evidence_text="image description", + evidence_chars_actual=17, + retrieved_nodes=["asset"], + ) + refs, scores = build_referenced_chunks(episode, snapshot) + assert refs == [ + { + "chunk_id": "asset", + "document_id": document_id, + "chunk_type": "image", + "section_path": "Root / image", + "file_path": "images/asset.png", + "job_id": job_id, + "score": 0.75, + } + ] + assert scores == {"asset": 0.75} + snapshot.close() + + def test_incomplete_index_falls_back_for_duplicate_unit_ids() -> None: first_sections = [ SectionRow("root-a", None, "Root A", "Root A", 0, "", 0), diff --git a/packages/shared-python/shared/models/database/document.py b/packages/shared-python/shared/models/database/document.py index 33dda2f9..ad59857e 100644 --- a/packages/shared-python/shared/models/database/document.py +++ b/packages/shared-python/shared/models/database/document.py @@ -301,6 +301,13 @@ class DocumentMapUnitToken(Base): "token_hash", "map_unit_id", ), + Index( + "idx_document_map_unit_tokens_unit_lookup", + "map_unit_id", + "channel", + "token_hash", + postgresql_include=["token", "frequency"], + ), Index("idx_document_map_unit_tokens_unit", "map_unit_id", "channel"), ) 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 c056a2d2..b7f3f119 100644 --- a/packages/shared-python/shared/services/retrieval/nav/nav_knowhere.py +++ b/packages/shared-python/shared/services/retrieval/nav/nav_knowhere.py @@ -53,6 +53,7 @@ ROOT_SECTION_PATH = "Root" _DEFAULT_DSN = "postgresql://root:root123@127.0.0.1:5433/Knowhere" _MAP_UNIT_INDEX_FORMAT_VERSION = 1 +_MAP_SCORE_CHANNELS: Tuple[str, str] = ("path", "content") _logger = logging.getLogger(__name__) @@ -164,6 +165,13 @@ def knowhere_database_url() -> str: class ChunkStore(Protocol): + def load_chunk_reference_metadata( + self, + document_id: str, + chunk_id: str, + ) -> Optional[Mapping[str, Any]]: + raise NotImplementedError + def load_persisted_score_corpus( self, document_ids: Sequence[str], @@ -265,6 +273,30 @@ def load_section_units( finally: cur.close() + def load_chunk_reference_metadata( + self, + document_id: str, + chunk_id: str, + ) -> Optional[Mapping[str, Any]]: + """Resolve deferred reference fields for one selected chunk.""" + doc_id = str(document_id).strip() + cid = str(chunk_id).strip() + job_result_id = self._revisions.get(doc_id) + if not doc_id or not cid or not job_result_id: + return None + cur = self._connection().cursor() + try: + cur.execute( + "SELECT file_path FROM document_chunks " + "WHERE document_id = %s AND job_result_id = %s AND chunk_id = %s " + "ORDER BY sort_order DESC, id DESC LIMIT 1", + (doc_id, job_result_id, cid), + ) + row = cur.fetchone() + return {"file_path": str(row[0] or "") or None} if row else None + finally: + cur.close() + def load_persisted_score_corpus( self, document_ids: Sequence[str], @@ -404,8 +436,13 @@ def load_persisted_score_corpus( "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), + "AND token = ANY(%s) AND channel = ANY(%s)", + ( + map_unit_ids, + query_token_hashes, + query_tokens, + list(_MAP_SCORE_CHANNELS), + ), ) for map_unit_id, channel, token, frequency in cur.fetchall(): frequencies.setdefault((str(map_unit_id), str(channel)), {})[ diff --git a/packages/shared-python/shared/services/retrieval/nav_snapshot.py b/packages/shared-python/shared/services/retrieval/nav_snapshot.py index 17668767..4eec0404 100644 --- a/packages/shared-python/shared/services/retrieval/nav_snapshot.py +++ b/packages/shared-python/shared/services/retrieval/nav_snapshot.py @@ -11,6 +11,7 @@ import logging import time from dataclasses import dataclass +from collections.abc import Callable, Iterator, Mapping from typing import Any, Protocol from sqlalchemy import Executable, literal, select, tuple_ @@ -53,7 +54,7 @@ class NavSnapshot: """In-memory corpus for one map-nav episode.""" provider: NamespaceKnowhereProvider - chunk_ref_index: dict[str, dict[str, Any]] + chunk_ref_index: Mapping[str, dict[str, Any]] document_ids: list[str] document_titles: dict[str, str] @@ -199,11 +200,9 @@ async def load_nav_snapshot( 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") + chunk_id: document_id + for document_id, chunk_ids in chunk_ids_by_doc.items() + for chunk_id in chunk_ids }, ) except Exception: @@ -211,7 +210,10 @@ async def load_nav_snapshot( raise return NavSnapshot( provider=provider, - chunk_ref_index=dict(chunk_ref_index), + chunk_ref_index=LazyChunkRefIndex( + chunk_ref_index, + resolver=store.load_chunk_reference_metadata, + ), document_ids=list(provider.document_ids()), document_titles={did: kept_titles.get(did, did) for did in provider.document_ids()}, ) @@ -252,7 +254,7 @@ async def _load_chunk_index( section_path_by_id: dict[str, str], job_id_by_result_id: dict[str, str], ) -> tuple[dict[str, list[str]], dict[str, dict[str, Any]], dict[str, dict[str, Any]]]: - """Load only IDs/reference metadata; content remains lazy.""" + """Load only IDs/reference metadata; content and asset paths remain lazy.""" ids_by_doc: dict[str, list[str]] = {} ref_index: dict[str, dict[str, Any]] = {} root_assets_by_doc: dict[str, set[str]] = {} @@ -271,7 +273,8 @@ async def _load_chunk_index( DocumentChunk.chunk_id, DocumentChunk.section_id, DocumentChunk.chunk_type, - DocumentChunk.file_path, + # Asset paths are only needed for selected references and + # are resolved by ``LazyChunkRefIndex`` at bridge time. DocumentChunk.chunk_metadata["connect_to"].label("connect_to"), DocumentChunk.sort_order, DocumentChunk.id, @@ -323,7 +326,7 @@ async def _load_chunk_index( "document_id": document_id, "section_path": section_path, "chunk_type": chunk_type, - "file_path": str(row[5] or "") or None, + "file_path": None, "job_id": job_id_by_result_id.get(job_result_id), } ids_by_doc.setdefault(document_id, []).append(chunk_id) @@ -335,7 +338,7 @@ async def _load_chunk_index( and section_path == "Root" ): root_assets_by_doc.setdefault(document_id, set()).add(chunk_id) - connections = row[6] + connections = row[5] if isinstance(connections, str) and connections.strip(): try: connections = json.loads(connections) @@ -356,9 +359,9 @@ async def _load_chunk_index( last_key = ( str(last[0]), str(last[1]), - int(last[7] or 0), + int(last[6] or 0), str(last[2] or ""), - str(last[8]), + str(last[7]), ) if len(rows) < _CHUNK_BATCH_SIZE: break @@ -379,6 +382,48 @@ async def _load_chunk_index( return ids_by_doc, ref_index, remounted +class LazyChunkRefIndex(Mapping[str, dict[str, Any]]): + """Reference metadata map that resolves selected asset paths on demand. + + Snapshot construction still records every chunk identity, section path, + type, and job id so navigation ownership is unchanged. ``file_path`` is + fetched only when the exit bridge asks for a selected reference. + """ + + def __init__( + self, + base: Mapping[str, dict[str, Any]], + *, + resolver: Callable[[str, str], Mapping[str, Any] | None], + ) -> None: + self._base = {str(key): dict(value) for key, value in base.items()} + self._resolver = resolver + + def __getitem__(self, key: str) -> dict[str, Any]: + value = self.get(key) + if value is None: + raise KeyError(key) + return value + + def __iter__(self) -> Iterator[str]: + return iter(self._base) + + def __len__(self) -> int: + return len(self._base) + + def get(self, key: str, default: Any = None) -> dict[str, Any] | Any: + value = self._base.get(str(key)) + if value is None: + return default + if value.get("file_path") is None: + document_id = str(value.get("document_id") or "").strip() + chunk_id = str(key).split(":", 1)[-1].strip() + resolved = self._resolver(document_id, chunk_id) + if resolved is not None: + value.update({"file_path": resolved.get("file_path") or None}) + return value + + async def _load_sections( db: SnapshotSession, *, From 00fc7e2981a51ca34a8b5331b7f483efe4a4b9fb Mon Sep 17 00:00:00 2001 From: suguanYang Date: Sat, 29 Aug 2026 17:35:37 +0800 Subject: [PATCH 04/19] perf: trust transactional map index marker --- .../test_retrieval_map_unit_index_contract.py | 5 +++ .../services/retrieval/nav/nav_knowhere.py | 44 +++---------------- 2 files changed, 10 insertions(+), 39 deletions(-) diff --git a/apps/api/tests/contract/test_retrieval_map_unit_index_contract.py b/apps/api/tests/contract/test_retrieval_map_unit_index_contract.py index 56a2dced..eb124fb2 100644 --- a/apps/api/tests/contract/test_retrieval_map_unit_index_contract.py +++ b/apps/api/tests/contract/test_retrieval_map_unit_index_contract.py @@ -274,6 +274,11 @@ def reject_payload_read( await db.execute( delete(DocumentMapUnitToken).where(DocumentMapUnitToken.id == token_id) ) + await db.execute( + delete(DocumentMapUnitIndex).where( + DocumentMapUnitIndex.document_id == document_id + ) + ) await db.commit() incomplete_snapshot = await load_nav_snapshot( db, 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 b7f3f119..2aad0b29 100644 --- a/packages/shared-python/shared/services/retrieval/nav/nav_knowhere.py +++ b/packages/shared-python/shared/services/retrieval/nav/nav_knowhere.py @@ -338,45 +338,11 @@ def load_persisted_score_corpus( ): return None - stage_started = time.perf_counter() - 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 - _logger.info( - "retrieval map-index load stage=unit_count seconds=%.3f rows=%d", - time.perf_counter() - stage_started, - indexed_unit_count, - ) - 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 - stage_started = time.perf_counter() - 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 - _logger.info( - "retrieval map-index load stage=token_count seconds=%.3f rows=%d", - time.perf_counter() - stage_started, - indexed_token_count, - ) - expected_token_count = sum(int(row[4]) for row in manifests) - if indexed_token_count != expected_token_count: - return None + # The marker is written last in the same transaction that inserts + # all units and token rows. A committed marker therefore denotes + # one complete revision snapshot; avoid recounting millions of + # token rows on every request. Any rebuild deletes the marker + # first, so readers fall back to legacy scoring until completion. # The public chunk id is content-derived and may repeat within a # revision. The persisted scorer keys scores by that id, so use # the legacy payload path whenever ambiguity would change results. From 1f335316e1edb1ff262138504272332ffad38a28 Mon Sep 17 00:00:00 2001 From: suguanYang Date: Sun, 30 Aug 2026 12:40:39 +0800 Subject: [PATCH 05/19] perf: add revision-pinned retrieval serving index --- CONTEXT.md | 70 +++ ...6c7d8e_add_section_snapshot_order_index.py | 41 ++ ...d8e9f_add_retrieval_serving_generations.py | 171 +++++++ .../5b6c7d8e9f0a_add_term_trigram_indexes.py | 35 ++ .../services/documents/lifecycle_service.py | 47 +- apps/api/scripts/backfill_map_unit_indexes.py | 70 ++- .../test_bm25_fts_prefilter_contract.py | 38 +- .../tests/contract/test_documents_contract.py | 86 +++- .../test_retrieval_manifest_cache_contract.py | 56 +++ .../test_retrieval_mapnav_session_contract.py | 4 + ...test_retrieval_relit_map_cache_contract.py | 27 ++ .../test_retrieval_revision_races_contract.py | 102 ++++ .../test_retrieval_rrf_duplicate_contract.py | 19 + ...est_retrieval_serving_manifest_contract.py | 51 ++ ...retrieval_snapshot_consistency_contract.py | 54 +++ .../test_retrieval_term_score_contract.py | 75 +++ ...0005-stream-retrieval-progress-over-sse.md | 47 ++ ...mically-publish-retrieval-serving-index.md | 6 + ...-coherent-retrieval-serving-generations.md | 6 + ...enance-window-for-serving-index-rollout.md | 6 + docs/adr/README.md | 5 +- docs/design/retrieval-serving-index-plan.md | 425 ++++++++++++++++ docs/design/retrieval-streaming-sse.md | 207 ++++++++ .../shared/models/database/document.py | 171 +++++++ .../services/retrieval/execution/plan.py | 26 +- .../retrieval/execution/reference_resolver.py | 3 + .../retrieval/execution/revision_pins.py | 92 ++++ .../retrieval/execution/route_types.py | 3 + .../services/retrieval/execution/routes.py | 81 +++- .../services/retrieval/hydration/connected.py | 22 +- .../services/retrieval/hydration/reference.py | 44 +- .../retrieval/hydration/result_assembly.py | 3 + .../services/retrieval/manifest_cache.py | 47 ++ .../services/retrieval/nav/nav_knowhere.py | 308 +++++++++--- .../services/retrieval/nav/nav_orchestrate.py | 3 + .../services/retrieval/nav/nav_types.py | 6 + .../shared/services/retrieval/nav_snapshot.py | 453 ++++++++++++++---- .../services/retrieval/publication_content.py | 3 + .../services/retrieval/publication_service.py | 41 ++ .../services/retrieval/search/channels.py | 53 +- .../services/retrieval/search/discovery.py | 56 ++- .../services/retrieval/search/ranking.py | 109 +++-- .../retrieval/search/scoped_corpus.py | 119 ++++- .../services/retrieval/search/scoring.py | 10 +- .../services/retrieval/serving_generation.py | 60 +++ .../services/retrieval/serving_manifest.py | 405 ++++++++++++++++ 46 files changed, 3500 insertions(+), 266 deletions(-) create mode 100644 apps/api/alembic/versions/3f4a5b6c7d8e_add_section_snapshot_order_index.py create mode 100644 apps/api/alembic/versions/4a5b6c7d8e9f_add_retrieval_serving_generations.py create mode 100644 apps/api/alembic/versions/5b6c7d8e9f0a_add_term_trigram_indexes.py create mode 100644 apps/api/tests/contract/test_retrieval_manifest_cache_contract.py create mode 100644 apps/api/tests/contract/test_retrieval_relit_map_cache_contract.py create mode 100644 apps/api/tests/contract/test_retrieval_revision_races_contract.py create mode 100644 apps/api/tests/contract/test_retrieval_rrf_duplicate_contract.py create mode 100644 apps/api/tests/contract/test_retrieval_serving_manifest_contract.py create mode 100644 apps/api/tests/contract/test_retrieval_term_score_contract.py create mode 100644 docs/adr/0005-stream-retrieval-progress-over-sse.md create mode 100644 docs/adr/0006-atomically-publish-retrieval-serving-index.md create mode 100644 docs/adr/0007-use-coherent-retrieval-serving-generations.md create mode 100644 docs/adr/0008-use-a-maintenance-window-for-serving-index-rollout.md create mode 100644 docs/design/retrieval-serving-index-plan.md create mode 100644 docs/design/retrieval-streaming-sse.md create mode 100644 packages/shared-python/shared/services/retrieval/execution/revision_pins.py create mode 100644 packages/shared-python/shared/services/retrieval/manifest_cache.py create mode 100644 packages/shared-python/shared/services/retrieval/serving_generation.py create mode 100644 packages/shared-python/shared/services/retrieval/serving_manifest.py diff --git a/CONTEXT.md b/CONTEXT.md index c8903a58..b67e16d0 100644 --- a/CONTEXT.md +++ b/CONTEXT.md @@ -173,6 +173,76 @@ The query workflow that returns cited evidence from published documents. The typed retrieval request that owns cache-shaping fields and route policy: scope, filters, data type, channels, ranking options, and agentic toggle. +### Retrieval Run + +One execution of a Retrieval Query, recorded across classic, map-nav, +small-corpus, cache-hit, failed, and cancelled outcomes with route, timing, +and terminal-status metadata. + +### Retrieval Progress Event + +A safe, user-facing update about the current phase of a Retrieval run. It uses +the fixed phases `started`, `planning`, `searching`, `reviewing_sources`, and +`finalizing`; it never contains chain-of-thought or raw planner output. + +### Retrieval Stream + +The server-to-client event stream for a Retrieval Query. It carries +Retrieval Progress Events during execution and one authoritative final result +or terminal failure, while leaving answer generation to downstream clients. + +### Retrieval Duration + +The end-to-end server time for a Retrieval Query, measured from retrieval +execution start through final public-result assembly. It includes cache lookup +and excludes authentication, network transfer, SSE delivery time, and +downstream answer generation. + +### Retrieval Non-LLM Work + +The database and retrieval-engine work for a Retrieval Query: snapshot or +serving-index loading, lexical scoring, ranking, result hydration, citation +assembly, and asset-reference resolution. It excludes planner, harvest, +control, and answer-generation model time, which are measured separately. + +### Retrieval Serving Index + +The publication-derived read model used to load retrieval structure and +scoring inputs without rebuilding them from the full document corpus for each +query. It is revision-pinned and complete before its document revision becomes +active. + +### Retrieval Serving Fallback + +The exact legacy retrieval path used when a serving index is missing, +incomplete, or inconsistent. It preserves retrieval quality while sacrificing +the serving-index latency target until the derived data is repaired. + +### Retrieval Serving Generation + +The namespace-scoped version that identifies one coherent set of active +document revisions and their serving-index statistics. Retrieval captures one +generation and retries or falls back if publication changes it during capture. + +### Retrieval Semantic Parity + +The compatibility requirement that a serving-index retrieval returns the same +selected chunk IDs, ordering, rounded scores, citations, and asset references +as the legacy retrieval path for the same request. + +### Retrieval Revision Pin + +The set of document revision IDs captured at retrieval start and used for the +entire retrieval run, including lazy content and asset resolution. A later +publication affects subsequent runs, not the run already in progress. + +### Online Retrieval Serving Rollout + +The additive rollout of retrieval-serving schema and derived data while +retrieval and document publication remain available. Incomplete or +inconsistent revisions use the exact legacy retrieval path until backfill and +validation finish. + ### Workflow Run Request The agentic Retrieval request passed through planning and step execution. It diff --git a/apps/api/alembic/versions/3f4a5b6c7d8e_add_section_snapshot_order_index.py b/apps/api/alembic/versions/3f4a5b6c7d8e_add_section_snapshot_order_index.py new file mode 100644 index 00000000..4dcc5011 --- /dev/null +++ b/apps/api/alembic/versions/3f4a5b6c7d8e_add_section_snapshot_order_index.py @@ -0,0 +1,41 @@ +"""Add the index used by lazy map-nav section pagination.""" + +from __future__ import annotations + +from alembic import op + + +revision = "3f4a5b6c7d8e" +down_revision = "2e3f4a5b6c7d" +branch_labels = None +depends_on = None + +_INDEX_NAME = "idx_document_sections_revision_snapshot_order" + + +def upgrade() -> None: + external_transaction = bool( + op.get_context().opts.get("knowhere_external_transaction", False) + ) + statement = ( + f"CREATE INDEX {{concurrently}}IF NOT EXISTS {_INDEX_NAME} " + "ON document_sections " + "(document_id, job_result_id, sort_order, section_id)" + ) + if external_transaction: + op.execute(statement.format(concurrently="")) + return + with op.get_context().autocommit_block(): + op.execute(statement.format(concurrently="CONCURRENTLY ")) + + +def downgrade() -> None: + external_transaction = bool( + op.get_context().opts.get("knowhere_external_transaction", False) + ) + statement = f"DROP INDEX {{concurrently}}IF EXISTS {_INDEX_NAME}" + if external_transaction: + op.execute(statement.format(concurrently="")) + return + with op.get_context().autocommit_block(): + op.execute(statement.format(concurrently="CONCURRENTLY ")) diff --git a/apps/api/alembic/versions/4a5b6c7d8e9f_add_retrieval_serving_generations.py b/apps/api/alembic/versions/4a5b6c7d8e9f_add_retrieval_serving_generations.py new file mode 100644 index 00000000..426ff6e0 --- /dev/null +++ b/apps/api/alembic/versions/4a5b6c7d8e9f_add_retrieval_serving_generations.py @@ -0,0 +1,171 @@ +"""Add revision-pinned serving manifests and namespace statistics.""" + +from __future__ import annotations + +from collections.abc import Sequence + +from alembic import op +import sqlalchemy as sa + + +revision: str = "4a5b6c7d8e9f" +down_revision: str | None = "3f4a5b6c7d8e" +branch_labels: Sequence[str] | None = None +depends_on: Sequence[str] | None = None + + +def upgrade() -> None: + if not sa.inspect(op.get_bind()).has_table("retrieval_namespace_generations"): + op.create_table( + "retrieval_namespace_generations", + sa.Column("id", sa.String(length=100), nullable=False), + sa.Column("user_id", sa.Text(), nullable=False), + sa.Column("namespace", sa.String(length=255), nullable=False), + sa.Column( + "generation", sa.BigInteger(), nullable=False, server_default="0" + ), + sa.Column("updated_at", sa.DateTime(), nullable=False), + sa.PrimaryKeyConstraint("id"), + sa.UniqueConstraint( + "user_id", + "namespace", + name="uq_retrieval_namespace_generations_scope", + ), + ) + if not sa.inspect(op.get_bind()).has_table("retrieval_serving_revision_manifests"): + op.create_table( + "retrieval_serving_revision_manifests", + sa.Column("id", sa.String(length=100), nullable=False), + sa.Column("user_id", sa.Text(), nullable=False), + sa.Column("namespace", sa.String(length=255), nullable=False), + sa.Column("document_id", sa.String(length=36), nullable=False), + sa.Column("job_result_id", sa.String(length=36), nullable=False), + sa.Column("format_version", sa.Integer(), nullable=False), + sa.Column("payload_zlib", sa.LargeBinary(), nullable=False), + sa.Column("checksum", sa.String(length=64), nullable=False), + sa.Column("created_at", sa.DateTime(), nullable=False), + sa.ForeignKeyConstraint( + ["document_id"], ["documents.document_id"], ondelete="CASCADE" + ), + sa.ForeignKeyConstraint( + ["job_result_id"], ["job_results.id"], ondelete="CASCADE" + ), + sa.PrimaryKeyConstraint("id"), + sa.UniqueConstraint( + "document_id", + "job_result_id", + name="uq_retrieval_serving_revision_manifests_revision", + ), + ) + if not _index_exists("idx_retrieval_serving_revision_manifests_scope"): + op.create_index( + "idx_retrieval_serving_revision_manifests_scope", + "retrieval_serving_revision_manifests", + ["user_id", "namespace", "document_id", "job_result_id"], + ) + if not sa.inspect(op.get_bind()).has_table("retrieval_serving_revision_stats"): + op.create_table( + "retrieval_serving_revision_stats", + sa.Column("id", sa.String(length=100), nullable=False), + sa.Column("user_id", sa.Text(), nullable=False), + sa.Column("namespace", sa.String(length=255), nullable=False), + sa.Column("document_id", sa.String(length=36), nullable=False), + sa.Column("job_result_id", sa.String(length=36), nullable=False), + sa.Column("format_version", sa.Integer(), nullable=False), + sa.Column("payload_zlib", sa.LargeBinary(), nullable=False), + sa.Column("checksum", sa.String(length=64), nullable=False), + sa.Column("created_at", sa.DateTime(), nullable=False), + sa.ForeignKeyConstraint( + ["document_id"], ["documents.document_id"], ondelete="CASCADE" + ), + sa.ForeignKeyConstraint( + ["job_result_id"], ["job_results.id"], ondelete="CASCADE" + ), + sa.PrimaryKeyConstraint("id"), + sa.UniqueConstraint( + "document_id", + "job_result_id", + name="uq_retrieval_serving_revision_stats_revision", + ), + ) + if not _index_exists("idx_retrieval_serving_revision_stats_scope"): + op.create_index( + "idx_retrieval_serving_revision_stats_scope", + "retrieval_serving_revision_stats", + ["user_id", "namespace", "document_id", "job_result_id"], + ) + if not sa.inspect(op.get_bind()).has_table("retrieval_namespace_stats"): + op.create_table( + "retrieval_namespace_stats", + sa.Column("id", sa.String(length=100), nullable=False), + sa.Column("user_id", sa.Text(), nullable=False), + sa.Column("namespace", sa.String(length=255), nullable=False), + sa.Column("generation", sa.BigInteger(), nullable=False), + sa.Column("payload_zlib", sa.LargeBinary(), nullable=False), + sa.Column("checksum", sa.String(length=64), nullable=False), + sa.Column("updated_at", sa.DateTime(), nullable=False), + sa.PrimaryKeyConstraint("id"), + sa.UniqueConstraint( + "user_id", "namespace", name="uq_retrieval_namespace_stats_scope" + ), + ) + if not sa.inspect(op.get_bind()).has_table("retrieval_namespace_token_stats"): + op.create_table( + "retrieval_namespace_token_stats", + sa.Column("id", sa.String(length=100), nullable=False), + sa.Column("user_id", sa.Text(), nullable=False), + sa.Column("namespace", sa.String(length=255), nullable=False), + sa.Column("generation", sa.BigInteger(), nullable=False), + sa.Column("channel", sa.String(length=32), nullable=False), + sa.Column("token_hash", sa.String(length=64), nullable=False), + sa.Column("document_frequency", sa.Integer(), nullable=False), + sa.PrimaryKeyConstraint("id"), + sa.UniqueConstraint( + "user_id", + "namespace", + "channel", + "token_hash", + name="uq_retrieval_namespace_token_stats_key", + ), + ) + if not _index_exists("idx_retrieval_namespace_token_stats_lookup"): + op.create_index( + "idx_retrieval_namespace_token_stats_lookup", + "retrieval_namespace_token_stats", + ["user_id", "namespace", "generation", "channel", "token_hash"], + ) + + +def _index_exists(index_name: str) -> bool: + bind = op.get_bind() + inspector = sa.inspect(bind) + for table_name in inspector.get_table_names(): + if any( + index.get("name") == index_name + for index in inspector.get_indexes(table_name) + ): + return True + return False + + +def downgrade() -> None: + op.drop_index( + "idx_retrieval_namespace_token_stats_lookup", + table_name="retrieval_namespace_token_stats", + if_exists=True, + ) + op.drop_table("retrieval_namespace_token_stats", if_exists=True) + op.drop_table("retrieval_namespace_stats", if_exists=True) + op.drop_index( + "idx_retrieval_serving_revision_stats_scope", + table_name="retrieval_serving_revision_stats", + if_exists=True, + ) + op.drop_table("retrieval_serving_revision_stats", if_exists=True) + op.drop_index( + "idx_retrieval_serving_revision_manifests_scope", + table_name="retrieval_serving_revision_manifests", + if_exists=True, + ) + op.drop_table("retrieval_serving_revision_manifests", if_exists=True) + op.drop_table("retrieval_namespace_generations", if_exists=True) diff --git a/apps/api/alembic/versions/5b6c7d8e9f0a_add_term_trigram_indexes.py b/apps/api/alembic/versions/5b6c7d8e9f0a_add_term_trigram_indexes.py new file mode 100644 index 00000000..71720ff0 --- /dev/null +++ b/apps/api/alembic/versions/5b6c7d8e9f0a_add_term_trigram_indexes.py @@ -0,0 +1,35 @@ +"""Add trigram acceleration for exact term-channel candidate discovery.""" + +from __future__ import annotations + +from collections.abc import Sequence + +from alembic import op + + +revision: str = "5b6c7d8e9f0a" +down_revision: str | None = "4a5b6c7d8e9f" +branch_labels: Sequence[str] | None = None +depends_on: Sequence[str] | None = None + +_MAP_UNIT_INDEX = "idx_document_map_units_term_trgm" +_CHUNK_INDEX = "idx_document_chunks_term_trgm" + + +def upgrade() -> None: + op.execute("CREATE EXTENSION IF NOT EXISTS pg_trgm") + op.execute( + f"CREATE INDEX IF NOT EXISTS {_MAP_UNIT_INDEX} " + "ON document_map_units USING gin " + "(term_search_text_lower gin_trgm_ops)" + ) + op.execute( + f"CREATE INDEX IF NOT EXISTS {_CHUNK_INDEX} " + "ON document_chunks USING gin " + "(lower(COALESCE(term_search_text, '')) gin_trgm_ops)" + ) + + +def downgrade() -> None: + op.execute(f"DROP INDEX IF EXISTS {_CHUNK_INDEX}") + op.execute(f"DROP INDEX IF EXISTS {_MAP_UNIT_INDEX}") diff --git a/apps/api/app/services/documents/lifecycle_service.py b/apps/api/app/services/documents/lifecycle_service.py index 1bd5ae3a..6b39507d 100644 --- a/apps/api/app/services/documents/lifecycle_service.py +++ b/apps/api/app/services/documents/lifecycle_service.py @@ -8,13 +8,25 @@ from app.repositories.document_repository import DocumentRepository from loguru import logger +from sqlalchemy import delete from sqlalchemy.ext.asyncio import AsyncSession -from shared.models.database.document import DocumentChunk, DocumentSection +from shared.models.database.document import ( + DocumentChunk, + DocumentSection, + RetrievalServingRevisionStat, +) from shared.services.retrieval.cache_service import ( invalidate_retrieval_cache_namespaces, ) from shared.services.retrieval.graph.service import DocumentGraphService, GraphScope +from shared.services.retrieval.serving_generation import ( + advance_namespace_generation, + lock_namespace_generation, +) +from shared.services.retrieval.serving_manifest import ( + rebuild_namespace_serving_statistics, +) from shared.services.storage.result_storage import ResultStorage, get_result_storage _DOCUMENT_CHUNK_ASSET_URL_EXPIRES_SECONDS = 7 * 24 * 60 * 60 @@ -103,7 +115,7 @@ def _normalize_page_asset(raw_asset: dict[str, Any]) -> dict[str, Any] | None: "content_type": content_type, "source": source, } - if (asset_url := str(raw_asset.get("asset_url") or "").strip()): + if asset_url := str(raw_asset.get("asset_url") or "").strip(): asset["asset_url"] = asset_url if (width := _positive_int(raw_asset.get("width"))) is not None: asset["width"] = width @@ -446,7 +458,38 @@ async def archive_document( return document_payload(document) previous_namespace = document.namespace + await db.run_sync( + lambda sync_db: lock_namespace_generation( + sync_db, + user_id=user_id, + namespace=previous_namespace, + ) + ) await self._repository.archive_document(db, document=document) + current_revision = document.current_job_result_id + if current_revision: + await db.run_sync( + lambda sync_db: sync_db.execute( + delete(RetrievalServingRevisionStat).where( + RetrievalServingRevisionStat.document_id == document_id, + RetrievalServingRevisionStat.job_result_id == current_revision, + ) + ) + ) + await db.run_sync( + lambda sync_db: rebuild_namespace_serving_statistics( + sync_db, + user_id=user_id, + namespace=previous_namespace, + ) + ) + await db.run_sync( + lambda sync_db: advance_namespace_generation( + sync_db, + user_id=user_id, + namespace=previous_namespace, + ) + ) await db.run_sync( lambda sync_db: self._graph_service.remove_document_graph( sync_db, diff --git a/apps/api/scripts/backfill_map_unit_indexes.py b/apps/api/scripts/backfill_map_unit_indexes.py index 78a88f6c..6a8a79db 100644 --- a/apps/api/scripts/backfill_map_unit_indexes.py +++ b/apps/api/scripts/backfill_map_unit_indexes.py @@ -33,6 +33,14 @@ def _bootstrap_python_path() -> None: from shared.models.database.document import Document from shared.services.retrieval.map_unit_index import replace_document_map_units from shared.services.retrieval.publication_models import DocumentPublicationScope +from shared.services.retrieval.serving_generation import ( + advance_namespace_generation, + lock_namespace_generation, +) +from shared.services.retrieval.serving_manifest import ( + persist_revision_serving_state, + rebuild_namespace_serving_statistics, +) def _build_parser() -> argparse.ArgumentParser: @@ -44,14 +52,20 @@ def _build_parser() -> argparse.ArgumentParser: action="store_true", help="Build and commit each current revision index.", ) - parser.add_argument("--document-id", default="", help="Limit the backfill to one document.") + 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)) + statement = ( + select(Document) + .where(Document.status == "active") + .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) @@ -62,7 +76,9 @@ 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}") + print( + f"would backfill document={document.document_id} revision={document.current_job_result_id}" + ) return len(documents) session_factory = get_sync_session_factory() @@ -70,15 +86,49 @@ def backfill_map_unit_indexes(*, apply: bool, document_id: str = "") -> int: job_result_id = document.current_job_result_id if not job_result_id: continue - scope = DocumentPublicationScope( - user_id=document.user_id, - namespace=document.namespace, - document_id=document.document_id, - job_result_id=job_result_id, - source_file_name=str(document.source_file_name or ""), - ) with session_factory() as db: + lock_namespace_generation( + db, + user_id=document.user_id, + namespace=document.namespace, + ) + locked_document = db.execute( + select(Document) + .where(Document.document_id == document.document_id) + .with_for_update() + ).scalar_one_or_none() + if ( + locked_document is None + or locked_document.status != "active" + or locked_document.current_job_result_id != job_result_id + or locked_document.user_id != document.user_id + or locked_document.namespace != document.namespace + ): + db.rollback() + print( + f"skipped stale or inactive document={document.document_id} " + f"revision={job_result_id}" + ) + continue + scope = DocumentPublicationScope( + user_id=locked_document.user_id, + namespace=locked_document.namespace, + document_id=locked_document.document_id, + job_result_id=job_result_id, + source_file_name=str(locked_document.source_file_name or ""), + ) replace_document_map_units(db, scope=scope) + persist_revision_serving_state(db, scope=scope) + rebuild_namespace_serving_statistics( + db, + user_id=scope.user_id, + namespace=scope.namespace, + ) + advance_namespace_generation( + db, + user_id=scope.user_id, + namespace=scope.namespace, + ) db.commit() print(f"backfilled document={document.document_id} revision={job_result_id}") return len(documents) diff --git a/apps/api/tests/contract/test_bm25_fts_prefilter_contract.py b/apps/api/tests/contract/test_bm25_fts_prefilter_contract.py index 6b958329..56add8e5 100644 --- a/apps/api/tests/contract/test_bm25_fts_prefilter_contract.py +++ b/apps/api/tests/contract/test_bm25_fts_prefilter_contract.py @@ -23,10 +23,10 @@ user_id TEXT, namespace TEXT, status TEXT, - current_job_result_id INTEGER, + current_job_result_id TEXT, source_file_name TEXT ); -CREATE TABLE job_results (id INTEGER PRIMARY KEY, job_id TEXT); +CREATE TABLE job_results (id TEXT PRIMARY KEY, job_id TEXT); CREATE TABLE document_sections (section_id TEXT PRIMARY KEY, section_path TEXT); CREATE TABLE document_chunks ( id SERIAL PRIMARY KEY, @@ -38,7 +38,7 @@ source_chunk_path TEXT, file_path TEXT, chunk_metadata JSONB, - job_result_id INTEGER, + job_result_id TEXT, sort_order INTEGER, content_search_text TEXT, content_search_tsv TSVECTOR GENERATED ALWAYS AS @@ -197,3 +197,35 @@ async def test_exclusions_still_apply_under_the_prefilter( exclude_sections=[], ) assert rows == [] + + +@pytest.mark.asyncio +async def test_content_channel_uses_the_requested_revision_pin( + seeded_session: AsyncSession, +) -> None: + await seeded_session.execute(text("INSERT INTO job_results VALUES (2, 'job2')")) + await seeded_session.execute( + text( + "INSERT INTO document_chunks " + "(chunk_id, document_id, section_id, chunk_type, content, " + " job_result_id, sort_order, content_search_text, path_search_text) " + "VALUES ('new-hit', 'd1', 's1', 'text', 'new body', 2, 1, " + " 'alpha replacement', 'new path')" + ) + ) + await seeded_session.execute( + text("UPDATE documents SET current_job_result_id = 2 WHERE document_id = 'd1'") + ) + + rows = await content_channel( + seeded_session, + user_id="u1", + namespace="ns1", + query="alpha", + top_k=50, + exclude_document_ids=[], + exclude_sections=[], + revision_pins={"d1": "1"}, + ) + + assert [str(row["chunk_id"]) for row in rows] == ["hit-en"] diff --git a/apps/api/tests/contract/test_documents_contract.py b/apps/api/tests/contract/test_documents_contract.py index 6b8e33b0..a71d4029 100644 --- a/apps/api/tests/contract/test_documents_contract.py +++ b/apps/api/tests/contract/test_documents_contract.py @@ -13,6 +13,7 @@ from tests.support.contract_database import ContractDatabase from shared.testing.contract_runtime import get_contract_database_url +from shared.services.retrieval.serving_manifest import encode_serving_manifest async def _create_contract_engine() -> AsyncEngine: @@ -1019,9 +1020,7 @@ async def test_should_include_media_asset_urls_in_document_chunk_list_when_reque assert chunks[1]["asset_url"] == expected_asset_url assert default_response.status_code == 200 - default_chunks = cast( - list[dict[str, object]], default_response.json()["chunks"] - ) + default_chunks = cast(list[dict[str, object]], default_response.json()["chunks"]) assert default_chunks[1]["asset_url"] is None @@ -1334,3 +1333,84 @@ async def test_should_archive_a_document_via_the_legacy_archive_route( assert response_json["archived_at"] assert persisted_document["status"] == "archived" assert persisted_document["archived_at"] is not None + + +@pytest.mark.asyncio +async def test_archive_removes_revision_serving_stats_and_advances_generation( + developer_api_client_factory: Callable[ + [], AbstractAsyncContextManager[AsyncClient] + ], +) -> None: + document_id = f"doc_{uuid4().hex[:12]}" + namespace = f"archive-serving-{uuid4().hex[:8]}" + async with developer_api_client_factory() as api_client: + revision = await _insert_document_revision_with_chunks( + document_id=document_id, + namespace=namespace, + chunks=[ + { + "id": f"dchk_{uuid4().hex[:12]}", + "chunk_id": "archive-serving-chunk", + "chunk_type": "text", + "content": "serving contribution", + "source_chunk_path": "Archive/Serving", + "metadata": {}, + } + ], + ) + payload_bytes, checksum, version = encode_serving_manifest( + { + "document_id": document_id, + "job_result_id": revision["job_result_id"], + "unit_count": 1, + "path_token_count": 1, + "content_token_count": 2, + "token_frequencies": { + "path": {"archive": 1}, + "content": {"serving": 1}, + }, + } + ) + await ContractDatabase.execute( + """ + INSERT INTO retrieval_serving_revision_stats ( + id, user_id, namespace, document_id, job_result_id, + format_version, payload_zlib, checksum, created_at + ) VALUES ( + :id, :user_id, :namespace, :document_id, :job_result_id, + :format_version, :payload_zlib, :checksum, NOW() + ) + """, + { + "id": f"rss_{uuid4().hex[:12]}", + "user_id": "local-dev-user", + "namespace": namespace, + "document_id": document_id, + "job_result_id": revision["job_result_id"], + "format_version": version, + "payload_zlib": payload_bytes, + "checksum": checksum, + }, + ) + response = await api_client.post(f"/api/v1/documents/{document_id}/archive") + + assert response.status_code == 200 + remaining_revision_stats = await ContractDatabase.fetch_one( + """ + SELECT id + FROM retrieval_serving_revision_stats + WHERE document_id = :document_id AND job_result_id = :job_result_id + """, + {"document_id": document_id, "job_result_id": revision["job_result_id"]}, + ) + namespace_stats = await ContractDatabase.fetch_one( + """ + SELECT generation + FROM retrieval_namespace_stats + WHERE user_id = :user_id AND namespace = :namespace + """, + {"user_id": "local-dev-user", "namespace": namespace}, + ) + assert remaining_revision_stats is None + assert namespace_stats is not None + assert int(namespace_stats["generation"]) >= 1 diff --git a/apps/api/tests/contract/test_retrieval_manifest_cache_contract.py b/apps/api/tests/contract/test_retrieval_manifest_cache_contract.py new file mode 100644 index 00000000..e2f73e47 --- /dev/null +++ b/apps/api/tests/contract/test_retrieval_manifest_cache_contract.py @@ -0,0 +1,56 @@ +from __future__ import annotations + +import pytest +from sqlalchemy.exc import SQLAlchemyError + +from shared.services.retrieval.manifest_cache import ( + cache_manifest_payloads, + get_cached_manifest_payloads, +) +from shared.services.retrieval.search.scoped_corpus import count_manifest_chunks + + +class _SessionWithInfo: + def __init__(self) -> None: + self.info: dict[str, object] = {} + + +def test_manifest_payload_cache_is_scoped_to_revision_pin_set() -> None: + session = _SessionWithInfo() + revisions = {"doc-a": "result-a", "doc-b": "result-b"} + payloads = { + ("doc-a", "result-a"): {"chunks": [{"chunk_id": "chunk-a"}]}, + ("doc-b", "result-b"): {"chunks": [{"chunk_id": "chunk-b"}]}, + } + + cache_manifest_payloads(session, revisions=revisions, payloads=payloads) + + assert get_cached_manifest_payloads(session, revisions=revisions) == payloads + assert get_cached_manifest_payloads( + session, + revisions={"doc-a": "different-result"}, + ) is None + + +class _UnavailableManifestSession: + def __init__(self) -> None: + self.rollback_count = 0 + + async def execute(self, _statement: object) -> object: + raise SQLAlchemyError("serving manifest table is unavailable") + + async def rollback(self) -> None: + self.rollback_count += 1 + + +@pytest.mark.asyncio +async def test_manifest_count_falls_back_after_derived_table_error() -> None: + session = _UnavailableManifestSession() + + result = await count_manifest_chunks( + session, # type: ignore[arg-type] + revision_pins={"doc-a": "result-a"}, + ) + + assert result is None + assert session.rollback_count == 1 diff --git a/apps/api/tests/contract/test_retrieval_mapnav_session_contract.py b/apps/api/tests/contract/test_retrieval_mapnav_session_contract.py index a277c4c0..d267941b 100644 --- a/apps/api/tests/contract/test_retrieval_mapnav_session_contract.py +++ b/apps/api/tests/contract/test_retrieval_mapnav_session_contract.py @@ -135,12 +135,14 @@ async def fake_resolve_workflow_references( namespace: str, refs: list[RouteRow], score_by_chunk_id: dict[str, float] | None = None, + revision_pins: dict[str, str] | None = None, ) -> ResolvedWorkflowReferences: assert db is fresh_db assert user_id == "contract-user" assert namespace == "contract-namespace" assert refs assert score_by_chunk_id is not None + assert revision_pins is None events.append("resolve_references") row = { "document_id": "doc_contract", @@ -159,11 +161,13 @@ async def fake_assemble_retrieval_results( exclude_document_ids: list[str], exclude_sections: list[dict[str, str]], allowed_chunk_types: set[str] | None, + revision_pins: dict[str, str] | None = None, ) -> list[RouteRow]: assert db is fresh_db assert exclude_document_ids == [] assert exclude_sections == [] assert allowed_chunk_types is None + assert revision_pins is None events.append("assemble_results") return rows diff --git a/apps/api/tests/contract/test_retrieval_relit_map_cache_contract.py b/apps/api/tests/contract/test_retrieval_relit_map_cache_contract.py new file mode 100644 index 00000000..b8257c9c --- /dev/null +++ b/apps/api/tests/contract/test_retrieval_relit_map_cache_contract.py @@ -0,0 +1,27 @@ +from __future__ import annotations + +from shared.services.retrieval.nav.nav_orchestrate import _relit_map +from shared.services.retrieval.nav.nav_types import NavState + + +def test_relit_map_reuses_same_query_within_episode(monkeypatch) -> None: + calls: list[str] = [] + + def fake_relight_map_for_query(*_args, **kwargs): + calls.append(str(kwargs["query"])) + return {"section": 1.0}, {"unit": 2.0}, ["section"] + + monkeypatch.setattr( + "shared.services.retrieval.nav.nav_map_scores.relight_map_for_query", + fake_relight_map_for_query, + ) + state = NavState(doc_id="", query="retrieval") + config = type("Config", (), {"collect_top_k": 6})() + + with _relit_map(None, state, config, query="retrieval"): + pass + with _relit_map(None, state, config, query="retrieval"): + pass + + assert calls == ["retrieval"] + assert state.relit_map_cache["retrieval"][0] == {"section": 1.0} diff --git a/apps/api/tests/contract/test_retrieval_revision_races_contract.py b/apps/api/tests/contract/test_retrieval_revision_races_contract.py new file mode 100644 index 00000000..dd73d966 --- /dev/null +++ b/apps/api/tests/contract/test_retrieval_revision_races_contract.py @@ -0,0 +1,102 @@ +"""Deterministic contracts for revision and channel-session coherence.""" + +from __future__ import annotations + +from collections.abc import AsyncGenerator +from contextlib import asynccontextmanager +from typing import Any, cast + +import pytest +from sqlalchemy.ext.asyncio import AsyncSession + +from shared.services.retrieval.execution.revision_pins import ( + RetrievalRevisionPins, + is_revision_generation_stable, +) +from shared.services.retrieval.search import discovery + + +@pytest.mark.asyncio +async def test_classic_channels_share_pins_but_use_distinct_sessions( + monkeypatch: pytest.MonkeyPatch, +) -> None: + pins = RetrievalRevisionPins( + revisions={"doc-1": "revision-1"}, + generation=7, + ) + sessions: list[object] = [] + observed: list[tuple[object, object]] = [] + + @asynccontextmanager + async def fake_context() -> AsyncGenerator[object, None]: + session = object() + sessions.append(session) + yield session + + async def fake_channel( + db: AsyncSession, + **kwargs: Any, + ) -> list[dict[str, Any]]: + observed.append((db, kwargs["revision_pins"])) + return [] + + monkeypatch.setattr("shared.core.database.get_db_context", fake_context) + monkeypatch.setattr(discovery, "path_channel", fake_channel) + monkeypatch.setattr(discovery, "content_channel", fake_channel) + monkeypatch.setattr(discovery, "term_channel", fake_channel) + + result = await discovery.bottom_discovery( + cast(AsyncSession, object()), + user_id="user-1", + namespace="namespace-1", + query="coherent query", + top_k=3, + exclude_document_ids=[], + exclude_sections=[], + revision_pins=pins, + ) + + assert result.status == "discovery_done" + assert len(sessions) == 3 + assert len({id(session) for session in sessions}) == 3 + assert len(observed) == 3 + assert {id(session) for session, _pins in observed} == { + id(session) for session in sessions + } + assert all(observed_pins is pins for _session, observed_pins in observed) + + +class _GenerationResult: + def __init__(self, value: int | None) -> None: + self._value = value + + def scalar_one_or_none(self) -> int | None: + return self._value + + +class _GenerationSession: + def __init__(self, values: list[int | None]) -> None: + self._values = iter(values) + + async def execute(self, _statement: object) -> _GenerationResult: + return _GenerationResult(next(self._values)) + + +@pytest.mark.asyncio +async def test_generation_change_is_detected_before_scoring() -> None: + pins = RetrievalRevisionPins(revisions={"doc-1": "revision-1"}, generation=7) + stable = await is_revision_generation_stable( + cast(AsyncSession, _GenerationSession([7])), + user_id="user-1", + namespace="namespace-1", + pins=pins, + ) + changed = await is_revision_generation_stable( + cast(AsyncSession, _GenerationSession([8])), + user_id="user-1", + namespace="namespace-1", + pins=pins, + ) + + assert stable is True + assert changed is False diff --git a/apps/api/tests/contract/test_retrieval_rrf_duplicate_contract.py b/apps/api/tests/contract/test_retrieval_rrf_duplicate_contract.py new file mode 100644 index 00000000..e056af6c --- /dev/null +++ b/apps/api/tests/contract/test_retrieval_rrf_duplicate_contract.py @@ -0,0 +1,19 @@ +"""Contracts for duplicate chunk handling in reciprocal-rank fusion.""" + +from __future__ import annotations + +from shared.services.retrieval.search.scoring import merge_channels_rrf + + +def test_rrf_counts_each_chunk_once_per_channel() -> None: + rows = [ + {"chunk_id": "shared", "document_id": "doc-a"}, + {"chunk_id": "shared", "document_id": "doc-b"}, + {"chunk_id": "other", "document_id": "doc-c"}, + ] + + result = merge_channels_rrf([rows], [1.0], top_k=3) + + assert [row["chunk_id"] for row in result] == ["shared", "other"] + assert result[0]["score"] == round(1.0 / 61, 6) + assert result[1]["score"] == round(1.0 / 62, 6) diff --git a/apps/api/tests/contract/test_retrieval_serving_manifest_contract.py b/apps/api/tests/contract/test_retrieval_serving_manifest_contract.py new file mode 100644 index 00000000..5bb68c06 --- /dev/null +++ b/apps/api/tests/contract/test_retrieval_serving_manifest_contract.py @@ -0,0 +1,51 @@ +"""Contract tests for serving-manifest integrity and version handling.""" + +from __future__ import annotations + +import pytest + +from shared.services.retrieval.serving_manifest import ( + SERVING_MANIFEST_FORMAT_VERSION, + decode_serving_manifest, + encode_serving_manifest, +) + + +def test_serving_manifest_round_trip_preserves_canonical_payload() -> None: + payload = { + "document_id": "doc_contract", + "job_result_id": "result_contract", + "sections": [{"section_id": "sec_1", "sort_order": 0}], + "chunks": [{"chunk_id": "chunk_1", "connect_to": []}], + } + + compressed, checksum, version = encode_serving_manifest(payload) + + assert version == SERVING_MANIFEST_FORMAT_VERSION + assert decode_serving_manifest( + compressed, + checksum=checksum, + format_version=version, + ) == payload + + +def test_serving_manifest_rejects_checksum_mismatch() -> None: + compressed, _, version = encode_serving_manifest({"document_id": "doc"}) + + with pytest.raises(ValueError, match="checksum mismatch"): + decode_serving_manifest( + compressed, + checksum="0" * 64, + format_version=version, + ) + + +def test_serving_manifest_rejects_unknown_version() -> None: + compressed, checksum, _ = encode_serving_manifest({"document_id": "doc"}) + + with pytest.raises(ValueError, match="unsupported serving manifest version"): + decode_serving_manifest( + compressed, + checksum=checksum, + format_version=SERVING_MANIFEST_FORMAT_VERSION + 1, + ) diff --git a/apps/api/tests/contract/test_retrieval_snapshot_consistency_contract.py b/apps/api/tests/contract/test_retrieval_snapshot_consistency_contract.py index 64aaf4c0..bbe03e05 100644 --- a/apps/api/tests/contract/test_retrieval_snapshot_consistency_contract.py +++ b/apps/api/tests/contract/test_retrieval_snapshot_consistency_contract.py @@ -6,8 +6,12 @@ from uuid import uuid4 from httpx import AsyncClient +import pytest from sqlalchemy import Executable, Result +from shared.services.retrieval.execution.reference_resolver import ( + resolve_workflow_references, +) from shared.services.retrieval.nav_snapshot import SnapshotSession, load_nav_snapshot from tests.support.retrieval_snapshot_support import contract_db_session from tests.support.contract_database import ContractDatabase @@ -180,3 +184,53 @@ async def publish_new_revision() -> None: assert snapshot.chunk_ref_index[chunks[0].chunk_id]["section_path"] == ( "republished.pdf/old" ) + + +@pytest.mark.asyncio +async def test_reference_hydration_keeps_the_captured_revision_after_republish( + developer_api_client_factory: Callable[ + [], AbstractAsyncContextManager[AsyncClient] + ], +) -> None: + namespace = f"revision-hydration-race-{uuid4().hex[:8]}" + async with developer_api_client_factory(): + document_id, new_result_id = await _seed_republished_document(namespace) + revision_rows = await ContractDatabase.fetch_all( + """ + SELECT job_result_id, chunk_id + FROM document_chunks + WHERE document_id = :document_id + ORDER BY job_result_id + """, + {"document_id": document_id}, + ) + old_revision = next( + row for row in revision_rows if row["job_result_id"] != new_result_id + ) + await ContractDatabase.execute( + """ + UPDATE documents + SET current_job_result_id = :new_result_id + WHERE document_id = :document_id + """, + {"new_result_id": new_result_id, "document_id": document_id}, + ) + + async with contract_db_session() as db: + resolved = await resolve_workflow_references( + db=db, + user_id=_USER_ID, + namespace=namespace, + refs=[ + { + "document_id": document_id, + "chunk_id": old_revision["chunk_id"], + } + ], + revision_pins={document_id: old_revision["job_result_id"]}, + ) + + assert [row["content"] for row in resolved.rows] == ["old content"] + assert [row["job_result_id"] for row in resolved.rows] == [ + old_revision["job_result_id"] + ] diff --git a/apps/api/tests/contract/test_retrieval_term_score_contract.py b/apps/api/tests/contract/test_retrieval_term_score_contract.py new file mode 100644 index 00000000..425f4de0 --- /dev/null +++ b/apps/api/tests/contract/test_retrieval_term_score_contract.py @@ -0,0 +1,75 @@ +"""Contracts for persisted map-unit term scoring.""" + +from __future__ import annotations + +from shared.services.retrieval.nav.nav_knowhere import ReadOnlyChunkStore + + +class _FakeCursor: + def __init__(self, rows: list[tuple[str, str]]) -> None: + self.rows = rows + self.statement = "" + self.parameters: object = None + + def execute(self, statement: str, parameters: object) -> None: + self.statement = statement + self.parameters = parameters + + def fetchall(self) -> list[tuple[str, str]]: + return self.rows + + +def _store() -> ReadOnlyChunkStore: + return ReadOnlyChunkStore.__new__(ReadOnlyChunkStore) + + +def test_term_scores_keep_literal_substring_and_token_hit_semantics() -> None: + cursor = _FakeCursor( + [ + ("unit-full", "prefix alpha beta suffix"), + ("unit-token", "prefix alpha gamma suffix"), + ] + ) + queries = ["alpha beta", "", "alpha"] + query_tokens = { + "alpha beta": ["alpha", "beta"], + "": [], + "alpha": ["alpha"], + } + + scores = _store()._load_term_scores( + cursor, # type: ignore[arg-type] + map_unit_ids=["unit-full", "unit-token", "unit-miss"], + queries=queries, + query_tokens_by_query=query_tokens, + ) + + assert scores == { + "unit-full": (100.0, 0.0, 100.0), + "unit-token": (1.0, 0.0, 100.0), + } + assert "LIKE ANY" in cursor.statement + assert "POSITION" not in cursor.statement + parameters = cursor.parameters + assert isinstance(parameters, tuple) + assert parameters[0] == ["unit-full", "unit-token", "unit-miss"] + assert parameters[1] == ["%alpha beta%", "%alpha%", "%beta%"] + + +def test_long_query_uses_constant_shape_candidate_sql() -> None: + cursor = _FakeCursor([]) + tokens = [f"token-{index}" for index in range(300)] + query = " ".join(tokens) + + _store()._load_term_scores( + cursor, # type: ignore[arg-type] + map_unit_ids=["unit-1"], + queries=[query], + query_tokens_by_query={query: tokens}, + ) + + assert cursor.statement.count("POSITION") == 0 + assert "LIKE ANY" not in cursor.statement + parameters = cursor.parameters + assert isinstance(parameters, tuple) + assert parameters == (["unit-1"],) diff --git a/docs/adr/0005-stream-retrieval-progress-over-sse.md b/docs/adr/0005-stream-retrieval-progress-over-sse.md new file mode 100644 index 00000000..dfaafbd8 --- /dev/null +++ b/docs/adr/0005-stream-retrieval-progress-over-sse.md @@ -0,0 +1,47 @@ +# Stream Retrieval Progress Over SSE + +## Status + +Accepted + +## Context + +Online Brain users currently wait for the complete Retrieval response while +map-nav planning, searching, and source review run. Knowhere owns retrieval and +evidence, while answer generation belongs to downstream clients. A streaming +contract must improve perceived latency without exposing chain-of-thought or +making partial citations authoritative. + +## Decision + +Add a v2-only `POST /v2/retrieval/query/stream` endpoint using Server-Sent +Events. The endpoint accepts the full `RetrievalQueryRequest` shape and is +live-only: disconnecting cooperatively cancels the retrieval run, and retries +start a new run. In-progress events use a fixed, route-aware phase vocabulary +(`started`, `planning`, `searching`, `reviewing_sources`, `finalizing`) and may +include only safe aggregate counts. The stream terminates with a versioned +envelope containing either the existing retrieval response as the authoritative +result or a typed, user-safe failure (`failed`, `cancelled`, or `no_results`). + +The synchronous map-nav engine remains intact and publishes sanitized progress +through an optional step callback bridged to the SSE route by an async queue. +Answer-token streaming remains downstream. The endpoint sends heartbeats, +disables proxy buffering, and uses per-connection event IDs without promising +replay in the first version. + +Retrieval duration is measured from the execution plan's start through final +public-result assembly. The same duration definition is used for persisted +`retrieval_runs.latency_ms` and latency aggregates; cache hits are recorded too. +Trace persistence must receive the execution start timestamp rather than +starting its own timer after navigation and hydration. Timing and terminal +outcomes must cover classic, map-nav, small-corpus, cache-hit, failed, and +cancelled routes. `retrieval_runs` is the ledger for all of those routes and +stores explicit route, cache, latency, and terminal-status fields. + +## Consequences + +The existing JSON retrieval endpoint remains backward compatible, while SDKs +and Online Brain clients need a new streaming adapter and UI state model. A +future resumable stream would require durable event replay and is deliberately +out of scope. Partial evidence and provisional citations are also deferred +until their grounding and revision semantics are defined. diff --git a/docs/adr/0006-atomically-publish-retrieval-serving-index.md b/docs/adr/0006-atomically-publish-retrieval-serving-index.md new file mode 100644 index 00000000..6f0b0800 --- /dev/null +++ b/docs/adr/0006-atomically-publish-retrieval-serving-index.md @@ -0,0 +1,6 @@ +# Atomically publish the retrieval-serving index + +- Status: Accepted +- Context: Retrieval will use a persistent derived serving index to avoid rebuilding a large namespace on every first request. A document revision without a complete index would have unpredictable latency and could produce inconsistent scoring metadata. +- Decision: Build the serving manifest and scoring statistics in the same database transaction as the document revision. Write the completeness marker last. If serving-index construction fails, roll back the publication and retry the job; do not expose an active revision with a partial serving index. +- Consequences: Active revisions have a simple completeness invariant and predictable first-request behavior. Publication takes more work and storage, and an index failure can delay publication, but retrieval can retain a guarded legacy fallback for migrations or already-existing incomplete revisions. diff --git a/docs/adr/0007-use-coherent-retrieval-serving-generations.md b/docs/adr/0007-use-coherent-retrieval-serving-generations.md new file mode 100644 index 00000000..89d2ca05 --- /dev/null +++ b/docs/adr/0007-use-coherent-retrieval-serving-generations.md @@ -0,0 +1,6 @@ +# Use coherent retrieval-serving generations + +- Status: Accepted +- Context: A namespace can contain many active document revisions, and publication can replace them while a retrieval request is loading serving metadata and scoring statistics. +- Decision: Assign each namespace a serving generation. Retrieval captures one generation and verifies it across serving reads; if it changes, retry once and use the exact legacy path if consistency cannot be established. +- Consequences: Retrieval never combines incompatible revision metadata and scoring statistics. Publication and retrieval need a small amount of generation bookkeeping, and rare concurrent updates may cause a retry or slower fallback. diff --git a/docs/adr/0008-use-a-maintenance-window-for-serving-index-rollout.md b/docs/adr/0008-use-a-maintenance-window-for-serving-index-rollout.md new file mode 100644 index 00000000..e7db06bf --- /dev/null +++ b/docs/adr/0008-use-a-maintenance-window-for-serving-index-rollout.md @@ -0,0 +1,6 @@ +# Roll out the serving index online + +- Status: Accepted +- Context: The server and document publication must remain available while serving-index schema changes and backfill are introduced. +- Decision: Use additive online migrations and bounded idempotent backfill. The retrieval reader automatically uses the serving index only when a revision is complete and consistent; otherwise it uses the exact legacy reader. New publication continues online and builds serving data atomically before activating a revision. Do not expose partial serving data. +- Consequences: There is no planned retrieval or publication downtime. Backfill consumes bounded database resources and some revisions remain on the slower legacy path until complete. Generation checks and stale-revision guards are required while backfill and publication run concurrently. diff --git a/docs/adr/README.md b/docs/adr/README.md index 68e15b22..00d7b9da 100644 --- a/docs/adr/README.md +++ b/docs/adr/README.md @@ -19,4 +19,7 @@ Use this shape: | [0002](0002-use-typed-workflow-outcomes.md) | Use typed workflow outcomes | | [0003](0003-keep-retrieval-workflow-policy-explicit.md) | Keep retrieval workflow policy explicit | | [0004](0004-anonymous-self-hosted-telemetry.md) | Anonymous self-hosted telemetry | -| \ No newline at end of file +| [0005](0005-stream-retrieval-progress-over-sse.md) | Stream retrieval progress over SSE | +| [0006](0006-atomically-publish-retrieval-serving-index.md) | Atomically publish the retrieval-serving index | +| [0007](0007-use-coherent-retrieval-serving-generations.md) | Use coherent retrieval-serving generations | +| [0008](0008-use-a-maintenance-window-for-serving-index-rollout.md) | Roll out the serving index online | diff --git a/docs/design/retrieval-serving-index-plan.md b/docs/design/retrieval-serving-index-plan.md new file mode 100644 index 00000000..5e66f6d6 --- /dev/null +++ b/docs/design/retrieval-serving-index-plan.md @@ -0,0 +1,425 @@ +# Retrieval-serving index: online performance plan + +**Status:** Proposed for review +**Reviewed against:** current `knowhere` retrieval and publication code, 2026-08-29 +**Scope:** first-request retrieval performance; LLM planner/harvest/control time is excluded + +## 1. Goal and non-goals + +The target is predictable, bounded **Retrieval Non-LLM Work** on the current +production-sized namespace (about 643 active documents, 50k sections, and 60k +chunks). The measured baseline is roughly 160-170 seconds before map-nav's LLM +episode begins. We do not require a one-second absolute target in this phase; +we require that the request path avoid repeated full-corpus work and have a +clear linear/bounded complexity profile. + +The target includes: + +- snapshot or serving-index loading; +- classic or map-nav lexical scoring; +- ranking; +- selected-result hydration; +- citation and asset-reference assembly. + +It excludes planner, harvest, control, and answer-generation model time. Those +stages must continue to have separate timings. + +The plan does not change prompts, models, tokenization, BM25 formulas, RRF +weights, cache semantics, citation rules, or public HTTP response shapes. +It does not add LLM-response caching, process-wide serving caches, or startup +prewarming. The acceptance benchmark is a cold, uncached retrieval request, +and planner/harvest/control model time remains a separately reported +dependency. Episode-local reuse is allowed only while one request is active. + +## 2. Verified current behavior + +The current code has two retrieval routes in +`shared/services/retrieval/execution/routes.py`: + +- `use_agentic=false` runs `bottom_discovery()` and then ranking/hydration. +- The default map-nav route calls `load_nav_snapshot(..., lazy=True)`, runs the + synchronous navigation episode, then opens a fresh database context for + reference resolution and result assembly. + +`load_nav_snapshot()` currently: + +1. Reads active documents and their current job-result IDs. +2. Loads all matching sections into memory. +3. Loads all chunk identities and `connect_to` metadata into memory. +4. Uses a lazy store for selected chunk content and asset paths. + +The current persisted map scorer in +`shared/services/retrieval/nav/nav_knowhere.py` still loads all eligible map +units, then reads query frequencies and term scores from the persisted tables. +The scorer itself is fast; the broad database projection is not. + +`bottom_discovery()` currently executes path, content, and term channels +sequentially. `term_channel()` uses substring predicates over lowercased text, +without a trigram index. + +Publication currently writes sections/chunks and then calls +`replace_document_map_units()` in the same SQLAlchemy transaction. The existing +`document_map_unit_indexes` row is a per-document-revision completeness marker. +The existing backfill script only rebuilds that map-unit index; it does not yet +build the proposed serving manifest or namespace statistics. + +The working tree already contains an uncommitted keyset-pagination change and +the `3f4a5b6c7d8e` section-order migration. Keep those changes separate from +the serving-index implementation and from unrelated documentation edits. + +## 3. Architecture decision + +Use PostgreSQL as the source of truth and add a persistent, revision-pinned +serving read model. Do not add OpenSearch, Elasticsearch, Tantivy, or another +search service in this version. PostgreSQL's existing FTS plus `pg_trgm` keeps +the current scoring and tie-breaking behavior directly testable. + +### 3.1 Revision serving manifest + +Add one compressed manifest row per `(document_id, job_result_id)`. The payload +contains ordered metadata only: + +- document, revision, source filename, and job identity; +- section IDs, parent IDs, paths, titles, levels, summaries, and sort order; +- chunk IDs, section IDs, types, sort order, and `connect_to` target IDs; +- map-unit row IDs, unit IDs, unit kinds, token lengths, and sort order; +- root-asset IDs and remounted asset owners. + +It must not contain full chunk content or asset file paths. Those remain in the +canonical tables and are loaded lazily for selected evidence. + +Store the payload as canonical JSON compressed with the standard-library zlib +implementation, with a format version and checksum. The serving loader must +reject an unknown version, checksum mismatch, or incomplete payload. + +### 3.2 Serving generations and statistics + +Add namespace-scoped generation metadata and persistent scoring statistics: + +- `retrieval_namespace_generations`: current generation per user/namespace; +- `retrieval_serving_revision_stats`: compressed per-revision contributions; +- `retrieval_namespace_stats`: aggregate unit counts, total lengths, vocabulary + frequency histograms, and generation; +- `retrieval_namespace_token_stats`: queryable document frequency per channel + and token hash. + +The generation is a consistency marker, not a replacement for revision IDs. +Retrieval captures active revision IDs and one generation. If generation changes +while the snapshot is being captured or before scoring starts, retry once; if +consistency still cannot be proven, use the exact legacy reader. + +Every retrieval route must carry that capture as an immutable revision pin set: +`{document_id, namespace, job_result_id}` plus the captured generation. The pin +set is the source of truth for the request after capture. Downstream queries +must constrain sections, map units, chunks, connected assets, ranking lookups, +reference resolution, and result assembly by the pinned `job_result_id`; they +must not re-join through the live `Document.current_job_result_id`. A generation +change after scoring has started must never cause a mix of old and new rows: +finish against the captured pins (or return an exact legacy result), and only +retry before work that depends on the snapshot begins. +Snapshot admission is therefore decided at capture time. If a later archive +must suppress an in-flight result, discard/retry the whole request; do not +replace its pinned revision with the document's new current revision. + +Cache hits occur before route execution, so every operation that changes the +serving generation (publication, republish, archive, or namespace move) must +also advance the namespace retrieval-cache version, or store the generation in +the cache entry and reject mismatches. This keeps cached responses from +outliving the generation they represent without changing the public response +shape. + +### 3.3 Indexes + +Additive migrations should provide: + +- a token-first covering index for map-unit token candidates. The existing + `idx_document_map_unit_tokens_lookup` is token-first but does not cover the + selected columns; the pending `2e3f4a5b6c7d` migration adds a unit-first + covering index for a different access pattern, so the serving reader may + need one additional token-first covering index; +- a revision/section lookup index for map units; +- a trigram GIN index on `document_map_units.term_search_text_lower`; +- a generated lowercased term field and trigram GIN index for + `document_chunks.term_search_text`; +- the existing chunk ordering index plus the pending token-covering and + section-order migrations (`2e3f4a5b6c7d` and `3f4a5b6c7d8e`). + +Enable PostgreSQL's built-in `pg_trgm` extension. No separate search service is +required. + +## 4. Retrieval changes + +Capture the revision pin set and generation at retrieval-route entry, before +the small-corpus count or route selection. Pass that capture into whichever +route is selected; a route-local capture is allowed only when it is performed +as the same snapshot transaction. This prevents the count/load pair in the +small-corpus optimization from straddling a publication. + +### 4.1 Fast map-nav snapshot + +Extend `load_nav_snapshot()` to try the serving manifest first: + +1. Capture active documents, current revision IDs, and namespace generation in a + short read-only transaction, returning the immutable revision pin set with + the snapshot. +2. Fetch one manifest row per active revision. +3. Decode and validate manifests. +4. Apply the existing document and section exclusion predicates. +5. Build the current `LazyKnowhereProvider` and `LazyChunkRefIndex` from the + decoded metadata. +6. Pin the lazy chunk store to the captured revision IDs. +7. Verify generation stability before returning the snapshot. + +For an unfiltered map-nav request, route selection may count chunks from these +same validated manifests instead of scanning `document_chunks`; filtered and +classic requests retain the exact SQL counter. The count shortcut must fall +back when any manifest is missing or invalid. + +If any manifest is absent or invalid, use the existing legacy snapshot loader. +The legacy loader must return the same revision pin set and apply the same +downstream predicates. This fallback is automatic and exact; it is not a public +feature flag. + +The serving path must preserve current ordering, duplicate bare/document-scoped +reference keys, root-asset remounting, section filtering, and revision pinning. +Keep the pin set available through the complete map-nav request. After the LLM +episode, either materialize selected rows (including connected assets) from the +pinned lazy store before closing it, or pass the pin set to +`resolve_workflow_references()` and `assemble_retrieval_results()`. Their SQL +must select the captured `(document_id, job_result_id)` rows directly, so a +republish during the episode cannot make final citations resolve against the +new current revision. + +### 4.2 Exact persisted map scoring + +Keep `PersistedScoreCorpus` and the existing scorer unchanged wherever possible. +Replace only the data-loading strategy: + +- Prepare the immutable, revision-pinned unit projection and namespace scoring + statistics once per retrieval episode. Checklist relight waves must reuse that + projection; they may fetch or compute only query-specific postings/scores. + A wave must not issue another full-namespace unit/statistics load for the same + pin set. Instrument the loader call count and include it in the benchmark + report so repeated projection loads cannot hide behind separate wave timings. +- use manifest map-unit metadata to represent all units, including zero-score + units; the serving reader should not re-query `document_map_units` for these + IDs, lengths, or section membership; +- query token postings only for tokens in the request; +- filter postings by captured revisions and allowed sections; +- discover term-channel candidates through the trigram index while retaining the + current exact substring/token-hit scoring. The candidate predicate must use + the trigram-indexed `LIKE '%term%'` form (with the same lowercased query and + tokens), then apply the existing exact score expression; do not scan every + unit's term text in Python. +- obtain normal-corpus lengths, document frequencies, and IDF-flooring data from + persistent statistics; +- preserve the existing lexical sort key and RRF ranking. + +Queries with document or section exclusions must remain exact. If adjusted +statistics cannot be calculated with certainty, use the legacy scorer for that +request rather than approximating them. + +### 4.3 Classic retrieval — one pinned revision snapshot + +Keep the existing channel implementations and result projection. In +`bottom_discovery()`: + +- capture one revision pin set and generation before starting any channel; +- execute enabled channels concurrently; +- give each channel its own short-lived database session; +- pass the same pin set to every channel and constrain every channel query to + those revisions; +- preserve channel limits, Python BM25, term scoring, RRF merge, score + normalization, and all-or-error behavior; +- use the new trigram index only to narrow term candidates. + +Do not share one `AsyncSession` across concurrent channel tasks. +Ranking lookups, duplicate suppression, connected-target hydration, and final +assembly must receive the same pin set as discovery. The classic result must +therefore contain rows from one revision per document even if publication +replaces a document while one of the channel sessions is running. The +small-corpus optimization must use this same captured snapshot/pin contract (or +the exact legacy equivalent), rather than loading all rows through live current +revision joins. + +## 5. Publication and lifecycle behavior + +Refactor publication so the same build pass produces: + +- canonical sections/chunks; +- existing map-unit rows and completeness marker; +- the revision serving manifest; +- revision statistics and namespace-statistics deltas. + +All of this happens synchronously in the existing publication transaction. The +completeness marker and generation update are written last. If serving-index +construction fails, the publication transaction rolls back. + +New publication remains online during backfill. First publication, republish, +archive, and namespace-move paths must update statistics under the same +namespace generation row lock. The lock covers the active revision set, +namespace membership, revision contributions, and the generation increment, so +readers and writers have one lifecycle ordering. + +Backfill must rebuild the complete derived serving state (map units, manifest, +revision contribution, and namespace-statistics delta), not only the existing +map-unit index. It must select only documents with `status = 'active'`, a non-null +`current_job_result_id`, and the intended user/namespace. Immediately before +writing a contribution, it must hold the namespace lock and re-read the +document, then require all of the following to remain true: active status, +unchanged user/namespace, and `current_job_result_id` equal to the captured +revision. Otherwise it skips that revision without adding statistics. This +active-status predicate is required in the selector as well as in the +commit-time guard; update `apps/api/scripts/backfill_map_unit_indexes.py` to +include it in the existing selector. + +Archiving must atomically remove or invalidate that document revision's serving +statistics contribution while holding the same lock and advance the namespace +generation. `archive` currently changes `status` without clearing +`current_job_result_id`, so checking the revision pointer alone is insufficient +and would allow an in-flight backfill to re-add an archived revision. + +## 6. Online rollout + +There is no planned downtime, runtime feature flag, or production shadow-read +mode. + +1. Deploy additive schema/index migrations, beginning with the pending + `2e3f4a5b6c7d` and `3f4a5b6c7d8e` migrations. +2. Deploy code that automatically uses the serving reader only for complete, + valid revisions and otherwise uses the legacy reader. +3. Run an explicit, idempotent, bounded backfill for existing active revisions. +4. Keep retrieval and publication online while backfill runs. +5. Verify manifest checksums, revision coverage, namespace statistics, and + generation consistency. +6. Run strict legacy-versus-serving differential checks before considering the + rollout complete. + +If backfill is incomplete, affected revisions continue on the exact legacy +path. If online serving data is corrupted, reject it, alert, repair it with the +backfill/rebuild script, and do not serve partial data. + +Before enabling the serving reader for a namespace, record an inventory of +active `(document_id, current_job_result_id)` pairs, manifest completeness, and +expected per-revision and aggregate unit counts. After backfill, reconcile those +same values and verify that every aggregate includes only active, namespace- +member revisions. Abort the fast-path rollout on any missing/extra revision, +checksum failure, count mismatch, archived contribution, or generation +discontinuity. + +Any migration, backfill, or other database write—especially against +production—requires explicit approval immediately before execution. Read-only +inspection and benchmarking may proceed without that approval. + +## 6.1 DevOps operations runbook + +DevOps owns the production rollout mechanics; application code does not run a +startup backfill or create serving tables implicitly. Execute the following in +order: + +1. **Preflight (read-only):** confirm the target account, database, migration + head, available disk, connection headroom, and a recent rollback point. Record + the active `(document_id, current_job_result_id)` inventory for each namespace + that will be backfilled. +2. **Schema rollout:** with explicit approval immediately beforehand, apply the + additive migrations in dependency order: `2e3f4a5b6c7d`, + `3f4a5b6c7d8e`, `4a5b6c7d8e9f`, then `5b6c7d8e9f0a`. Run the trigram-index + migration during a low-traffic window and monitor for blocking locks. +3. **Application rollout:** deploy the API and worker versions containing the + serving reader and atomic publication changes. Verify health, error rate, and + legacy fallback before starting the backfill. +4. **Bounded backfill:** with separate approval, run + `uv run python apps/api/scripts/backfill_map_unit_indexes.py --apply` from a + controlled operator environment. Limit concurrency, pause on database + saturation, and resume safely; the operation is idempotent and stale or + inactive revisions must be skipped. +5. **Reconciliation:** compare the preflight inventory with serving manifests, + checksums, per-revision unit counts, namespace aggregates, and generation + values. Confirm aggregates contain only active documents still belonging to + the namespace. Investigate every missing, extra, stale, or invalid revision. +6. **Acceptance:** run the production read-only legacy-versus-serving + differential harness and record latency, selected IDs, order, scores, + citations, section paths, asset references, and fallback behavior. Declare + the rollout complete only after zero semantic mismatches. + +If migration or backfill must be stopped, leave the serving tables in place and +stop the operator job. The reader will continue using the exact legacy path for +incomplete revisions. Roll back application code first if necessary; do not +drop serving tables or indexes as an emergency rollback action. Repair a failed +revision by rerunning the bounded backfill after the cause is understood. + +## 7. Contract tests and benchmarks + +Use contract tests only. Add contracts for: + +- manifest round-trip, checksum, version, and revision pinning; +- eager versus serving snapshot equivalence; +- exclusions, duplicate chunk IDs, document-scoped references, and root assets; +- exact Latin/CJK, empty, no-hit, phrase, token-only, and negative-IDF cases; +- incomplete serving data falling back to legacy; +- publication replacement, archive deltas, concurrent generation changes, and + stale backfill protection; +- cache invalidation racing with a generation change, proving an old cached + response is not returned for a newer serving generation; +- map-nav republish during the LLM episode, proving final hydration and + connected-asset resolution stay on the captured revisions; +- classic publication replacement during concurrent channels, ranking, and + final assembly, proving every returned row shares the channel's pin set; +- archive/backfill races proving archived or namespace-moved revisions never + contribute to serving statistics; +- concurrent classic channels preserving IDs, order, scores, citations, and + fallback behavior. + +Race tests must use barriers or an equivalent deterministic hook to force a +republish during the map-nav episode, a publication between classic channel +sessions, an archive during backfill, and a namespace move during backfill. +Each test must assert both the returned evidence and the persisted statistics, +not merely that the request completed. + +The validation harness must also inspect the generated SQL/query plans (or an +equivalent query-boundary assertion) to prove pinned reads do not use live +`Document.current_job_result_id` joins. Run cache-version/generation races and +verify an old cached response is rejected after a lifecycle change. + +Run a differential harness against the production read-only database and +compare selected IDs, ordering, rounded scores, citations, section paths, +asset references, and fallback behavior. + +Benchmark fresh processes and uncached queries. Report separately: + +- serving capture/decode; +- map index projection and scoring; +- episode-local projection reuse (number of full projection loads and per-wave + query-only scoring time); +- classic discovery; +- ranking; +- hydration/assembly; +- total Retrieval Non-LLM Work; +- planner/harvest/control LLM time. + +The complexity check is explicit: one request may perform one full pinned +snapshot/projection pass, relight work should scale with query postings rather +than reloading the corpus, and hydration should scale with selected evidence +(`top_k`/references), not namespace size. Navigation wave count must not +multiply full-corpus database loads. + +Record peak resident memory for a fresh worker during the same benchmark and +repeat it with the expected concurrent-request level. Memory is reported as an +operational trade-off rather than a latency acceptance gate for this phase; +before production rollout, any episode-local or process-local reuse still needs +an explicit byte/item budget and an agreed worker ceiling. + +The fast path is accepted only after zero semantic mismatches and evidence that +the cold request performs one bounded serving projection, does not repeat +full-corpus loads per navigation wave, and meets an agreed latency budget for +the current production-sized corpus. + +## 8. Main tradeoffs and risks + +- Publication becomes slower and uses more storage because derived data is built + synchronously. +- Existing documents need an explicit backfill before they use the fast path. +- A serving-index inconsistency causes a slower legacy request, not approximate + evidence. +- PostgreSQL remains a scaling dependency; a future search-engine migration + would require a new semantic-parity review. diff --git a/docs/design/retrieval-streaming-sse.md b/docs/design/retrieval-streaming-sse.md new file mode 100644 index 00000000..b0b5a73e --- /dev/null +++ b/docs/design/retrieval-streaming-sse.md @@ -0,0 +1,207 @@ +# Retrieval Streaming over SSE + +**Status:** Accepted design +**Related issue:** [#330](https://github.com/Ontos-AI/knowhere/issues/330) +**Related ADR:** [0005](../adr/0005-stream-retrieval-progress-over-sse.md) + +## Purpose + +Online Brain users currently wait for a complete retrieval response while +map-nav planning, searching, source review, and final hydration run. This +design makes that work visible without exposing chain-of-thought or changing +who owns answer generation. + +## Ownership boundary + +Knowhere owns retrieval, safe progress, evidence, and authoritative citations. +The downstream Online Brain client owns answer synthesis and answer-token +streaming. Knowhere does not begin generating the final answer as part of this +feature. + +## Public API + +Add `POST /v2/retrieval/query/stream` with `Content-Type: text/event-stream`. +The request body is the full existing +`RetrievalQueryRequest`; streaming changes delivery, not retrieval semantics. +The existing JSON endpoints remain unchanged as the fallback. + +The stream is live-only. A disconnect cooperatively cancels the run; retrying +starts a new run. Event IDs are monotonic per connection but do not imply +replay or resumability. + +The route remains behind the normal authenticated-user dependency and v2 +route-admission policy. The guest-key allowlist, system-limit configuration, +OpenAPI registration, and any CORS policy must explicitly include the stream +route. v2 BYOK `llm_config` is accepted exactly as on the JSON route and is +never copied into an event payload or log message. + +## Event contract + +Each SSE frame has an event name (`progress`, `heartbeat`, or `terminal`) and a +JSON payload with `schema_version`, `stream_id`, `sequence`, and server +`elapsed_ms`. Progress payloads use the fixed, route-aware phases: + +```text +started → planning? → searching → reviewing_sources → finalizing +``` + +Classic and small-corpus routes omit `planning`; they must not emit phases that +did not occur. In-progress events contain only safe aggregate counts such as +`candidate_source_count` and `reviewed_source_count`. They do not contain +document names, chunk content, query rewrites, raw planner output, citations, +or chain-of-thought. + +The terminal event is a versioned envelope containing the existing retrieval +response as the authoritative result: + +```json +{ + "schema_version": 1, + "stream_id": "rst_...", + "sequence": 7, + "elapsed_ms": 2410, + "status": "completed", + "response": { "namespace": "default", "query": "...", "router_used": "mapnav", "evidence_text": "...", "referenced_chunks": [], "results": [] } +} +``` + +Failure terminals use `failed`, `cancelled`, or `no_results`. Failures expose a +stable user-safe `code` and `message`; detailed provider, database, +authentication, and planner errors remain in server logs. + +Wire requirements are part of the contract: frames use UTF-8 SSE `id`, +`event`, and one `data` line followed by a blank line; a heartbeat is an SSE +comment or named heartbeat event and carries no retrieval data; exactly one +terminal event is sent before the connection closes; no `retry` directive is +promised because the stream is not resumable. Progress events may be +coalesced when a bounded queue is full, but terminal events must never be +dropped. + +Cache hits still emit `started` and a terminal `completed` event, with a safe +`cache_hit: true` indicator. They omit phases for work that did not run. +`no_results` is reserved for a successful retrieval execution that produces no +evidence; provider, timeout, validation, and internal failures remain +`failed`. + +HTTP authentication, request validation, and route-admission failures happen +before the stream starts and use ordinary HTTP error responses. Once a `200` +SSE response has been opened, execution failures must be represented by a +terminal SSE event because the HTTP status can no longer be changed. + +## Internal implementation seam + +Keep the synchronous map-nav implementation. Add an optional callback that +receives a sanitized progress projection after each completed planner, +search, or review step. The SSE route bridges this callback to an +`asyncio.Queue` using a thread-safe loop handoff while retrieval continues in +its existing worker thread. + +The queue is bounded and owns cleanup of the worker task, callback, heartbeat +task, and database session. The callback must not touch an `AsyncSession` from +the worker thread. The route polls request disconnect state and propagates a +cancellation token; all producer tasks are joined or cancelled in a `finally` +block so abandoned streams cannot leak threads or connections. + +Add cooperative cancellation checks between steps. An in-flight synchronous +provider call may finish before cancellation takes effect. Cancelled runs do +not perform final hydration when cancellation is observed in time. + +Phase ownership is explicit: the route emits `started`; the map-nav adapter +emits `planning` before `plan_query` and `searching` before navigation or +classic discovery; the route emits `reviewing_sources` after retrieval +selection and before reference hydration; and it emits `finalizing` before +public projection. Counts are sourced from existing snapshot, reference, and +assembled-result counts and are omitted when not yet known. + +## Correct duration accounting + +The execution plan already starts a monotonic timer before cache lookup and +logs elapsed time after the route. However, `TraceRecorder` currently starts a +second timer in its constructor, and the map-nav route constructs it only +after navigation, reference resolution, and result assembly. Persisted +`retrieval_runs.latency_ms` and its aggregates therefore under-report retrieval +latency and mostly measure trace flush time. + +The canonical `Retrieval Duration` is the execution-plan timer from retrieval +start through final public-result assembly. It includes cache lookup and +applies to cache hits and misses. It excludes authentication, network/SSE +delivery, and downstream answer generation. + +Required changes: + +- pass the execution start timestamp into `TraceRecorder`; +- set `retrieval_runs.latency_ms` from that timestamp; +- record cache-hit runs with the same definition; +- ensure classic, map-nav, small-corpus, cache-hit, failed, and cancelled + retrievals all have an explicit timing/observability outcome; +- expose separate `time_to_first_event_ms`, `retrieval_latency_ms`, and + downstream `time_to_first_token_ms` measurements; +- retain per-step `elapsed_ms` as step latency, not total request latency. + +`retrieval_runs` is the ledger for every retrieval execution, not only +map-nav. Each row records the route type, `agentic_enabled`, `cache_hit`, +canonical latency, and terminal status for classic, map-nav, small-corpus, +cache-hit, failed, and cancelled runs. Add a backward-compatible status field +and migration rather than overloading free-form error text. + +The execution timer must end after public response projection, not merely when +the internal route outcome is assembled. If trace persistence is best-effort, +its latency update must still use the captured execution timestamps and must +not extend the user-visible retrieval duration with an unbounded database +flush. + +Because the map-nav route deliberately rolls back its request session before +the synchronous LLM episode, a trace row must not be created in that session +before the rollback. Capture the execution start immediately, then create or +update the trace record at the terminal persistence point with the captured +start and an explicit end/duration supplied by the execution plan (or use a +separate trace session). `TraceRecorder.complete()` must not silently choose a +later local constructor time or include its own flush duration. + +## Operational requirements + +- heartbeat every 15 seconds while active; +- `Cache-Control: no-cache` and `X-Accel-Buffering: no`; +- flush after every event; +- use `fetch`-style clients where POST authorization headers are required; +- propagate disconnects to the cancellation token. +- bound maximum stream lifetime and enforce the same request/rate-limit policy + as the JSON endpoint; +- document worker-thread, database-connection, and concurrent-stream limits. +- count one stream request as one retrieval request under the existing user and + system limits; retries count as new requests and cannot bypass quota. + +## Delivery slices + +1. **Knowhere contract and timing:** typed event models, stream route, callback + bridge, cancellation, route admission, cache semantics, corrected + `RetrievalRun` timing across all route types, and API contract tests. +2. **SDK adapters:** typed Python and Node stream consumers with fallback to + the existing JSON query; parse named SSE events, expose abort/error + handling, and preserve the full terminal response shape. +3. **Online Brain UX:** phase state model, progress view, safe error states, + and integration with existing downstream answer-token streaming. +4. **Production verification:** proxy-path e2e tests and latency/buffering + instrumentation before setting hard p50/p95 targets. + +## Verification gates + +- correct phase order for map-nav, classic, and small-corpus routes; +- cache-hit streams emit only applicable phases and identify the cache hit; +- no sensitive planner or evidence data before the terminal event; +- terminal citations match the existing JSON endpoint; +- cancellation, timeout, no-results, provider failure, and disconnect are + distinguishable; +- persisted latency includes the full retrieval path for every supported route + and matches API timing within an expected tolerance; +- cache-hit latency and route type are visible in observability data; +- heartbeats pass through the deployed proxy without buffering; +- first-event and first-token times are measured separately. + +## Explicit non-goals + +- v1 API streaming endpoint; +- resumable/replayed streams; +- partial authoritative evidence or provisional citation revision; +- moving answer generation into Knowhere; +- invented latency SLAs before a production baseline exists. diff --git a/packages/shared-python/shared/models/database/document.py b/packages/shared-python/shared/models/database/document.py index ad59857e..c396283d 100644 --- a/packages/shared-python/shared/models/database/document.py +++ b/packages/shared-python/shared/models/database/document.py @@ -12,8 +12,10 @@ DateTime, Float, ForeignKey, + BigInteger, Index, Integer, + LargeBinary, String, Text, UniqueConstraint, @@ -130,6 +132,13 @@ class DocumentSection(Base): ), Index("idx_document_sections_scope", "user_id", "namespace"), Index("idx_document_sections_doc_revision", "document_id", "job_result_id"), + Index( + "idx_document_sections_revision_snapshot_order", + "document_id", + "job_result_id", + "sort_order", + "section_id", + ), ) @@ -347,6 +356,168 @@ class DocumentMapUnitIndex(Base): ) +class RetrievalNamespaceGeneration(Base): + """Monotonic serving generation for one user-owned namespace.""" + + __tablename__ = "retrieval_namespace_generations" + + id: Mapped[str] = mapped_column( + String(100), primary_key=True, default=lambda: f"rng_{uuid4().hex}" + ) + user_id: Mapped[str] = mapped_column(Text, nullable=False) + namespace: Mapped[str] = mapped_column(String(255), nullable=False) + generation: Mapped[int] = mapped_column(BigInteger, nullable=False, default=0) + updated_at: Mapped[datetime] = mapped_column( + DateTime, default=utc_now_naive, onupdate=utc_now_naive, nullable=False + ) + + __table_args__ = ( + UniqueConstraint( + "user_id", + "namespace", + name="uq_retrieval_namespace_generations_scope", + ), + ) + + +class RetrievalServingRevisionManifest(Base): + """Compressed ordered metadata for one document revision.""" + + __tablename__ = "retrieval_serving_revision_manifests" + + id: Mapped[str] = mapped_column( + String(100), primary_key=True, default=lambda: f"rsm_{uuid4().hex}" + ) + user_id: Mapped[str] = mapped_column(Text, nullable=False) + namespace: Mapped[str] = mapped_column(String(255), nullable=False) + document_id: Mapped[str] = mapped_column( + String(36), ForeignKey("documents.document_id", ondelete="CASCADE"), nullable=False + ) + job_result_id: Mapped[str] = mapped_column( + String(36), ForeignKey("job_results.id", ondelete="CASCADE"), nullable=False + ) + format_version: Mapped[int] = mapped_column(Integer, nullable=False) + payload_zlib: Mapped[bytes] = mapped_column(LargeBinary, nullable=False) + checksum: Mapped[str] = mapped_column(String(64), nullable=False) + created_at: Mapped[datetime] = mapped_column( + DateTime, default=utc_now_naive, nullable=False + ) + + __table_args__ = ( + UniqueConstraint( + "document_id", + "job_result_id", + name="uq_retrieval_serving_revision_manifests_revision", + ), + Index( + "idx_retrieval_serving_revision_manifests_scope", + "user_id", + "namespace", + "document_id", + "job_result_id", + ), + ) + + +class RetrievalServingRevisionStat(Base): + """Compressed scoring contribution for one document revision.""" + + __tablename__ = "retrieval_serving_revision_stats" + + id: Mapped[str] = mapped_column( + String(100), primary_key=True, default=lambda: f"rss_{uuid4().hex}" + ) + user_id: Mapped[str] = mapped_column(Text, nullable=False) + namespace: Mapped[str] = mapped_column(String(255), nullable=False) + document_id: Mapped[str] = mapped_column( + String(36), ForeignKey("documents.document_id", ondelete="CASCADE"), nullable=False + ) + job_result_id: Mapped[str] = mapped_column( + String(36), ForeignKey("job_results.id", ondelete="CASCADE"), nullable=False + ) + format_version: Mapped[int] = mapped_column(Integer, nullable=False) + payload_zlib: Mapped[bytes] = mapped_column(LargeBinary, nullable=False) + checksum: Mapped[str] = mapped_column(String(64), nullable=False) + created_at: Mapped[datetime] = mapped_column( + DateTime, default=utc_now_naive, nullable=False + ) + + __table_args__ = ( + UniqueConstraint( + "document_id", + "job_result_id", + name="uq_retrieval_serving_revision_stats_revision", + ), + Index( + "idx_retrieval_serving_revision_stats_scope", + "user_id", + "namespace", + "document_id", + "job_result_id", + ), + ) + + +class RetrievalNamespaceStat(Base): + """Compressed aggregate scoring statistics for one namespace generation.""" + + __tablename__ = "retrieval_namespace_stats" + + id: Mapped[str] = mapped_column( + String(100), primary_key=True, default=lambda: f"rns_{uuid4().hex}" + ) + user_id: Mapped[str] = mapped_column(Text, nullable=False) + namespace: Mapped[str] = mapped_column(String(255), nullable=False) + generation: Mapped[int] = mapped_column(BigInteger, nullable=False) + payload_zlib: Mapped[bytes] = mapped_column(LargeBinary, nullable=False) + checksum: Mapped[str] = mapped_column(String(64), nullable=False) + updated_at: Mapped[datetime] = mapped_column( + DateTime, default=utc_now_naive, onupdate=utc_now_naive, nullable=False + ) + + __table_args__ = ( + UniqueConstraint( + "user_id", + "namespace", + name="uq_retrieval_namespace_stats_scope", + ), + ) + + +class RetrievalNamespaceTokenStat(Base): + """Document frequency for one token/channel in a namespace generation.""" + + __tablename__ = "retrieval_namespace_token_stats" + + id: Mapped[str] = mapped_column( + String(100), primary_key=True, default=lambda: f"rnt_{uuid4().hex}" + ) + user_id: Mapped[str] = mapped_column(Text, nullable=False) + namespace: Mapped[str] = mapped_column(String(255), nullable=False) + generation: Mapped[int] = mapped_column(BigInteger, nullable=False) + channel: Mapped[str] = mapped_column(String(32), nullable=False) + token_hash: Mapped[str] = mapped_column(String(64), nullable=False) + document_frequency: Mapped[int] = mapped_column(Integer, nullable=False) + + __table_args__ = ( + UniqueConstraint( + "user_id", + "namespace", + "channel", + "token_hash", + name="uq_retrieval_namespace_token_stats_key", + ), + Index( + "idx_retrieval_namespace_token_stats_lookup", + "user_id", + "namespace", + "generation", + "channel", + "token_hash", + ), + ) + + class GraphNode(Base): """Persisted derived graph node used for routing and expansion.""" diff --git a/packages/shared-python/shared/services/retrieval/execution/plan.py b/packages/shared-python/shared/services/retrieval/execution/plan.py index 8d05753a..adb2a6da 100644 --- a/packages/shared-python/shared/services/retrieval/execution/plan.py +++ b/packages/shared-python/shared/services/retrieval/execution/plan.py @@ -1,6 +1,7 @@ from __future__ import annotations import time +from dataclasses import replace from typing import Any from loguru import logger @@ -16,6 +17,10 @@ set_cached_retrieval_query_result, ) from shared.services.retrieval.execution.routes import run_retrieval_route +from shared.services.retrieval.execution.revision_pins import ( + capture_revision_pins, + is_revision_generation_stable, +) from shared.services.retrieval.stats.recorder import ( schedule_retrieval_hit_stats_update, ) @@ -142,7 +147,26 @@ async def _execute_with_overrides(self, request: RetrievalQuery) -> dict[str, An logger.debug(f" 📦 Cache miss (version={cache_version}), running full pipeline") - outcome = await run_retrieval_route(request.build_route_context()) + route_context = request.build_route_context() + revision_pins = await capture_revision_pins( + request.db, + user_id=request.user_id, + namespace=request.namespace, + ) + if not await is_revision_generation_stable( + request.db, + user_id=request.user_id, + namespace=request.namespace, + pins=revision_pins, + ): + revision_pins = await capture_revision_pins( + request.db, + user_id=request.user_id, + namespace=request.namespace, + ) + outcome = await run_retrieval_route( + replace(route_context, revision_pins=revision_pins) + ) if cache_version is not None: await _write_cached_response( diff --git a/packages/shared-python/shared/services/retrieval/execution/reference_resolver.py b/packages/shared-python/shared/services/retrieval/execution/reference_resolver.py index 802fa278..459ad273 100644 --- a/packages/shared-python/shared/services/retrieval/execution/reference_resolver.py +++ b/packages/shared-python/shared/services/retrieval/execution/reference_resolver.py @@ -1,6 +1,7 @@ from __future__ import annotations from dataclasses import dataclass +from collections.abc import Mapping from typing import Any from sqlalchemy.ext.asyncio import AsyncSession @@ -28,6 +29,7 @@ async def resolve_workflow_references( namespace: str, refs: list[dict[str, Any]], score_by_chunk_id: dict[str, float] | None = None, + revision_pins: Mapping[str, str] | None = None, ) -> ResolvedWorkflowReferences: hydrated_rows = await hydrate_referenced_chunk_rows( db=db, @@ -35,6 +37,7 @@ async def resolve_workflow_references( namespace=namespace, refs=refs, score_by_chunk_id=score_by_chunk_id, + revision_pins=revision_pins, ) resolved = _select_matching_references(refs, hydrated_rows) enriched_rows = await enrich_referenced_chunks_with_asset_url(resolved.rows) diff --git a/packages/shared-python/shared/services/retrieval/execution/revision_pins.py b/packages/shared-python/shared/services/retrieval/execution/revision_pins.py new file mode 100644 index 00000000..e7e9efcc --- /dev/null +++ b/packages/shared-python/shared/services/retrieval/execution/revision_pins.py @@ -0,0 +1,92 @@ +"""Capture and carry one immutable revision set through a retrieval request.""" + +from __future__ import annotations + +from collections.abc import Iterator, Mapping +from dataclasses import dataclass +from types import MappingProxyType + +from sqlalchemy import select +from sqlalchemy.exc import SQLAlchemyError +from sqlalchemy.ext.asyncio import AsyncSession + +from shared.models.database.document import Document +from shared.models.database.document import RetrievalNamespaceGeneration + + +@dataclass(frozen=True) +class RetrievalRevisionPins(Mapping[str, str]): + """The active document revisions admitted to one retrieval request.""" + + revisions: Mapping[str, str] + generation: int | None = None + + def __post_init__(self) -> None: + object.__setattr__(self, "revisions", MappingProxyType(dict(self.revisions))) + + def __getitem__(self, document_id: str) -> str: + return self.revisions[document_id] + + def __iter__(self) -> Iterator[str]: + return iter(self.revisions) + + def __len__(self) -> int: + return len(self.revisions) + + +async def capture_revision_pins( + db: AsyncSession, + *, + user_id: str, + namespace: str, +) -> RetrievalRevisionPins: + """Capture active document revisions in one database read transaction.""" + statement = ( + select(Document.document_id, Document.current_job_result_id) + .where(Document.user_id == user_id) + .where(Document.namespace == namespace) + .where(Document.status == "active") + .where(Document.current_job_result_id.is_not(None)) + .order_by(Document.document_id) + ) + try: + generation_result = await db.execute( + select(RetrievalNamespaceGeneration.generation) + .where(RetrievalNamespaceGeneration.user_id == user_id) + .where(RetrievalNamespaceGeneration.namespace == namespace) + ) + generation_row = generation_result.scalar_one_or_none() + except SQLAlchemyError: + await db.rollback() + generation_row = None + rows = (await db.execute(statement)).all() + revisions = { + str(document_id): str(job_result_id) + for document_id, job_result_id in rows + if document_id and job_result_id + } + return RetrievalRevisionPins( + revisions=revisions, + generation=int(generation_row) if generation_row is not None else 0, + ) + + +async def is_revision_generation_stable( + db: AsyncSession, + *, + user_id: str, + namespace: str, + pins: RetrievalRevisionPins, +) -> bool: + """Return whether the namespace generation is unchanged since capture.""" + try: + result = await db.execute( + select(RetrievalNamespaceGeneration.generation) + .where(RetrievalNamespaceGeneration.user_id == user_id) + .where(RetrievalNamespaceGeneration.namespace == namespace) + ) + current_generation = result.scalar_one_or_none() + except SQLAlchemyError: + await db.rollback() + return True + return int(current_generation or 0) == int(pins.generation or 0) diff --git a/packages/shared-python/shared/services/retrieval/execution/route_types.py b/packages/shared-python/shared/services/retrieval/execution/route_types.py index 59c173e4..13c21c46 100644 --- a/packages/shared-python/shared/services/retrieval/execution/route_types.py +++ b/packages/shared-python/shared/services/retrieval/execution/route_types.py @@ -5,6 +5,8 @@ from sqlalchemy.ext.asyncio import AsyncSession +from shared.services.retrieval.execution.revision_pins import RetrievalRevisionPins + @dataclass(frozen=True) class RetrievalRouteContext: @@ -26,6 +28,7 @@ class RetrievalRouteContext: internal_recall_k: int | None effective_recall_k: int use_agentic: bool | None + revision_pins: RetrievalRevisionPins | None = None @dataclass(frozen=True) diff --git a/packages/shared-python/shared/services/retrieval/execution/routes.py b/packages/shared-python/shared/services/retrieval/execution/routes.py index 97961074..d2b6d211 100644 --- a/packages/shared-python/shared/services/retrieval/execution/routes.py +++ b/packages/shared-python/shared/services/retrieval/execution/routes.py @@ -8,18 +8,29 @@ from sqlalchemy.ext.asyncio import AsyncSession from shared.services.retrieval.search.discovery import bottom_discovery -from shared.services.retrieval.execution.reference_resolver import resolve_workflow_references -from shared.services.retrieval.hydration.result_assembly import assemble_retrieval_results -from shared.services.retrieval.hydration.legacy_evidence import render_legacy_evidence_text +from shared.services.retrieval.execution.reference_resolver import ( + resolve_workflow_references, +) +from shared.services.retrieval.hydration.result_assembly import ( + assemble_retrieval_results, +) +from shared.services.retrieval.hydration.legacy_evidence import ( + render_legacy_evidence_text, +) from shared.services.retrieval.execution.route_types import ( RetrievalRouteContext, RetrievalRouteOutcome, ) from shared.services.retrieval.search.ranking import rank_retrieval_candidates from shared.services.retrieval.search.scoped_corpus import ( + count_manifest_chunks, count_scoped_chunks, load_all_scoped_chunks, ) +from shared.services.retrieval.execution.revision_pins import ( + capture_revision_pins, + is_revision_generation_stable, +) def open_fresh_database_context() -> AbstractAsyncContextManager[AsyncSession]: @@ -46,13 +57,28 @@ async def run_retrieval_route( async def _try_run_small_corpus_route( context: RetrievalRouteContext, ) -> RetrievalRouteOutcome | None: - total_chunk_count = await count_scoped_chunks( - context.db, - user_id=context.user_id, - namespace=context.namespace, - exclude_document_ids=context.exclude_document_ids, - allowed_chunk_types=context.allowed_chunk_types, - ) + total_chunk_count: int | None = None + if ( + context.use_agentic is not False + and context.revision_pins is not None + and not context.exclude_document_ids + and not context.exclude_sections + and context.allowed_chunk_types is None + and not context.signal_paths + ): + total_chunk_count = await count_manifest_chunks( + context.db, + revision_pins=context.revision_pins, + ) + if total_chunk_count is None: + total_chunk_count = await count_scoped_chunks( + context.db, + user_id=context.user_id, + namespace=context.namespace, + exclude_document_ids=context.exclude_document_ids, + allowed_chunk_types=context.allowed_chunk_types, + revision_pins=context.revision_pins, + ) logger.info(f"\n Total chunks in scope: {total_chunk_count}") if total_chunk_count > context.top_k: @@ -71,6 +97,7 @@ async def _try_run_small_corpus_route( allowed_chunk_types=context.allowed_chunk_types, signal_paths=context.signal_paths or [], filter_mode=context.filter_mode, + revision_pins=context.revision_pins, ) logger.info( f" small_corpus load: loaded={len(all_rows)} rows after signal/exclude filters" @@ -81,6 +108,7 @@ async def _try_run_small_corpus_route( exclude_document_ids=context.exclude_document_ids, exclude_sections=context.exclude_sections, allowed_chunk_types=context.allowed_chunk_types, + revision_pins=context.revision_pins, ) results = assembled_rows response = { @@ -117,6 +145,7 @@ async def _run_classic_topk_route( channels=context.channels, channel_weights=context.channel_weights, internal_recall_k=context.internal_recall_k, + revision_pins=context.revision_pins, ) fused_rows = ( @@ -132,6 +161,7 @@ async def _run_classic_topk_route( discovery_rows=fused_rows, routed_rows=[], top_k=context.top_k, + revision_pins=context.revision_pins, ) assembled_rows = await assemble_retrieval_results( @@ -140,6 +170,7 @@ async def _run_classic_topk_route( exclude_document_ids=context.exclude_document_ids, exclude_sections=context.exclude_sections, allowed_chunk_types=context.allowed_chunk_types, + revision_pins=context.revision_pins, ) results = assembled_rows response = { @@ -181,7 +212,9 @@ async def _run_mapnav_route( episode_token_count, episode_workflow_plan, ) + snapshot_started = time.perf_counter() + snapshot_pins = context.revision_pins snapshot = await load_nav_snapshot( context.db, user_id=context.user_id, @@ -189,7 +222,29 @@ async def _run_mapnav_route( exclude_document_ids=context.exclude_document_ids, exclude_sections=context.exclude_sections, lazy=True, + revision_pins=snapshot_pins, ) + if snapshot_pins is not None and not await is_revision_generation_stable( + context.db, + user_id=context.user_id, + namespace=context.namespace, + pins=snapshot_pins, + ): + snapshot.close() + snapshot_pins = await capture_revision_pins( + context.db, + user_id=context.user_id, + namespace=context.namespace, + ) + snapshot = await load_nav_snapshot( + context.db, + user_id=context.user_id, + namespace=context.namespace, + exclude_document_ids=context.exclude_document_ids, + exclude_sections=context.exclude_sections, + lazy=True, + revision_pins=snapshot_pins, + ) snapshot_seconds = time.perf_counter() - snapshot_started logger.info( "retrieval mapnav stage=snapshot_load seconds={:.3f} documents={} refs={}".format( @@ -239,6 +294,7 @@ async def _run_mapnav_route( namespace=context.namespace, refs=refs, score_by_chunk_id=score_by_chunk_id or None, + revision_pins=snapshot.document_revisions, ) assembled_rows = await assemble_retrieval_results( db=final_db, @@ -246,6 +302,7 @@ async def _run_mapnav_route( exclude_document_ids=context.exclude_document_ids, exclude_sections=context.exclude_sections, allowed_chunk_types=context.allowed_chunk_types, + revision_pins=snapshot.document_revisions, ) decision_steps = build_decision_trace( @@ -299,9 +356,7 @@ async def _run_mapnav_route( "decision_trace": decision_trace, } - completion_detail = ( - f"chunks | evidence={len(evidence_text)} chars | router=mapnav" - ) + completion_detail = f"chunks | evidence={len(evidence_text)} chars | router=mapnav" return RetrievalRouteOutcome( response=response, hit_stats_results=resolved.refs, diff --git a/packages/shared-python/shared/services/retrieval/hydration/connected.py b/packages/shared-python/shared/services/retrieval/hydration/connected.py index fcbc647b..2a1f3127 100644 --- a/packages/shared-python/shared/services/retrieval/hydration/connected.py +++ b/packages/shared-python/shared/services/retrieval/hydration/connected.py @@ -1,5 +1,6 @@ from __future__ import annotations +from collections.abc import Mapping from typing import Any from sqlalchemy import and_, or_, select @@ -20,6 +21,7 @@ async def hydrate_connected_target_rows( rows: list[dict[str, Any]], exclude_document_ids: list[str], exclude_sections: list[dict[str, str]], + revision_pins: Mapping[str, str] | None = None, ) -> list[dict[str, Any]]: if db is None: return [] @@ -61,7 +63,25 @@ async def hydrate_connected_target_rows( stmt = ( select(Document, DocumentChunk, DocumentSection, JobResult) - .join(DocumentChunk, DocumentChunk.document_id == Document.document_id) + .join( + DocumentChunk, + ( + (DocumentChunk.document_id == Document.document_id) + if revision_pins is None + else and_( + DocumentChunk.document_id == Document.document_id, + or_( + *[ + and_( + DocumentChunk.document_id == document_id, + DocumentChunk.job_result_id == job_result_id, + ) + for document_id, job_result_id in target_ids_by_revision + ] + ), + ) + ), + ) .outerjoin(DocumentSection, DocumentSection.section_id == DocumentChunk.section_id) .join(JobResult, JobResult.id == DocumentChunk.job_result_id) .where(or_(*revision_filters)) diff --git a/packages/shared-python/shared/services/retrieval/hydration/reference.py b/packages/shared-python/shared/services/retrieval/hydration/reference.py index 46208df5..cd240654 100644 --- a/packages/shared-python/shared/services/retrieval/hydration/reference.py +++ b/packages/shared-python/shared/services/retrieval/hydration/reference.py @@ -1,8 +1,9 @@ from __future__ import annotations +from collections.abc import Mapping from typing import Any -from sqlalchemy import select +from sqlalchemy import and_, or_, select from sqlalchemy.ext.asyncio import AsyncSession from shared.models.database.document import Document, DocumentChunk, DocumentSection @@ -20,6 +21,7 @@ async def hydrate_referenced_chunk_rows( namespace: str, refs: list[dict[str, Any]], score_by_chunk_id: dict[str, float] | None = None, + revision_pins: Mapping[str, str] | None = None, ) -> list[dict[str, Any]]: if db is None or not refs: return [] @@ -39,22 +41,48 @@ async def hydrate_referenced_chunk_rows( document_ids = sorted({document_id for document_id, _, _, _ in ref_keys}) chunk_ids = sorted({chunk_id for _, chunk_id, _, _ in ref_keys}) - stmt = ( - select(Document, DocumentChunk, DocumentSection, JobResult) - .join( - DocumentChunk, + pinned_document_ids = [ + document_id for document_id in document_ids if revision_pins and document_id in revision_pins + ] + if revision_pins is not None and not pinned_document_ids: + return [] + + if revision_pins is None: + chunk_join = ( (DocumentChunk.document_id == Document.document_id) - & (DocumentChunk.job_result_id == Document.current_job_result_id), + & (DocumentChunk.job_result_id == Document.current_job_result_id) + ) + else: + chunk_join = and_( + DocumentChunk.document_id == Document.document_id, + or_( + *[ + and_( + DocumentChunk.document_id == document_id, + DocumentChunk.job_result_id == str(revision_pins[document_id]), + ) + for document_id in pinned_document_ids + ] + ), ) + + stmt = ( + select(Document, DocumentChunk, DocumentSection, JobResult) + .join(DocumentChunk, chunk_join) .outerjoin(DocumentSection, DocumentSection.section_id == DocumentChunk.section_id) .join(JobResult, JobResult.id == DocumentChunk.job_result_id) .where(Document.user_id == user_id) .where(Document.namespace == namespace) - .where(Document.status == 'active') - .where(Document.document_id.in_(document_ids)) + .where( + Document.document_id.in_( + document_ids if revision_pins is None else pinned_document_ids + ) + ) .where(DocumentChunk.chunk_id.in_(chunk_ids)) .order_by(DocumentChunk.sort_order) ) + if revision_pins is None: + stmt = stmt.where(Document.status == 'active') result = await db.execute(stmt) rows_by_key: dict[ReferenceLookupKey, dict[str, Any]] = {} diff --git a/packages/shared-python/shared/services/retrieval/hydration/result_assembly.py b/packages/shared-python/shared/services/retrieval/hydration/result_assembly.py index ed448704..863bfadb 100644 --- a/packages/shared-python/shared/services/retrieval/hydration/result_assembly.py +++ b/packages/shared-python/shared/services/retrieval/hydration/result_assembly.py @@ -1,5 +1,6 @@ from __future__ import annotations +from collections.abc import Mapping from typing import Any from sqlalchemy.ext.asyncio import AsyncSession @@ -21,6 +22,7 @@ async def assemble_retrieval_results( exclude_document_ids: list[str], exclude_sections: list[dict[str, str]], allowed_chunk_types: set[str] | None = None, + revision_pins: Mapping[str, str] | None = None, ) -> list[dict[str, Any]]: filtered_rows = filter_excluded_rows( rows, @@ -37,6 +39,7 @@ async def assemble_retrieval_results( rows=filtered_rows, exclude_document_ids=exclude_document_ids, exclude_sections=exclude_sections, + revision_pins=revision_pins, ) rows_by_chunk_id = { str(row.get('chunk_id') or ''): row diff --git a/packages/shared-python/shared/services/retrieval/manifest_cache.py b/packages/shared-python/shared/services/retrieval/manifest_cache.py new file mode 100644 index 00000000..3861903e --- /dev/null +++ b/packages/shared-python/shared/services/retrieval/manifest_cache.py @@ -0,0 +1,47 @@ +"""Request-scoped decoded serving-manifest reuse.""" + +from __future__ import annotations + +from collections.abc import Mapping +from typing import Any + + +_CACHE_INFO_KEY = "retrieval_serving_manifest_payloads" + + +def manifest_revision_key(revisions: Mapping[str, str]) -> tuple[tuple[str, str], ...]: + """Return a stable key for one immutable revision pin set.""" + return tuple(sorted((str(document_id), str(job_result_id)) for document_id, job_result_id in revisions.items())) + + +def get_cached_manifest_payloads( + db: object, + *, + revisions: Mapping[str, str], +) -> dict[tuple[str, str], dict[str, Any]] | None: + """Return decoded manifests cached on this request's SQLAlchemy session.""" + info = getattr(db, "info", None) + if not isinstance(info, dict): + return None + batches = info.get(_CACHE_INFO_KEY) + if not isinstance(batches, dict): + return None + payloads = batches.get(manifest_revision_key(revisions)) + if not isinstance(payloads, dict): + return None + return payloads + + +def cache_manifest_payloads( + db: object, + *, + revisions: Mapping[str, str], + payloads: dict[tuple[str, str], dict[str, Any]], +) -> None: + """Store one complete decoded manifest batch for this request only.""" + info = getattr(db, "info", None) + if not isinstance(info, dict): + return + batches = info.setdefault(_CACHE_INFO_KEY, {}) + if isinstance(batches, dict): + batches[manifest_revision_key(revisions)] = payloads diff --git a/packages/shared-python/shared/services/retrieval/nav/nav_knowhere.py b/packages/shared-python/shared/services/retrieval/nav/nav_knowhere.py index 2aad0b29..e25234ea 100644 --- a/packages/shared-python/shared/services/retrieval/nav/nav_knowhere.py +++ b/packages/shared-python/shared/services/retrieval/nav/nav_knowhere.py @@ -47,6 +47,7 @@ PersistedScoreUnit, tokenize_query_for_ranker, ) +from shared.services.retrieval.serving_manifest import decode_serving_manifest _ASSET_TYPES = ("table", "image") # Knowhere sentinel path for the virtual document container (not a collectable leaf). @@ -54,6 +55,11 @@ _DEFAULT_DSN = "postgresql://root:root123@127.0.0.1:5433/Knowhere" _MAP_UNIT_INDEX_FORMAT_VERSION = 1 _MAP_SCORE_CHANNELS: Tuple[str, str] = ("path", "content") +# PostgreSQL's trigram planner evaluates one LIKE branch per pattern. Once a +# query has more than a handful of patterns, fetching the bounded pinned unit +# set once and evaluating literal substring hits in Python is faster and keeps +# SQL planning time from growing with the planner's subgoal length. +_TERM_SCORE_FULL_SCAN_PATTERN_THRESHOLD = 8 _logger = logging.getLogger(__name__) @@ -220,6 +226,19 @@ def __init__( self._revisions = dict(revisions) self._excluded_sections = set(excluded_sections or ()) self._conn: Optional[_SyncConnection] = None + self._score_manifest_cache: Optional[ + tuple[tuple[tuple[str, str], ...], list[Sequence[object]]] + ] = None + self._score_unit_rows_cache: dict[ + tuple[tuple[str, str], ...], list[Sequence[object]] + ] = {} + self._score_frequency_cache: dict[ + tuple[tuple[str, str], ...], + dict[tuple[str, str], dict[str, int]], + ] = {} + self._score_term_cache: dict[ + tuple[tuple[str, str], ...], dict[str, Tuple[float, ...]] + ] = {} def _connection(self) -> "_SyncConnection": if self._conn is None: @@ -317,26 +336,71 @@ def load_persisted_score_corpus( ] cur = self._connection().cursor() try: - stage_started = time.perf_counter() - 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()) + revision_key = tuple(revisions) + cached_manifests = self._score_manifest_cache + if cached_manifests is not None and cached_manifests[0] == revision_key: + manifests = cached_manifests[1] + _logger.info( + "retrieval map-index load stage=manifests cache_hit rows=%d", + len(manifests), + ) + else: + stage_started = time.perf_counter() + try: + cur.execute( + "SELECT indexes.document_id, indexes.job_result_id, " + "indexes.format_version, indexes.unit_count, indexes.token_count, " + "manifests.payload_zlib, manifests.checksum, manifests.format_version, " + "statistics.payload_zlib, statistics.checksum, statistics.format_version " + "FROM document_map_unit_indexes AS indexes " + f"JOIN (VALUES {values_sql}) AS revisions(document_id, job_result_id) " + "ON indexes.document_id = revisions.document_id " + "AND indexes.job_result_id = revisions.job_result_id " + "JOIN retrieval_serving_revision_manifests AS manifests " + "ON manifests.document_id = indexes.document_id " + "AND manifests.job_result_id = indexes.job_result_id " + "JOIN retrieval_serving_revision_stats AS statistics " + "ON statistics.document_id = indexes.document_id " + "AND statistics.job_result_id = indexes.job_result_id", + revision_params, + ) + except Exception as exc: + if getattr(exc, "pgcode", None) == "42P01": + return None + raise + manifests = list(cur.fetchall()) + self._score_manifest_cache = (revision_key, manifests) _logger.info( "retrieval map-index load stage=manifests seconds=%.3f rows=%d", - time.perf_counter() - stage_started, + time.perf_counter() - stage_started if cached_manifests is None else 0.0, len(manifests), ) if len(manifests) != len(revisions) or any( - int(row[2]) != _MAP_UNIT_INDEX_FORMAT_VERSION for row in manifests + len(row) < 11 + or int(row[2]) != _MAP_UNIT_INDEX_FORMAT_VERSION + or not row[5] + or not row[6] + or not row[8] + or not row[9] + for row in manifests ): return None + decoded_manifests: dict[tuple[str, str], dict[str, Any]] = {} + decoded_statistics: dict[tuple[str, str], dict[str, Any]] = {} + try: + for row in manifests: + decoded_manifests[(str(row[0]), str(row[1]))] = decode_serving_manifest( + bytes(row[5]), + checksum=str(row[6]), + format_version=int(row[7]), + ) + decoded_statistics[(str(row[0]), str(row[1]))] = decode_serving_manifest( + bytes(row[8]), + checksum=str(row[9]), + format_version=int(row[10]), + ) + except ValueError: + return None # The marker is written last in the same transaction that inserts # all units and token rows. A committed marker therefore denotes @@ -355,30 +419,63 @@ def load_persisted_score_corpus( for document_id, section_ids in allowed_by_document.items() for section_id in section_ids ] + unit_cache_key = revision_key + all_unit_rows = self._score_unit_rows_cache.get(unit_cache_key, []) + units_cache_hit = unit_cache_key in self._score_unit_rows_cache unit_rows: list[Sequence[object]] = [] if allowed_pairs: - allowed_document_ids = [pair[0] for pair in allowed_pairs] - allowed_section_ids = [pair[1] for pair in allowed_pairs] - stage_started = time.perf_counter() - cur.execute( - "SELECT units.id, units.document_id, units.unit_id, units.section_id, " - "units.path_token_count, units.content_token_count " - "FROM document_map_units AS units " - f"JOIN (VALUES {values_sql}) AS revisions(document_id, job_result_id) " - "ON units.document_id = revisions.document_id " - "AND units.job_result_id = revisions.job_result_id " - "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()) + if not units_cache_hit: + stage_started = time.perf_counter() + manifest_rows: list[Sequence[object]] = [] + for document_id, job_result_id in revisions: + payload = decoded_manifests.get((document_id, job_result_id), {}) + raw_units = payload.get("map_units") + if not isinstance(raw_units, list): + manifest_rows = [] + break + for raw_unit in raw_units: + if not isinstance(raw_unit, dict): + manifest_rows = [] + break + row_id = str(raw_unit.get("row_id") or "").strip() + unit_id = str(raw_unit.get("unit_id") or "").strip() + if not row_id or not unit_id: + manifest_rows = [] + break + manifest_rows.append( + ( + row_id, + document_id, + unit_id, + str(raw_unit.get("section_id") or ""), + int(raw_unit.get("path_token_count") or 0), + int(raw_unit.get("content_token_count") or 0), + ) + ) + if not manifest_rows and raw_units: + break + expected_unit_count = sum(int(row[3]) for row in manifests) + all_unit_rows = ( + manifest_rows + if len(manifest_rows) == expected_unit_count + else [] + ) + if expected_unit_count and not all_unit_rows: + return None + self._score_unit_rows_cache[unit_cache_key] = all_unit_rows + else: + stage_started = time.perf_counter() + allowed_pairs_set = set(allowed_pairs) + unit_rows = [ + row + for row in all_unit_rows + if (str(row[1]), str(row[3])) in allowed_pairs_set + ] _logger.info( - "retrieval map-index load stage=units seconds=%.3f rows=%d", - time.perf_counter() - stage_started, + "retrieval map-index load stage=units seconds=%.3f rows=%d cache_hit=%s", + time.perf_counter() - stage_started if not units_cache_hit else 0.0, len(unit_rows), + units_cache_hit, ) map_unit_ids = [str(row[0]) for row in unit_rows] unique_queries = list(dict.fromkeys(str(query) for query in queries)) @@ -394,6 +491,48 @@ def load_persisted_score_corpus( ) frequencies: Dict[Tuple[str, str], Dict[str, int]] = {} if map_unit_ids and query_tokens: + full_frequency_map = self._score_frequency_cache.get(revision_key) + if full_frequency_map is None: + full_frequency_map = {} + for document_id, job_result_id in revisions: + payload = decoded_statistics.get((document_id, job_result_id), {}) + unit_frequencies = payload.get("unit_frequencies", {}) + if not isinstance(unit_frequencies, dict): + continue + for map_unit_id, by_channel in unit_frequencies.items(): + if not isinstance(by_channel, dict): + continue + for channel in _MAP_SCORE_CHANNELS: + values = by_channel.get(channel, {}) + if isinstance(values, dict): + full_frequency_map[(str(map_unit_id), channel)] = { + str(token): int(value) + for token, value in values.items() + } + self._score_frequency_cache[revision_key] = full_frequency_map + frequencies = { + key: { + token: value + for token, value in values.items() + if token in query_tokens + } + for key, values in full_frequency_map.items() + if key[0] in map_unit_ids + } + expected_units_by_revision = { + (str(row[0]), str(row[1])): int(row[3]) for row in manifests + } + statistics_complete = all( + int( + decoded_statistics.get((document_id, job_result_id), {}).get( + "unit_count", -1 + ) + ) + == expected_units_by_revision.get((document_id, job_result_id), -1) + for document_id, job_result_id in revisions + ) + manifest_frequency_complete = bool(map_unit_ids) and statistics_complete + if map_unit_ids and query_tokens and not manifest_frequency_complete: query_token_hashes = [ sha256(token.encode("utf-8")).hexdigest() for token in query_tokens ] @@ -420,12 +559,16 @@ def load_persisted_score_corpus( len(map_unit_ids), ) - term_scores = self._load_term_scores( - cur, - map_unit_ids=map_unit_ids, - queries=unique_queries, - query_tokens_by_query=query_tokens_by_query, - ) + term_cache_key = (revision_key, tuple(unique_queries)) + term_scores = self._score_term_cache.get(term_cache_key) + if term_scores is None: + term_scores = self._load_term_scores( + cur, + map_unit_ids=map_unit_ids, + queries=unique_queries, + query_tokens_by_query=query_tokens_by_query, + ) + self._score_term_cache[term_cache_key] = term_scores _logger.info( "retrieval map-index load stage=complete units=%d queries=%d", len(unit_rows), @@ -481,42 +624,81 @@ def _load_term_scores( ) -> Dict[str, Tuple[float, ...]]: if not map_unit_ids or not queries: return {} - expressions: List[str] = [] - params: List[object] = [] + # Keep candidate selection in PostgreSQL so the trigram index can + # discard non-matching units, but compute the exact score once per + # returned row in Python. The previous query generated one POSITION + # expression and one OR predicate for every query token. Long planner + # subgoals therefore produced very large SQL statements and repeated + # substring evaluation for the same row. LIKE ANY keeps the SQL shape + # constant while preserving the existing literal substring semantics in + # the final Python scoring pass (LIKE may over-select wildcard matches, + # which are rejected by the literal checks below). + candidate_patterns: list[str] = [] for query in queries: query_lower = query.lower().strip() if not query_lower: - 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" + candidate_patterns.append(f"%{query_lower}%") + candidate_patterns.extend( + f"%{token}%" for token in query_tokens_by_query[query] if token ) - params.append(query_lower) - params.extend(query_tokens_by_query[query]) - params.append(list(map_unit_ids)) + candidate_patterns = list(dict.fromkeys(candidate_patterns)) + if not candidate_patterns: + return {} stage_started = time.perf_counter() - cur.execute( - "SELECT id, " + ", ".join(expressions) + " " - "FROM document_map_units WHERE id = ANY(%s)", - params, - ) + if len(candidate_patterns) > _TERM_SCORE_FULL_SCAN_PATTERN_THRESHOLD: + # Long subgoals make LIKE ANY increasingly expensive even with the + # trigram index. The map-unit id list is already bounded by the + # pinned serving projection, so one text fetch plus literal Python + # checks avoids a query whose shape grows with token count. + cur.execute( + "SELECT id, term_search_text_lower " + "FROM document_map_units WHERE id = ANY(%s)", + (list(map_unit_ids),), + ) + candidate_mode = "full_scan" + else: + cur.execute( + "SELECT id, term_search_text_lower " + "FROM document_map_units " + "WHERE id = ANY(%s) AND term_search_text_lower LIKE ANY(%s)", + (list(map_unit_ids), candidate_patterns), + ) + candidate_mode = "trigram" rows = cur.fetchall() _logger.info( - "retrieval map-index load stage=term_scores units=%d queries=%d seconds=%.3f", + "retrieval map-index load stage=term_scores units=%d queries=%d mode=%s seconds=%.3f", len(map_unit_ids), len(queries), + candidate_mode, time.perf_counter() - stage_started, ) - return { - str(row[0]): tuple(float(value) for value in row[1:]) - for row in rows - } + scores_by_unit: Dict[str, Tuple[float, ...]] = {} + for row in rows: + if len(row) < 2: + continue + unit_id = str(row[0]) + haystack = str(row[1] or "").lower() + scores: list[float] = [] + for query in queries: + query_lower = query.lower().strip() + if not query_lower: + scores.append(0.0) + elif query_lower in haystack: + scores.append(100.0) + else: + scores.append( + float( + sum( + 1 + for token in query_tokens_by_query[query] + if token in haystack + ) + ) + ) + if any(score > 0.0 for score in scores): + scores_by_unit[unit_id] = tuple(scores) + return scores_by_unit def _load_persisted_bm25_stats( self, diff --git a/packages/shared-python/shared/services/retrieval/nav/nav_orchestrate.py b/packages/shared-python/shared/services/retrieval/nav/nav_orchestrate.py index 2e82be08..52cf8867 100644 --- a/packages/shared-python/shared/services/retrieval/nav/nav_orchestrate.py +++ b/packages/shared-python/shared/services/retrieval/nav/nav_orchestrate.py @@ -203,6 +203,8 @@ def _relit_map( """ relit = prepared q = (query or "").strip() + if relit is None and q: + relit = state.relit_map_cache.get(q) if relit is None and q: try: from .nav_map_scores import relight_map_for_query @@ -215,6 +217,7 @@ def _relit_map( ) if scores: relit = (scores, units, highlights) + state.relit_map_cache[q] = relit except Exception: relit = None if relit is None: diff --git a/packages/shared-python/shared/services/retrieval/nav/nav_types.py b/packages/shared-python/shared/services/retrieval/nav/nav_types.py index 8f3366c0..3c6b5241 100644 --- a/packages/shared-python/shared/services/retrieval/nav/nav_types.py +++ b/packages/shared-python/shared/services/retrieval/nav/nav_types.py @@ -305,6 +305,12 @@ class NavState: # PLAN's rewritten retrieval_query per subgoal after widen. Overrides the # planned query for the next harvest, and re-scores the shared map with it. subgoal_refined_queries: Dict[str, str] = field(default_factory=dict) + # Episode-local map scores keyed by retrieval query. Checklist waves may + # revisit the same subgoal query; reuse the exact score snapshot instead of + # rebuilding the persisted index and rescoring the corpus. + relit_map_cache: Dict[ + str, Tuple[Dict[str, float], Dict[str, float], List[str]] + ] = field(default_factory=dict) # Per-subgoal "seen but not selected" section ids — hidden from later map # views for that subgoal so widen surfaces siblings instead of dead ends. subgoal_dismissed_section_ids: Dict[str, set[str]] = field(default_factory=dict) diff --git a/packages/shared-python/shared/services/retrieval/nav_snapshot.py b/packages/shared-python/shared/services/retrieval/nav_snapshot.py index 4eec0404..c68d6d8c 100644 --- a/packages/shared-python/shared/services/retrieval/nav_snapshot.py +++ b/packages/shared-python/shared/services/retrieval/nav_snapshot.py @@ -14,10 +14,25 @@ from collections.abc import Callable, Iterator, Mapping from typing import Any, Protocol -from sqlalchemy import Executable, literal, select, tuple_ +from sqlalchemy import ( + ARRAY, + Executable, + String, + bindparam, + cast, + func, + literal, + select, + tuple_, +) from sqlalchemy.engine import Result -from shared.models.database.document import Document, DocumentChunk, DocumentSection +from shared.models.database.document import ( + Document, + DocumentChunk, + DocumentSection, + RetrievalServingRevisionManifest, +) from shared.models.database.job_result import JobResult from shared.services.retrieval.nav.nav_knowhere import ( LazyKnowhereProvider, @@ -29,12 +44,13 @@ knowhere_database_url, ) from shared.services.retrieval.search.section_filters import is_excluded_section +from shared.services.retrieval.serving_manifest import decode_serving_manifest +from shared.services.retrieval.manifest_cache import get_cached_manifest_payloads # Keep each payload query bounded under the API's 30-second statement timeout. -# Ten-thousand-row keyset pages avoid OFFSET scans while keeping each payload -# statement bounded. The contract benchmark verifies this page size against -# the full 2 KiB content and metadata payload. +# Ten-thousand-row keyset pages avoid OFFSET scans while keeping the reference +# payload bounded under the API's 30-second asyncpg command timeout. _CHUNK_BATCH_SIZE = 10_000 # Keep revision predicates bounded while reducing round trips for large # namespaces. Keyset paging still caps each payload query at 10,000 rows. @@ -48,6 +64,9 @@ class SnapshotSession(Protocol): async def execute(self, statement: Executable) -> Result[tuple[object, ...]]: raise NotImplementedError + async def rollback(self) -> None: + raise NotImplementedError + @dataclass(frozen=True) class NavSnapshot: @@ -57,6 +76,7 @@ class NavSnapshot: chunk_ref_index: Mapping[str, dict[str, Any]] document_ids: list[str] document_titles: dict[str, str] + document_revisions: Mapping[str, str] | None def close(self) -> None: close = getattr(self.provider, "close", None) @@ -70,9 +90,12 @@ def build_nav_snapshot( sections_by_doc: dict[str, list[SectionRow]], units_by_doc: dict[str, list[UnitRow]], chunk_ref_index: dict[str, dict[str, Any]], + document_revisions: Mapping[str, str] | None = None, ) -> NavSnapshot: """Assemble provider + index from already-fetched rows (also used by tests).""" - doc_ids = [did for did in document_titles if did in sections_by_doc or did in units_by_doc] + doc_ids = [ + did for did in document_titles if did in sections_by_doc or did in units_by_doc + ] if not doc_ids: raise ValueError("nav snapshot requires at least one active document") @@ -90,7 +113,18 @@ def build_nav_snapshot( provider=provider, chunk_ref_index=dict(chunk_ref_index), document_ids=list(provider.document_ids()), - document_titles={did: document_titles.get(did, did) for did in provider.document_ids()}, + document_titles={ + did: document_titles.get(did, did) for did in provider.document_ids() + }, + document_revisions=( + { + did: str(document_revisions.get(did, "")) + for did in provider.document_ids() + if document_revisions.get(did) + } + if document_revisions is not None + else None + ), ) @@ -102,19 +136,38 @@ async def load_nav_snapshot( exclude_document_ids: list[str] | None = None, exclude_sections: list[dict[str, str]] | None = None, lazy: bool = False, + revision_pins: Mapping[str, str] | None = None, ) -> NavSnapshot: """Preload namespace current revision into a sync map-nav snapshot.""" - excluded_docs = [str(x).strip() for x in (exclude_document_ids or ()) if str(x).strip()] + excluded_docs = [ + str(x).strip() for x in (exclude_document_ids or ()) if str(x).strip() + ] excluded_secs = list(exclude_sections or ()) - doc_stmt = ( - select(Document.document_id, Document.source_file_name, Document.current_job_result_id) - .where(Document.user_id == user_id) - .where(Document.namespace == namespace) - .where(Document.status == "active") - .where(Document.current_job_result_id.is_not(None)) - .order_by(Document.document_id) - ) + if revision_pins is None: + doc_stmt = ( + select( + Document.document_id, + Document.source_file_name, + Document.current_job_result_id, + ) + .where(Document.user_id == user_id) + .where(Document.namespace == namespace) + .where(Document.status == "active") + .where(Document.current_job_result_id.is_not(None)) + .order_by(Document.document_id) + ) + else: + pinned_document_ids = [str(document_id) for document_id in revision_pins] + if not pinned_document_ids: + raise ValueError("revision pins must include at least one document") + doc_stmt = ( + select(Document.document_id, Document.source_file_name) + .where(Document.user_id == user_id) + .where(Document.namespace == namespace) + .where(Document.document_id.in_(pinned_document_ids)) + .order_by(Document.document_id) + ) if excluded_docs: doc_stmt = doc_stmt.where(Document.document_id.notin_(excluded_docs)) doc_rows = list((await db.execute(doc_stmt)).all()) @@ -127,12 +180,18 @@ async def load_nav_snapshot( document_titles: dict[str, str] = {} current_job_result_ids: set[str] = set() document_revisions: list[tuple[str, str]] = [] - for document_id, source_file_name, current_job_result_id in doc_rows: + for doc_row in doc_rows: + document_id = doc_row[0] + source_file_name = doc_row[1] did = str(document_id) title = str(source_file_name or "").strip() or did document_titles[did] = title - if current_job_result_id: - job_result_id = str(current_job_result_id) + job_result_id = ( + str(revision_pins.get(did, "")) + if revision_pins is not None + else str(doc_row[2] or "") + ) + if job_result_id: current_job_result_ids.add(job_result_id) document_revisions.append((did, job_result_id)) @@ -147,28 +206,47 @@ async def load_nav_snapshot( if job_result_id and job_id } - sections_by_doc, section_path_by_id = await _load_sections( + manifest_sections = await _load_manifest_sections( db, document_revisions=document_revisions, exclude_sections=excluded_secs, + job_id_by_result_id=job_id_by_result_id, ) - if lazy: - chunk_ids_by_doc, chunk_ref_index, remounted_assets = await _load_chunk_index( + if manifest_sections is None: + sections_by_doc, section_path_by_id = await _load_sections( db, document_revisions=document_revisions, exclude_sections=excluded_secs, - section_path_by_id=section_path_by_id, - job_id_by_result_id=job_id_by_result_id, ) + manifest_chunk_index = None + else: + sections_by_doc, section_path_by_id, manifest_chunk_index = manifest_sections + if lazy: + if manifest_chunk_index is None: + chunk_ids_by_doc, chunk_ref_index, remounted_assets = await _load_chunk_index( + db, + document_revisions=document_revisions, + exclude_sections=excluded_secs, + section_path_by_id=section_path_by_id, + job_id_by_result_id=job_id_by_result_id, + ) + else: + chunk_ids_by_doc, chunk_ref_index, remounted_assets = manifest_chunk_index kept_titles = { - did: title for did, title in document_titles.items() if sections_by_doc.get(did) + did: title + for did, title in document_titles.items() + if sections_by_doc.get(did) } if not kept_titles: raise ValueError( f"nav snapshot empty after excludes for " f"user_id={user_id!r} namespace={namespace!r}" ) - revisions = {did: result_id for did, result_id in document_revisions if did in kept_titles} + revisions = { + did: result_id + for did, result_id in document_revisions + if did in kept_titles + } store = ReadOnlyChunkStore( dsn=knowhere_database_url(), revisions=revisions, @@ -192,7 +270,9 @@ async def load_nav_snapshot( chunk_store=store, known_chunk_ids=chunk_ids_by_doc.get(did, ()), root_asset_ids=remounted_assets.get(did, {}).get("root", ()), - remounted_assets_by_section=remounted_assets.get(did, {}).get("owners", {}), + remounted_assets_by_section=remounted_assets.get(did, {}).get( + "owners", {} + ), ) for did in kept_titles ] @@ -215,7 +295,10 @@ async def load_nav_snapshot( resolver=store.load_chunk_reference_metadata, ), document_ids=list(provider.document_ids()), - document_titles={did: kept_titles.get(did, did) for did in provider.document_ids()}, + document_titles={ + did: kept_titles.get(did, did) for did in provider.document_ids() + }, + document_revisions=dict(revisions), ) units_by_doc, chunk_ref_index = await _load_chunks( @@ -228,9 +311,7 @@ async def load_nav_snapshot( # Keep only documents that still have sections after exclude filters. kept_titles = { - did: title - for did, title in document_titles.items() - if sections_by_doc.get(did) + did: title for did, title in document_titles.items() if sections_by_doc.get(did) } if not kept_titles: raise ValueError( @@ -243,9 +324,174 @@ async def load_nav_snapshot( sections_by_doc={did: sections_by_doc.get(did, []) for did in kept_titles}, units_by_doc={did: units_by_doc.get(did, []) for did in kept_titles}, chunk_ref_index=chunk_ref_index, + document_revisions={ + document_id: job_result_id + for document_id, job_result_id in document_revisions + if document_id in kept_titles + }, ) +async def _load_manifest_sections( + db: SnapshotSession, + *, + document_revisions: list[tuple[str, str]], + exclude_sections: list[dict[str, str]], + job_id_by_result_id: dict[str, str], +) -> tuple[ + dict[str, list[SectionRow]], + dict[str, str], + tuple[dict[str, list[str]], dict[str, dict[str, Any]], dict[str, dict[str, Any]]], +] | None: + """Load section metadata from complete serving manifests when available.""" + cached_payloads = get_cached_manifest_payloads( + db, + revisions=dict(document_revisions), + ) + if cached_payloads is not None and len(cached_payloads) == len(document_revisions): + manifest_entries: list[tuple[object, ...]] = [ + (document_id, job_result_id, payload, None, None) + for document_id, job_result_id in document_revisions + for payload in [cached_payloads.get((document_id, job_result_id))] + if payload is not None + ] + if len(manifest_entries) != len(document_revisions): + return None + else: + statement = select( + RetrievalServingRevisionManifest.document_id, + RetrievalServingRevisionManifest.job_result_id, + RetrievalServingRevisionManifest.payload_zlib, + RetrievalServingRevisionManifest.checksum, + RetrievalServingRevisionManifest.format_version, + ).where( + tuple_( + RetrievalServingRevisionManifest.document_id, + RetrievalServingRevisionManifest.job_result_id, + ).in_(document_revisions) + ) + try: + rows = (await db.execute(statement)).all() + except Exception: + await db.rollback() + return None + if len(rows) != len(document_revisions): + return None + manifest_entries = [tuple(row) for row in rows] + by_doc: dict[str, list[SectionRow]] = {} + path_by_id: dict[str, str] = {} + ids_by_doc: dict[str, list[str]] = {} + ref_index: dict[str, dict[str, Any]] = {} + root_assets_by_doc: dict[str, set[str]] = {} + text_connections_by_doc: dict[str, list[tuple[str, str]]] = {} + try: + for document_id, _job_result_id, payload_zlib, checksum, format_version in manifest_entries: + if isinstance(payload_zlib, dict): + payload = payload_zlib + else: + if not isinstance(payload_zlib, (bytes, bytearray, memoryview)): + return None + if checksum is None or format_version is None: + return None + payload = decode_serving_manifest( + bytes(payload_zlib), + checksum=str(checksum), + format_version=int(str(format_version)), + ) + raw_sections = payload.get("sections") + if not isinstance(raw_sections, list): + return None + for raw_section in raw_sections: + if not isinstance(raw_section, dict): + return None + section_path = str(raw_section.get("section_path") or "") + if is_excluded_section( + document_id=str(document_id), + section_path=section_path, + exclude_sections=exclude_sections, + ): + continue + section_id = str(raw_section.get("section_id") or "") + if not section_id or not section_path: + return None + section = SectionRow( + section_id=section_id, + parent_section_id=( + str(raw_section["parent_section_id"]) + if raw_section.get("parent_section_id") + else None + ), + section_path=section_path, + section_title=str(raw_section.get("section_title") or "").strip(), + section_level=int(raw_section.get("section_level") or 0), + summary=str(raw_section.get("summary") or "").strip(), + sort_order=int(raw_section.get("sort_order") or 0), + ) + by_doc.setdefault(str(document_id), []).append(section) + path_by_id[section_id] = section_path + raw_chunks = payload.get("chunks") + if not isinstance(raw_chunks, list): + return None + for raw_chunk in raw_chunks: + if not isinstance(raw_chunk, dict): + return None + chunk_id = str(raw_chunk.get("chunk_id") or "").strip() + if not chunk_id: + return None + section_id = ( + str(raw_chunk["section_id"]) + if raw_chunk.get("section_id") + else None + ) + section_path = path_by_id.get(section_id) if section_id else None + if is_excluded_section( + document_id=str(document_id), + section_path=section_path, + exclude_sections=exclude_sections, + ) or (section_id and section_id not in path_by_id): + continue + chunk_type = str(raw_chunk.get("chunk_type") or "text") + meta = { + "document_id": str(document_id), + "section_path": section_path, + "chunk_type": chunk_type, + "file_path": None, + "job_id": job_id_by_result_id.get(str(_job_result_id)), + } + document_key = str(document_id) + ids_by_doc.setdefault(document_key, []).append(chunk_id) + ref_index[chunk_id] = meta + ref_index[f"{document_key}:{chunk_id}"] = meta + if ( + chunk_type in {"image", "table"} + and section_id + and section_path == "Root" + ): + root_assets_by_doc.setdefault(document_key, set()).add(chunk_id) + connections = raw_chunk.get("connect_to") + if chunk_type == "text" and isinstance(connections, list): + for connection in connections: + target = ( + str(connection.get("target") or "").strip() + if isinstance(connection, dict) + else str(connection or "").strip() + ) + if target: + text_connections_by_doc.setdefault(document_key, []).append( + (section_id or "", target) + ) + except (TypeError, ValueError, KeyError): + return None + remounted: dict[str, dict[str, Any]] = {} + for document_id, asset_ids in root_assets_by_doc.items(): + owners: dict[str, list[str]] = {} + for section_id, target in text_connections_by_doc.get(document_id, ()): + if target in asset_ids: + owners.setdefault(section_id, []).append(target) + remounted[document_id] = {"root": sorted(asset_ids), "owners": owners} + return by_doc, path_by_id, (ids_by_doc, ref_index, remounted) + + async def _load_chunk_index( db: SnapshotSession, *, @@ -431,62 +677,109 @@ async def _load_sections( exclude_sections: list[dict[str, str]], ) -> tuple[dict[str, list[SectionRow]], dict[str, str]]: # Captured pairs replace DocumentSection.job_result_id == Document.current_job_result_id. - stmt = ( - select( - DocumentSection.document_id, - DocumentSection.section_id, - DocumentSection.parent_section_id, - DocumentSection.section_path, - DocumentSection.section_title, - DocumentSection.section_level, - DocumentSection.summary, - DocumentSection.sort_order, + by_doc: dict[str, list[SectionRow]] = {} + path_by_id: dict[str, str] = {} + document_ids = [document_id for document_id, _ in document_revisions] + job_result_ids = [job_result_id for _, job_result_id in document_revisions] + revision_rows = ( + func.unnest( + cast( + bindparam("section_document_ids", value=document_ids), ARRAY(String()) + ), + cast( + bindparam("section_job_result_ids", value=job_result_ids), + ARRAY(String()), + ), ) - .where( - tuple_( + .table_valued("document_id", "job_result_id") + .render_derived(name="revisions") + ) + last_key: tuple[str, str, int, str] | None = None + query_seconds = 0.0 + assembly_seconds = 0.0 + query_count = 0 + row_count = 0 + while True: + stmt = ( + select( DocumentSection.document_id, + DocumentSection.section_id, + DocumentSection.parent_section_id, + DocumentSection.section_path, + DocumentSection.section_title, + DocumentSection.section_level, + DocumentSection.summary, + DocumentSection.sort_order, DocumentSection.job_result_id, - ).in_(document_revisions) - ) - .order_by( - DocumentSection.document_id, - DocumentSection.sort_order, - DocumentSection.section_id, + ) + .join( + revision_rows, + (DocumentSection.document_id == revision_rows.c.document_id) + & (DocumentSection.job_result_id == revision_rows.c.job_result_id), + ) + .order_by( + DocumentSection.document_id, + DocumentSection.job_result_id, + DocumentSection.sort_order, + DocumentSection.section_id, + ) + .limit(_CHUNK_BATCH_SIZE) ) - ) - - by_doc: dict[str, list[SectionRow]] = {} - path_by_id: dict[str, str] = {} - query_started = time.perf_counter() - rows = (await db.execute(stmt)).all() - query_seconds = time.perf_counter() - query_started - assembly_started = time.perf_counter() - for row in rows: - document_id = str(row[0]) - section_path = str(row[3] or "") - if is_excluded_section( - document_id=document_id, - section_path=section_path, - exclude_sections=exclude_sections, - ): - continue - section_id = str(row[1]) - section = SectionRow( - section_id=section_id, - parent_section_id=str(row[2]) if row[2] else None, - section_path=section_path, - section_title=str(row[4] or "").strip(), - section_level=int(row[5] or 0), - summary=str(row[6] or "").strip(), - sort_order=int(row[7] or 0), + if last_key is not None: + stmt = stmt.where( + tuple_( + DocumentSection.document_id, + DocumentSection.job_result_id, + DocumentSection.sort_order, + DocumentSection.section_id, + ) + > tuple_(*[literal(value) for value in last_key]) + ) + query_started = time.perf_counter() + rows = (await db.execute(stmt)).all() + query_seconds += time.perf_counter() - query_started + query_count += 1 + if not rows: + break + row_count += len(rows) + assembly_started = time.perf_counter() + for row in rows: + document_id = str(row[0]) + section_path = str(row[3] or "") + if is_excluded_section( + document_id=document_id, + section_path=section_path, + exclude_sections=exclude_sections, + ): + continue + section_id = str(row[1]) + section = SectionRow( + section_id=section_id, + parent_section_id=str(row[2]) if row[2] else None, + section_path=section_path, + section_title=str(row[4] or "").strip(), + section_level=int(row[5] or 0), + summary=str(row[6] or "").strip(), + sort_order=int(row[7] or 0), + ) + by_doc.setdefault(document_id, []).append(section) + path_by_id[section_id] = section_path + assembly_seconds += time.perf_counter() - assembly_started + last = rows[-1] + last_key = ( + str(last[0]), + str(last[8]), + int(last[7] or 0), + str(last[1]), ) - by_doc.setdefault(document_id, []).append(section) - path_by_id[section_id] = section_path + if len(rows) < _CHUNK_BATCH_SIZE: + break _logger.info( - "retrieval snapshot phase=sections rows=%d query_seconds=%.3f assembly_seconds=%.3f", - len(rows), + "retrieval snapshot phase=sections rows=%d queries=%d query_seconds=%.3f assembly_seconds=%.3f", + row_count, + query_count, query_seconds, - time.perf_counter() - assembly_started, + assembly_seconds, ) return by_doc, path_by_id diff --git a/packages/shared-python/shared/services/retrieval/publication_content.py b/packages/shared-python/shared/services/retrieval/publication_content.py index 19a7f207..34092436 100644 --- a/packages/shared-python/shared/services/retrieval/publication_content.py +++ b/packages/shared-python/shared/services/retrieval/publication_content.py @@ -14,6 +14,7 @@ ) from shared.services.retrieval.map_unit_index import replace_document_map_units from shared.services.retrieval.publication_models import DocumentPublicationScope +from shared.services.retrieval.serving_manifest import persist_revision_serving_state from shared.services.retrieval.search.lexical_text import ( build_content_lexical_text, build_content_search_text, @@ -91,6 +92,8 @@ def replace_document_revision_content( ) db.flush() replace_document_map_units(db, scope=scope) + db.flush() + persist_revision_serving_state(db, scope=scope) class DocumentSectionPublisher: diff --git a/packages/shared-python/shared/services/retrieval/publication_service.py b/packages/shared-python/shared/services/retrieval/publication_service.py index e201916e..d7c879de 100644 --- a/packages/shared-python/shared/services/retrieval/publication_service.py +++ b/packages/shared-python/shared/services/retrieval/publication_service.py @@ -31,6 +31,13 @@ ExistingDocumentScope, PublishedDocumentState, ) +from shared.services.retrieval.serving_generation import ( + advance_namespace_generation, + lock_namespace_generation, +) +from shared.services.retrieval.serving_manifest import ( + rebuild_namespace_serving_statistics, +) def utc_now_naive() -> datetime: @@ -122,6 +129,23 @@ def _publish_document_state_for_job( skipped_all_duplicate=True, ) + existing_namespace = None + if document_id: + existing_namespace = db.execute( + select(Document.namespace).where( + Document.document_id == str(document_id) + ) + ).scalar_one_or_none() + namespaces_to_lock = {namespace} + if existing_namespace: + namespaces_to_lock.add(str(existing_namespace)) + for namespace_to_lock in sorted(namespaces_to_lock): + lock_namespace_generation( + db, + user_id=str(job.user_id), + namespace=namespace_to_lock, + ) + document = self._upsert_document_revision( db, job=job, @@ -156,6 +180,22 @@ def _publish_document_state_for_job( ) db.flush() + rebuild_namespace_serving_statistics( + db, + user_id=scope.user_id, + namespace=scope.namespace, + ) + if existing_namespace and str(existing_namespace) != scope.namespace: + rebuild_namespace_serving_statistics( + db, + user_id=scope.user_id, + namespace=str(existing_namespace), + ) + advance_namespace_generation( + db, + user_id=scope.user_id, + namespace=scope.namespace, + ) return PublishedDocumentState( user_id=str(job.user_id), namespace=namespace, @@ -209,6 +249,7 @@ def _upsert_document_revision( ) return None document.status = "active" + document.namespace = namespace document.archived_at = None document.current_job_result_id = job_result_id document.source_file_name = source_file_name or document.source_file_name diff --git a/packages/shared-python/shared/services/retrieval/search/channels.py b/packages/shared-python/shared/services/retrieval/search/channels.py index 1ddea1be..cdf67326 100644 --- a/packages/shared-python/shared/services/retrieval/search/channels.py +++ b/packages/shared-python/shared/services/retrieval/search/channels.py @@ -8,6 +8,7 @@ from __future__ import annotations import time +from collections.abc import Mapping from typing import Any from loguru import logger @@ -61,7 +62,7 @@ FROM document_chunks dc JOIN documents d ON d.document_id = dc.document_id - AND d.current_job_result_id = dc.job_result_id + {revision_join} LEFT JOIN document_sections ds ON ds.section_id = dc.section_id JOIN job_results jr @@ -69,12 +70,42 @@ WHERE d.user_id = :user_id AND d.namespace = :namespace AND d.status = 'active' + {revision_clause} {exclude_clause} {extra_filters} ) """ +def _build_revision_scope( + revision_pins: Mapping[str, str] | None, +) -> tuple[str, str, dict[str, Any]]: + if revision_pins is None: + return ( + "AND d.current_job_result_id = dc.job_result_id", + "", + {}, + ) + + pairs = [ + (str(document_id).strip(), str(job_result_id).strip()) + for document_id, job_result_id in revision_pins.items() + if str(document_id).strip() and str(job_result_id).strip() + ] + if not pairs: + return "", "AND FALSE", {} + + params: dict[str, Any] = {} + placeholders: list[str] = [] + for index, (document_id, job_result_id) in enumerate(pairs): + document_key = f"_pin_document_{index}" + revision_key = f"_pin_revision_{index}" + placeholders.append(f"(:{document_key}, :{revision_key})") + params[document_key] = document_id + params[revision_key] = job_result_id + return "", f"AND (dc.document_id, dc.job_result_id) IN ({', '.join(placeholders)})", params + + def _build_exclude_clause(exclude_document_ids: list[str]) -> str: if not exclude_document_ids: return "" @@ -228,6 +259,7 @@ async def path_channel( allowed_chunk_types: set[str] | None = None, signal_paths: list[str] | None = None, filter_mode: str = "delete", + revision_pins: Mapping[str, str] | None = None, ) -> list[dict[str, Any]]: """Path channel: BM25 over pre-tokenized path search text. @@ -246,6 +278,7 @@ async def path_channel( signal_paths=signal_paths, filter_mode=filter_mode, search_field="path_search_text", + revision_pins=revision_pins, ) @@ -261,6 +294,7 @@ async def content_channel( allowed_chunk_types: set[str] | None = None, signal_paths: list[str] | None = None, filter_mode: str = "delete", + revision_pins: Mapping[str, str] | None = None, ) -> list[dict[str, Any]]: """Content channel: BM25 over pre-tokenized content search text.""" return await _bm25_channel( @@ -275,6 +309,7 @@ async def content_channel( signal_paths=signal_paths, filter_mode=filter_mode, search_field="content_search_text", + revision_pins=revision_pins, ) @@ -291,6 +326,7 @@ async def _bm25_channel( signal_paths: list[str] | None, filter_mode: str, search_field: str, + revision_pins: Mapping[str, str] | None, ) -> list[dict[str, Any]]: if search_field not in {"content_search_text", "path_search_text"}: raise ValueError(f"Unsupported search_field: {search_field}") @@ -317,7 +353,14 @@ async def _bm25_channel( params.update(extra_params) params.update(section_params) + revision_join, revision_clause, revision_params = _build_revision_scope( + revision_pins + ) + params.update(revision_params) + corpus_cte = _SCOPED_CORPUS_CTE.format( + revision_join=revision_join, + revision_clause=revision_clause, exclude_clause=exclude_clause, extra_filters=extra_sql, ) @@ -409,6 +452,7 @@ async def term_channel( allowed_chunk_types: set[str] | None = None, signal_paths: list[str] | None = None, filter_mode: str = "delete", + revision_pins: Mapping[str, str] | None = None, ) -> list[dict[str, Any]]: """Term/grep channel: substring matching on term_search_text. @@ -435,6 +479,11 @@ async def term_channel( ) params.update(extra_params) + revision_join, revision_clause, revision_params = _build_revision_scope( + revision_pins + ) + params.update(revision_params) + ilike_conditions = [] for i, unit in enumerate(query_tokens): param_key = f"unit_{i}" @@ -448,6 +497,8 @@ async def term_channel( where_clause = " OR ".join(ilike_conditions) sql = ( _SCOPED_CORPUS_CTE.format( + revision_join=revision_join, + revision_clause=revision_clause, exclude_clause=exclude_clause, extra_filters=extra_sql ) + f""" diff --git a/packages/shared-python/shared/services/retrieval/search/discovery.py b/packages/shared-python/shared/services/retrieval/search/discovery.py index 2e3fdd3c..d71ccdd9 100644 --- a/packages/shared-python/shared/services/retrieval/search/discovery.py +++ b/packages/shared-python/shared/services/retrieval/search/discovery.py @@ -18,7 +18,9 @@ from __future__ import annotations +import asyncio import time +from collections.abc import Awaitable, Callable, Mapping from dataclasses import dataclass, field from typing import Any @@ -68,11 +70,13 @@ async def bottom_discovery( channels: list[str] | None = None, channel_weights: dict[str, float] | None = None, internal_recall_k: int | None = None, + revision_pins: Mapping[str, str] | None = None, **_kwargs: Any, ) -> DiscoveryResult: """Run 3-channel BM25 discovery plus RRF fusion.""" t0 = time.monotonic() try: + del db allowed_chunk_types = chunk_types effective_recall_k = ( internal_recall_k @@ -81,13 +85,10 @@ async def bottom_discovery( ) active_channels = set(channels) if channels else {"path", "content", "term"} - path_rows: list[dict[str, Any]] = [] - content_rows: list[dict[str, Any]] = [] - term_rows: list[dict[str, Any]] = [] - - if "path" in active_channels: - path_rows = await path_channel( - db, + path_rows, content_rows, term_rows = await asyncio.gather( + _run_channel( + path_channel, + enabled="path" in active_channels, user_id=user_id, namespace=namespace, query=query, @@ -97,11 +98,11 @@ async def bottom_discovery( allowed_chunk_types=allowed_chunk_types, signal_paths=signal_paths, filter_mode=filter_mode, - ) - - if "content" in active_channels: - content_rows = await content_channel( - db, + revision_pins=revision_pins, + ), + _run_channel( + content_channel, + enabled="content" in active_channels, user_id=user_id, namespace=namespace, query=query, @@ -111,11 +112,11 @@ async def bottom_discovery( allowed_chunk_types=allowed_chunk_types, signal_paths=signal_paths, filter_mode=filter_mode, - ) - - if "term" in active_channels: - term_rows = await term_channel( - db, + revision_pins=revision_pins, + ), + _run_channel( + term_channel, + enabled="term" in active_channels, user_id=user_id, namespace=namespace, query=query, @@ -125,7 +126,9 @@ async def bottom_discovery( allowed_chunk_types=allowed_chunk_types, signal_paths=signal_paths, filter_mode=filter_mode, - ) + revision_pins=revision_pins, + ), + ) default_weights = { "path": CHANNEL_WEIGHT_PATH, @@ -194,3 +197,20 @@ async def bottom_discovery( latency = int((time.monotonic() - t0) * 1000) logger.error(f" search.bottom_discovery failed: {exc}") return DiscoveryResult(status="error", error=str(exc), latency_ms=latency) + + +async def _run_channel( + channel: Callable[..., Awaitable[list[dict[str, Any]]]], + *, + enabled: bool, + **kwargs: Any, +) -> list[dict[str, Any]]: + if not enabled: + return [] + + # Import lazily so discovery remains usable by lightweight contract tests + # without creating a database context until a channel is actually enabled. + from shared.core.database import get_db_context + + async with get_db_context() as channel_db: + return await channel(channel_db, **kwargs) diff --git a/packages/shared-python/shared/services/retrieval/search/ranking.py b/packages/shared-python/shared/services/retrieval/search/ranking.py index 392aa1f1..c4906e46 100644 --- a/packages/shared-python/shared/services/retrieval/search/ranking.py +++ b/packages/shared-python/shared/services/retrieval/search/ranking.py @@ -1,6 +1,7 @@ from __future__ import annotations import math +from collections.abc import Mapping from typing import Any from loguru import logger @@ -15,9 +16,9 @@ def get_candidate_key(row: dict[str, Any]) -> str: path = get_row_path(row) if path: - return f'path:{path}' - chunk_id = str(row.get('chunk_id') or '').strip() - return f'chunk:{chunk_id}' if chunk_id else '' + return f"path:{path}" + chunk_id = str(row.get("chunk_id") or "").strip() + return f"chunk:{chunk_id}" if chunk_id else "" async def load_chunk_importance_scores( @@ -26,12 +27,11 @@ async def load_chunk_importance_scores( user_id: str, namespace: str, rows: list[dict[str, Any]], + revision_pins: Mapping[str, str] | None = None, ) -> dict[str, float]: - chunk_ids = sorted({ - str(row.get('chunk_id') or '').strip() - for row in rows - if row.get('chunk_id') - }) + chunk_ids = sorted( + {str(row.get("chunk_id") or "").strip() for row in rows if row.get("chunk_id")} + ) if not chunk_ids: return {} stmt = ( @@ -43,22 +43,31 @@ async def load_chunk_importance_scores( ) .where(RetrievalHitStat.user_id == user_id) .where(RetrievalHitStat.namespace == namespace) - .where(RetrievalHitStat.hit_kind == 'chunk') + .where(RetrievalHitStat.hit_kind == "chunk") .where(RetrievalHitStat.chunk_id.in_(chunk_ids)) ) + if revision_pins is not None: + pinned_document_ids = {str(document_id) for document_id in revision_pins} + stmt = stmt.where(RetrievalHitStat.document_id.in_(pinned_document_ids)) result = await db.execute(stmt) importance_scores: dict[str, float] = {} for chunk_id, hit_count, last_hit_at, created_at in result.all(): if not chunk_id: continue - importance_scores[str(chunk_id)] = compute_importance_score(hit_count, last_hit_at, created_at) + score = compute_importance_score(hit_count, last_hit_at, created_at) + importance_scores[str(chunk_id)] = score + for row in rows: + if str(row.get("chunk_id") or "") == str(chunk_id): + document_id = str(row.get("document_id") or "") + if document_id: + importance_scores[f"{document_id}:{chunk_id}"] = score return importance_scores def apply_importance_multiplier( rows: list[dict[str, Any]], *, - raw_field: str = 'importance_raw_score', + raw_field: str = "importance_raw_score", low: float = 0.1, high: float = 2.0, ) -> None: @@ -67,7 +76,11 @@ def apply_importance_multiplier( values = sorted(float(row.get(raw_field, 0.0) or 0.0) for row in rows) item_count = len(values) - median = values[item_count // 2] if item_count % 2 else (values[item_count // 2 - 1] + values[item_count // 2]) / 2 + median = ( + values[item_count // 2] + if item_count % 2 + else (values[item_count // 2 - 1] + values[item_count // 2]) / 2 + ) q1 = values[item_count // 4] if item_count >= 4 else values[0] q3 = values[3 * item_count // 4] if item_count >= 4 else values[-1] iqr = q3 - q1 @@ -80,13 +93,13 @@ def apply_importance_multiplier( z_score = (raw_score - median) / iqr sigmoid_score = 1.0 / (1.0 + math.exp(-z_score)) multiplier = low + (high - low) * sigmoid_score - row['importance_multiplier'] = round(multiplier, 4) - row['agent_score'] = round( - float(row.get('agent_score', 0.0) or 0.0) * multiplier, + row["importance_multiplier"] = round(multiplier, 4) + row["agent_score"] = round( + float(row.get("agent_score", 0.0) or 0.0) * multiplier, 6, ) - row['discovery_score'] = round( - float(row.get('discovery_score', 0.0) or 0.0) * multiplier, + row["discovery_score"] = round( + float(row.get("discovery_score", 0.0) or 0.0) * multiplier, 6, ) @@ -107,9 +120,9 @@ def rank_candidates_by_path( if not key: continue candidate = dict(row) - candidate['discovery_score'] = float(row.get('discovery_score', 0.0) or 0.0) - candidate['agent_score'] = 0.0 - candidate.setdefault('hydrate_mode', 'chunks') + candidate["discovery_score"] = float(row.get("discovery_score", 0.0) or 0.0) + candidate["agent_score"] = 0.0 + candidate.setdefault("hydrate_mode", "chunks") merged[key] = candidate insertion_order[key] = counter counter += 1 @@ -118,25 +131,33 @@ def rank_candidates_by_path( key = get_candidate_key(row) if not key: continue - routed_agent_score = float(row.get('agent_score', 0.0) or 0.0) + routed_agent_score = float(row.get("agent_score", 0.0) or 0.0) if key not in merged: candidate = dict(row) - candidate['discovery_score'] = float(row.get('discovery_score', 0.0) or 0.0) - candidate['agent_score'] = routed_agent_score + candidate["discovery_score"] = float(row.get("discovery_score", 0.0) or 0.0) + candidate["agent_score"] = routed_agent_score merged[key] = candidate insertion_order[key] = counter counter += 1 continue candidate = merged[key] - candidate['agent_score'] = max(float(candidate.get('agent_score', 0.0) or 0.0), routed_agent_score) - if not candidate.get('source_chunk_path') and row.get('source_chunk_path'): - candidate['source_chunk_path'] = row.get('source_chunk_path') - if not candidate.get('section_path') and row.get('section_path'): - candidate['section_path'] = row.get('section_path') + candidate["agent_score"] = max( + float(candidate.get("agent_score", 0.0) or 0.0), routed_agent_score + ) + if not candidate.get("source_chunk_path") and row.get("source_chunk_path"): + candidate["source_chunk_path"] = row.get("source_chunk_path") + if not candidate.get("section_path") and row.get("section_path"): + candidate["section_path"] = row.get("section_path") for row in merged.values(): - row['importance_raw_score'] = float( - (importance_scores or {}).get(str(row.get('chunk_id') or ''), 0.0) or 0.0 + document_id = str(row.get("document_id") or "") + chunk_id = str(row.get("chunk_id") or "") + row["importance_raw_score"] = float( + (importance_scores or {}).get( + f"{document_id}:{chunk_id}", + (importance_scores or {}).get(chunk_id, 0.0), + ) + or 0.0 ) apply_importance_multiplier(list(merged.values())) @@ -145,11 +166,13 @@ def rank_candidates_by_path( fallback_rows: list[dict[str, Any]] = [] for key, row in merged.items(): - agent_score = float(row.get('agent_score', 0.0) or 0.0) - discovery_score = float(row.get('discovery_score', 0.0) or 0.0) - row['evidence_score'] = round(agent_score if has_agent_results else max(discovery_score, agent_score), 6) - row['score'] = row['evidence_score'] - row['_candidate_order'] = insertion_order[key] + agent_score = float(row.get("agent_score", 0.0) or 0.0) + discovery_score = float(row.get("discovery_score", 0.0) or 0.0) + row["evidence_score"] = round( + agent_score if has_agent_results else max(discovery_score, agent_score), 6 + ) + row["score"] = row["evidence_score"] + row["_candidate_order"] = insertion_order[key] if has_agent_results and agent_score <= 0.0: fallback_rows.append(row) @@ -158,9 +181,9 @@ def rank_candidates_by_path( def get_sort_key(row: dict[str, Any]) -> tuple[float, float, int]: return ( - float(row.get('agent_score', 0.0) or 0.0), - float(row.get('discovery_score', 0.0) or 0.0), - -int(row.get('_candidate_order', 0) or 0), + float(row.get("agent_score", 0.0) or 0.0), + float(row.get("discovery_score", 0.0) or 0.0), + -int(row.get("_candidate_order", 0) or 0), ) primary_rows.sort(key=get_sort_key, reverse=True) @@ -168,10 +191,10 @@ def get_sort_key(row: dict[str, Any]) -> tuple[float, float, int]: if len(ranked_rows) < top_k and fallback_rows: fallback_rows.sort(key=get_sort_key, reverse=True) - ranked_rows.extend(fallback_rows[:top_k - len(ranked_rows)]) + ranked_rows.extend(fallback_rows[: top_k - len(ranked_rows)]) for row in ranked_rows: - row.pop('_candidate_order', None) + row.pop("_candidate_order", None) return ranked_rows @@ -183,6 +206,7 @@ async def rank_retrieval_candidates( discovery_rows: list[dict[str, Any]], routed_rows: list[dict[str, Any]], top_k: int, + revision_pins: Mapping[str, str] | None = None, ) -> list[dict[str, Any]]: try: importance_scores = await load_chunk_importance_scores( @@ -190,9 +214,12 @@ async def rank_retrieval_candidates( user_id=user_id, namespace=namespace, rows=[*discovery_rows, *routed_rows], + revision_pins=revision_pins, ) except Exception as exc: - logger.warning(f'Failed to load chunk importance scores, continuing without importance: {exc}') + logger.warning( + f"Failed to load chunk importance scores, continuing without importance: {exc}" + ) importance_scores = {} return rank_candidates_by_path( discovery_rows, diff --git a/packages/shared-python/shared/services/retrieval/search/scoped_corpus.py b/packages/shared-python/shared/services/retrieval/search/scoped_corpus.py index 4a60e65b..ed6f2719 100644 --- a/packages/shared-python/shared/services/retrieval/search/scoped_corpus.py +++ b/packages/shared-python/shared/services/retrieval/search/scoped_corpus.py @@ -1,13 +1,72 @@ from __future__ import annotations +from collections.abc import Mapping from typing import Any -from sqlalchemy import func, select +from sqlalchemy import and_, func, select, tuple_ from sqlalchemy.ext.asyncio import AsyncSession +from sqlalchemy.exc import SQLAlchemyError -from shared.models.database.document import Document, DocumentChunk, DocumentSection +from shared.models.database.document import ( + Document, + DocumentChunk, + DocumentSection, + RetrievalServingRevisionManifest, +) from shared.models.database.job_result import JobResult from shared.services.retrieval.search.section_filters import is_excluded_section +from shared.services.retrieval.serving_manifest import decode_serving_manifest +from shared.services.retrieval.manifest_cache import cache_manifest_payloads + + +async def count_manifest_chunks( + db: AsyncSession, + *, + revision_pins: Mapping[str, str], +) -> int | None: + """Count chunks exactly from complete pinned manifests when available.""" + if not revision_pins: + return None + statement = select( + RetrievalServingRevisionManifest.document_id, + RetrievalServingRevisionManifest.job_result_id, + RetrievalServingRevisionManifest.payload_zlib, + RetrievalServingRevisionManifest.checksum, + RetrievalServingRevisionManifest.format_version, + ).where( + tuple_( + RetrievalServingRevisionManifest.document_id, + RetrievalServingRevisionManifest.job_result_id, + ).in_(list(revision_pins.items())) + ) + try: + rows = (await db.execute(statement)).all() + if len(rows) != len(revision_pins): + return None + total = 0 + decoded_payloads: dict[tuple[str, str], dict[str, Any]] = {} + for document_id, job_result_id, payload_zlib, checksum, format_version in rows: + payload = decode_serving_manifest( + bytes(payload_zlib), + checksum=str(checksum), + format_version=int(format_version), + ) + chunks = payload.get("chunks") + if not isinstance(chunks, list): + return None + total += len(chunks) + decoded_payloads[(str(document_id), str(job_result_id))] = payload + cache_manifest_payloads( + db, + revisions=revision_pins, + payloads=decoded_payloads, + ) + return total + except SQLAlchemyError: + await db.rollback() + return None + except (TypeError, ValueError, KeyError): + return None async def count_scoped_chunks( @@ -17,18 +76,32 @@ async def count_scoped_chunks( namespace: str, exclude_document_ids: list[str], allowed_chunk_types: set[str] | None, + revision_pins: Mapping[str, str] | None = None, ) -> int: - stmt = ( - select(func.count(DocumentChunk.id)) - .join( - Document, - (Document.document_id == DocumentChunk.document_id) - & (Document.current_job_result_id == DocumentChunk.job_result_id), + if revision_pins is None: + stmt = ( + select(func.count(DocumentChunk.id)) + .join( + Document, + (Document.document_id == DocumentChunk.document_id) + & (Document.current_job_result_id == DocumentChunk.job_result_id), + ) + .where(Document.user_id == user_id) + .where(Document.namespace == namespace) + .where(Document.status == 'active') + ) + else: + stmt = ( + select(func.count(DocumentChunk.id)) + .join(Document, Document.document_id == DocumentChunk.document_id) + .where(Document.user_id == user_id) + .where(Document.namespace == namespace) + .where( + tuple_(DocumentChunk.document_id, DocumentChunk.job_result_id).in_( + list(revision_pins.items()) + ) + ) ) - .where(Document.user_id == user_id) - .where(Document.namespace == namespace) - .where(Document.status == 'active') - ) if exclude_document_ids: stmt = stmt.where(Document.document_id.notin_(list(exclude_document_ids))) if allowed_chunk_types is not None: @@ -47,21 +120,31 @@ async def load_all_scoped_chunks( allowed_chunk_types: set[str] | None, signal_paths: list[str], filter_mode: str, + revision_pins: Mapping[str, str] | None = None, ) -> list[dict[str, Any]]: - stmt = ( - select(Document, DocumentChunk, DocumentSection, JobResult) - .join( - DocumentChunk, + if revision_pins is None: + chunk_join = ( (DocumentChunk.document_id == Document.document_id) - & (DocumentChunk.job_result_id == Document.current_job_result_id), + & (DocumentChunk.job_result_id == Document.current_job_result_id) + ) + else: + chunk_join = and_( + DocumentChunk.document_id == Document.document_id, + tuple_(DocumentChunk.document_id, DocumentChunk.job_result_id).in_( + list(revision_pins.items()) + ), ) + stmt = ( + select(Document, DocumentChunk, DocumentSection, JobResult) + .join(DocumentChunk, chunk_join) .outerjoin(DocumentSection, DocumentSection.section_id == DocumentChunk.section_id) .join(JobResult, JobResult.id == DocumentChunk.job_result_id) .where(Document.user_id == user_id) .where(Document.namespace == namespace) - .where(Document.status == 'active') .order_by(DocumentChunk.sort_order) ) + if revision_pins is None: + stmt = stmt.where(Document.status == 'active') if exclude_document_ids: stmt = stmt.where(Document.document_id.notin_(list(exclude_document_ids))) if allowed_chunk_types is not None: diff --git a/packages/shared-python/shared/services/retrieval/search/scoring.py b/packages/shared-python/shared/services/retrieval/search/scoring.py index 848a3ada..2746de34 100644 --- a/packages/shared-python/shared/services/retrieval/search/scoring.py +++ b/packages/shared-python/shared/services/retrieval/search/scoring.py @@ -51,14 +51,18 @@ def merge_channels_rrf( for channel_idx, channel_rows in enumerate(channels): weight = weights[channel_idx] if channel_idx < len(weights) else 1.0 - for rank, row in enumerate(channel_rows): + seen_chunk_ids: set[str] = set() + unique_rank = 0 + for row in channel_rows: chunk_id = str(row.get('chunk_id') or '') - if not chunk_id: + if not chunk_id or chunk_id in seen_chunk_ids: continue - rrf_score = weight / (k + rank + 1) + seen_chunk_ids.add(chunk_id) + rrf_score = weight / (k + unique_rank + 1) score_dict[chunk_id] = score_dict.get(chunk_id, 0.0) + rrf_score if chunk_id not in row_by_chunk_id: row_by_chunk_id[chunk_id] = row + unique_rank += 1 ranked = sorted(score_dict.items(), key=lambda x: x[1], reverse=True) results: list[dict[str, Any]] = [] diff --git a/packages/shared-python/shared/services/retrieval/serving_generation.py b/packages/shared-python/shared/services/retrieval/serving_generation.py new file mode 100644 index 00000000..5dc7baf5 --- /dev/null +++ b/packages/shared-python/shared/services/retrieval/serving_generation.py @@ -0,0 +1,60 @@ +"""Namespace generation locking for serving-state lifecycle updates.""" + +from __future__ import annotations + +from hashlib import sha256 + +from sqlalchemy import select +from sqlalchemy.dialects.postgresql import insert +from sqlalchemy.orm import Session + +from shared.models.database.document import RetrievalNamespaceGeneration + + +def lock_namespace_generation( + db: Session, + *, + user_id: str, + namespace: str, +) -> RetrievalNamespaceGeneration: + """Create if needed, then lock and return one namespace generation row.""" + generation_id = f"rng_{sha256(f'{user_id}:{namespace}'.encode()).hexdigest()}" + db.execute( + insert(RetrievalNamespaceGeneration) + .values( + id=generation_id, + user_id=user_id, + namespace=namespace, + generation=0, + ) + .on_conflict_do_nothing( + index_elements=[ + RetrievalNamespaceGeneration.user_id, + RetrievalNamespaceGeneration.namespace, + ] + ) + ) + generation = db.execute( + select(RetrievalNamespaceGeneration) + .where(RetrievalNamespaceGeneration.user_id == user_id) + .where(RetrievalNamespaceGeneration.namespace == namespace) + .with_for_update() + ).scalar_one() + return generation + + +def advance_namespace_generation( + db: Session, + *, + user_id: str, + namespace: str, +) -> int: + """Increment a locked namespace generation and return its new value.""" + generation = lock_namespace_generation( + db, + user_id=user_id, + namespace=namespace, + ) + generation.generation += 1 + db.flush() + return generation.generation diff --git a/packages/shared-python/shared/services/retrieval/serving_manifest.py b/packages/shared-python/shared/services/retrieval/serving_manifest.py new file mode 100644 index 00000000..b2dd40c5 --- /dev/null +++ b/packages/shared-python/shared/services/retrieval/serving_manifest.py @@ -0,0 +1,405 @@ +"""Versioned compression and integrity checks for serving manifests.""" + +from __future__ import annotations + +import hashlib +import json +import zlib +from typing import Any + +from sqlalchemy import delete, select +from sqlalchemy.orm import Session + +from shared.models.database.document import ( + Document, + DocumentChunk, + DocumentMapUnit, + DocumentMapUnitToken, + DocumentSection, + RetrievalNamespaceGeneration, + RetrievalNamespaceStat, + RetrievalNamespaceTokenStat, + RetrievalServingRevisionManifest, + RetrievalServingRevisionStat, +) +from shared.models.database.job_result import JobResult +from shared.services.retrieval.publication_models import DocumentPublicationScope + +SERVING_MANIFEST_FORMAT_VERSION = 1 + + +def build_revision_serving_payload( + db: Session, + *, + scope: DocumentPublicationScope, +) -> dict[str, Any]: + """Build ordered metadata for one published document revision.""" + document = db.execute( + select(Document).where(Document.document_id == scope.document_id) + ).scalar_one() + job_result = db.execute( + select(JobResult).where(JobResult.id == scope.job_result_id) + ).scalar_one() + sections = list( + db.scalars( + select(DocumentSection) + .where(DocumentSection.document_id == scope.document_id) + .where(DocumentSection.job_result_id == scope.job_result_id) + .order_by(DocumentSection.sort_order, DocumentSection.section_id) + ) + ) + chunks = list( + db.scalars( + select(DocumentChunk) + .where(DocumentChunk.document_id == scope.document_id) + .where(DocumentChunk.job_result_id == scope.job_result_id) + .order_by( + DocumentChunk.sort_order, DocumentChunk.chunk_id, DocumentChunk.id + ) + ) + ) + map_units = list( + db.scalars( + select(DocumentMapUnit) + .where(DocumentMapUnit.document_id == scope.document_id) + .where(DocumentMapUnit.job_result_id == scope.job_result_id) + .order_by(DocumentMapUnit.sort_order, DocumentMapUnit.unit_id) + ) + ) + section_path_by_id = { + section.section_id: section.section_path for section in sections + } + root_asset_ids = { + chunk.chunk_id + for chunk in chunks + if chunk.chunk_type in {"image", "table"} + and chunk.section_id is not None + and section_path_by_id.get(chunk.section_id) == "Root" + } + remounted_assets: dict[str, list[str]] = {} + for chunk in chunks: + if chunk.chunk_type != "text" or not isinstance(chunk.chunk_metadata, dict): + continue + connections = chunk.chunk_metadata.get("connect_to") + if not isinstance(connections, list): + continue + targets = [ + str(connection.get("target") or "").strip() + for connection in connections + if isinstance(connection, dict) + and str(connection.get("target") or "").strip() in root_asset_ids + ] + if targets: + remounted_assets[chunk.section_id or ""] = targets + + return { + "document_id": scope.document_id, + "job_result_id": scope.job_result_id, + "job_id": str(job_result.job_id), + "source_file_name": str( + document.source_file_name or scope.source_file_name or "" + ), + "sections": [ + { + "section_id": section.section_id, + "parent_section_id": section.parent_section_id, + "section_path": section.section_path, + "section_title": section.section_title, + "section_level": section.section_level, + "summary": section.summary, + "sort_order": section.sort_order, + } + for section in sections + ], + "chunks": [ + { + "chunk_id": chunk.chunk_id, + "section_id": chunk.section_id, + "chunk_type": chunk.chunk_type, + "sort_order": chunk.sort_order, + "connect_to": _connection_target_ids(chunk.chunk_metadata), + } + for chunk in chunks + ], + "map_units": [ + { + "row_id": unit.id, + "unit_id": unit.unit_id, + "section_id": unit.section_id, + "unit_kind": unit.unit_kind, + "path_token_count": unit.path_token_count, + "content_token_count": unit.content_token_count, + "sort_order": unit.sort_order, + } + for unit in map_units + ], + "root_asset_ids": sorted(root_asset_ids), + "remounted_assets_by_section": remounted_assets, + } + + +def build_revision_statistics_payload( + db: Session, + *, + scope: DocumentPublicationScope, +) -> dict[str, Any]: + """Build compressed scoring contributions for one revision.""" + units = list( + db.scalars( + select(DocumentMapUnit) + .where(DocumentMapUnit.document_id == scope.document_id) + .where(DocumentMapUnit.job_result_id == scope.job_result_id) + ) + ) + unit_ids = [unit.id for unit in units] + frequencies: dict[str, dict[str, int]] = {"path": {}, "content": {}} + unit_frequencies: dict[str, dict[str, dict[str, int]]] = {} + if unit_ids: + for map_unit_id, channel, token, frequency in db.execute( + select( + DocumentMapUnitToken.map_unit_id, + DocumentMapUnitToken.channel, + DocumentMapUnitToken.token, + DocumentMapUnitToken.frequency, + ).where(DocumentMapUnitToken.map_unit_id.in_(unit_ids)) + ).all(): + channel_key = str(channel) + if channel_key in frequencies: + token_key = str(token) + frequency_value = int(frequency) + frequencies[channel_key][token_key] = ( + frequencies[channel_key].get(token_key, 0) + frequency_value + ) + unit_frequencies.setdefault(str(map_unit_id), {}).setdefault( + channel_key, {} + )[token_key] = frequency_value + return { + "document_id": scope.document_id, + "job_result_id": scope.job_result_id, + "unit_count": len(units), + "path_token_count": sum(int(unit.path_token_count or 0) for unit in units), + "content_token_count": sum( + int(unit.content_token_count or 0) for unit in units + ), + "token_frequencies": frequencies, + "unit_frequencies": unit_frequencies, + } + + +def persist_revision_serving_state( + db: Session, + *, + scope: DocumentPublicationScope, +) -> None: + """Replace manifest and statistics rows for one revision atomically.""" + manifest_payload = build_revision_serving_payload(db, scope=scope) + statistics_payload = build_revision_statistics_payload(db, scope=scope) + manifest_bytes, manifest_checksum, manifest_version = encode_serving_manifest( + manifest_payload + ) + statistics_bytes, statistics_checksum, statistics_version = encode_serving_manifest( + statistics_payload + ) + db.execute( + delete(RetrievalServingRevisionManifest) + .where(RetrievalServingRevisionManifest.document_id == scope.document_id) + .where(RetrievalServingRevisionManifest.job_result_id == scope.job_result_id) + ) + db.execute( + delete(RetrievalServingRevisionStat) + .where(RetrievalServingRevisionStat.document_id == scope.document_id) + .where(RetrievalServingRevisionStat.job_result_id == scope.job_result_id) + ) + db.add( + RetrievalServingRevisionManifest( + user_id=scope.user_id, + namespace=scope.namespace, + document_id=scope.document_id, + job_result_id=scope.job_result_id, + format_version=manifest_version, + payload_zlib=manifest_bytes, + checksum=manifest_checksum, + ) + ) + db.add( + RetrievalServingRevisionStat( + user_id=scope.user_id, + namespace=scope.namespace, + document_id=scope.document_id, + job_result_id=scope.job_result_id, + format_version=statistics_version, + payload_zlib=statistics_bytes, + checksum=statistics_checksum, + ) + ) + + +def rebuild_namespace_serving_statistics( + db: Session, + *, + user_id: str, + namespace: str, +) -> int: + """Recompute namespace aggregates from active current revisions. + + Callers hold the namespace generation lock. The aggregate is prepared for + the generation that the caller will publish next. + """ + generation = db.execute( + select(RetrievalNamespaceGeneration) + .where(RetrievalNamespaceGeneration.user_id == user_id) + .where(RetrievalNamespaceGeneration.namespace == namespace) + .with_for_update() + ).scalar_one() + target_generation = int(generation.generation) + 1 + revisions = { + (str(document_id), str(job_result_id)) + for document_id, job_result_id in db.execute( + select(Document.document_id, Document.current_job_result_id) + .where(Document.user_id == user_id) + .where(Document.namespace == namespace) + .where(Document.status == "active") + .where(Document.current_job_result_id.is_not(None)) + ).all() + if document_id and job_result_id + } + aggregate: dict[str, Any] = { + "document_count": 0, + "unit_count": 0, + "path_token_count": 0, + "content_token_count": 0, + "token_frequencies": {"path": {}, "content": {}}, + } + document_frequencies: dict[tuple[str, str], int] = {} + for row in db.scalars( + select(RetrievalServingRevisionStat) + .where(RetrievalServingRevisionStat.user_id == user_id) + .where(RetrievalServingRevisionStat.namespace == namespace) + ): + if (row.document_id, row.job_result_id) not in revisions: + continue + payload = decode_serving_manifest( + row.payload_zlib, + checksum=row.checksum, + format_version=row.format_version, + ) + aggregate["document_count"] += 1 + aggregate["unit_count"] += int(payload.get("unit_count", 0)) + aggregate["path_token_count"] += int(payload.get("path_token_count", 0)) + aggregate["content_token_count"] += int(payload.get("content_token_count", 0)) + token_frequencies = payload.get("token_frequencies", {}) + if not isinstance(token_frequencies, dict): + continue + for channel, values in token_frequencies.items(): + if channel not in aggregate["token_frequencies"] or not isinstance( + values, dict + ): + continue + for token, value in values.items(): + token_key = str(token) + aggregate["token_frequencies"][channel][token_key] = aggregate[ + "token_frequencies" + ][channel].get(token_key, 0) + int(value) + if int(value) > 0: + key = (str(channel), token_key) + document_frequencies[key] = document_frequencies.get(key, 0) + 1 + + encoded, checksum, _version = encode_serving_manifest(aggregate) + namespace_stat = db.execute( + select(RetrievalNamespaceStat) + .where(RetrievalNamespaceStat.user_id == user_id) + .where(RetrievalNamespaceStat.namespace == namespace) + ).scalar_one_or_none() + if namespace_stat is None: + db.add( + RetrievalNamespaceStat( + user_id=user_id, + namespace=namespace, + generation=target_generation, + payload_zlib=encoded, + checksum=checksum, + ) + ) + else: + namespace_stat.generation = target_generation + namespace_stat.payload_zlib = encoded + namespace_stat.checksum = checksum + db.execute( + delete(RetrievalNamespaceTokenStat) + .where(RetrievalNamespaceTokenStat.user_id == user_id) + .where(RetrievalNamespaceTokenStat.namespace == namespace) + ) + db.add_all( + [ + RetrievalNamespaceTokenStat( + user_id=user_id, + namespace=namespace, + generation=target_generation, + channel=channel, + token_hash=hashlib.sha256(token.encode("utf-8")).hexdigest(), + document_frequency=frequency, + ) + for (channel, token), frequency in document_frequencies.items() + ] + ) + db.flush() + return target_generation + + +def _connection_target_ids(metadata: Any) -> list[str]: + if not isinstance(metadata, dict): + return [] + connections = metadata.get("connect_to") + if not isinstance(connections, list): + return [] + return [ + target + for connection in connections + if isinstance(connection, dict) + for target in [str(connection.get("target") or "").strip()] + if target + ] + + +def encode_serving_manifest(payload: dict[str, Any]) -> tuple[bytes, str, int]: + """Return compressed canonical JSON, checksum, and format version.""" + canonical_payload = json.dumps( + payload, + ensure_ascii=False, + separators=(",", ":"), + sort_keys=True, + ).encode("utf-8") + checksum = hashlib.sha256(canonical_payload).hexdigest() + return ( + zlib.compress(canonical_payload), + checksum, + SERVING_MANIFEST_FORMAT_VERSION, + ) + + +def decode_serving_manifest( + payload_zlib: bytes, + *, + checksum: str, + format_version: int, +) -> dict[str, Any]: + """Validate and decode one persisted serving manifest.""" + if format_version != SERVING_MANIFEST_FORMAT_VERSION: + raise ValueError(f"unsupported serving manifest version: {format_version}") + + try: + canonical_payload = zlib.decompress(payload_zlib) + except zlib.error as exc: + raise ValueError("invalid serving manifest compression") from exc + + actual_checksum = hashlib.sha256(canonical_payload).hexdigest() + if actual_checksum != checksum: + raise ValueError("serving manifest checksum mismatch") + + try: + decoded = json.loads(canonical_payload.decode("utf-8")) + except (UnicodeDecodeError, json.JSONDecodeError) as exc: + raise ValueError("invalid serving manifest JSON") from exc + if not isinstance(decoded, dict): + raise ValueError("serving manifest payload must be an object") + return decoded From dabb25d10da4c47be6d9e06713c2d1cb1bafe93e Mon Sep 17 00:00:00 2001 From: suguanYang Date: Sun, 30 Aug 2026 12:58:01 +0800 Subject: [PATCH 06/19] test: patch active database module in revision race contract --- .../contract/test_retrieval_revision_races_contract.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/apps/api/tests/contract/test_retrieval_revision_races_contract.py b/apps/api/tests/contract/test_retrieval_revision_races_contract.py index dd73d966..3ebb0232 100644 --- a/apps/api/tests/contract/test_retrieval_revision_races_contract.py +++ b/apps/api/tests/contract/test_retrieval_revision_races_contract.py @@ -2,6 +2,7 @@ from __future__ import annotations +import importlib from collections.abc import AsyncGenerator from contextlib import asynccontextmanager from typing import Any, cast @@ -40,7 +41,12 @@ async def fake_channel( observed.append((db, kwargs["revision_pins"])) return [] - monkeypatch.setattr("shared.core.database.get_db_context", fake_context) + # Import the active module explicitly. Some API contract fixtures reload + # shared.core.database between tests, leaving the package attribute pointed + # at an old module object; dotted-string patching can then miss the module + # used by discovery's lazy import. + database_module = importlib.import_module("shared.core.database") + monkeypatch.setattr(database_module, "get_db_context", fake_context) monkeypatch.setattr(discovery, "path_channel", fake_channel) monkeypatch.setattr(discovery, "content_channel", fake_channel) monkeypatch.setattr(discovery, "term_channel", fake_channel) From 2e8e4c08bb5a4212f407a3f05db040b913e9fb8f Mon Sep 17 00:00:00 2001 From: suguanYang Date: Mon, 31 Aug 2026 03:11:07 +0800 Subject: [PATCH 07/19] perf: avoid eager job chunk hydration --- .../services/retrieval/hydration/reference.py | 85 ++++++++++--------- 1 file changed, 46 insertions(+), 39 deletions(-) diff --git a/packages/shared-python/shared/services/retrieval/hydration/reference.py b/packages/shared-python/shared/services/retrieval/hydration/reference.py index cd240654..753fc66f 100644 --- a/packages/shared-python/shared/services/retrieval/hydration/reference.py +++ b/packages/shared-python/shared/services/retrieval/hydration/reference.py @@ -28,10 +28,10 @@ async def hydrate_referenced_chunk_rows( ref_keys = [ build_reference_lookup_key( - document_id=ref.get('document_id'), - chunk_id=ref.get('chunk_id'), - section_path=ref.get('section_path'), - file_path=ref.get('file_path'), + document_id=ref.get("document_id"), + chunk_id=ref.get("chunk_id"), + section_path=ref.get("section_path"), + file_path=ref.get("file_path"), ) for ref in refs ] @@ -42,15 +42,16 @@ async def hydrate_referenced_chunk_rows( document_ids = sorted({document_id for document_id, _, _, _ in ref_keys}) chunk_ids = sorted({chunk_id for _, chunk_id, _, _ in ref_keys}) pinned_document_ids = [ - document_id for document_id in document_ids if revision_pins and document_id in revision_pins + document_id + for document_id in document_ids + if revision_pins and document_id in revision_pins ] if revision_pins is not None and not pinned_document_ids: return [] if revision_pins is None: - chunk_join = ( - (DocumentChunk.document_id == Document.document_id) - & (DocumentChunk.job_result_id == Document.current_job_result_id) + chunk_join = (DocumentChunk.document_id == Document.document_id) & ( + DocumentChunk.job_result_id == Document.current_job_result_id ) else: chunk_join = and_( @@ -67,9 +68,15 @@ async def hydrate_referenced_chunk_rows( ) stmt = ( - select(Document, DocumentChunk, DocumentSection, JobResult) + # Select only the job identifier needed for the public projection. + # Selecting the JobResult entity would trigger its ``chunks`` selectin + # relationship, loading the entire legacy job-chunk collection for + # every referenced revision during final hydration. + select(Document, DocumentChunk, DocumentSection, JobResult.job_id) .join(DocumentChunk, chunk_join) - .outerjoin(DocumentSection, DocumentSection.section_id == DocumentChunk.section_id) + .outerjoin( + DocumentSection, DocumentSection.section_id == DocumentChunk.section_id + ) .join(JobResult, JobResult.id == DocumentChunk.job_result_id) .where(Document.user_id == user_id) .where(Document.namespace == namespace) @@ -82,40 +89,40 @@ async def hydrate_referenced_chunk_rows( .order_by(DocumentChunk.sort_order) ) if revision_pins is None: - stmt = stmt.where(Document.status == 'active') + stmt = stmt.where(Document.status == "active") result = await db.execute(stmt) rows_by_key: dict[ReferenceLookupKey, dict[str, Any]] = {} rows_by_base_key: dict[tuple[str, str], list[dict[str, Any]]] = {} - for document, chunk, section, job_result in result.all(): + for document, chunk, section, job_id in result.all(): row = { - 'document_id': document.document_id, - 'chunk_id': chunk.chunk_id, - 'section_id': chunk.section_id, - 'section_path': section.section_path if section else None, - 'source_file_name': document.source_file_name, - 'chunk_type': chunk.chunk_type, - 'content': chunk.content, + "document_id": document.document_id, + "chunk_id": chunk.chunk_id, + "section_id": chunk.section_id, + "section_path": section.section_path if section else None, + "source_file_name": document.source_file_name, + "chunk_type": chunk.chunk_type, + "content": chunk.content, # Use the caller-supplied score when available (e.g. discovery RRF or # KG confidence). None signals "no score known" so consumers can # distinguish unscored chunks from genuinely high-scoring ones. - 'score': ( + "score": ( score_by_chunk_id.get(chunk.chunk_id) if score_by_chunk_id is not None else None ), - 'file_path': chunk.file_path, - 'chunk_metadata': chunk.chunk_metadata or {}, - 'job_result_id': chunk.job_result_id, - 'job_id': job_result.job_id if job_result else None, - 'source_chunk_path': chunk.source_chunk_path, - 'sort_order': chunk.sort_order, + "file_path": chunk.file_path, + "chunk_metadata": chunk.chunk_metadata or {}, + "job_result_id": chunk.job_result_id, + "job_id": job_id, + "source_chunk_path": chunk.source_chunk_path, + "sort_order": chunk.sort_order, } key = build_reference_lookup_key( - document_id=row['document_id'], - chunk_id=row['chunk_id'], - section_path=row['section_path'], - file_path=row['file_path'], + document_id=row["document_id"], + chunk_id=row["chunk_id"], + section_path=row["section_path"], + file_path=row["file_path"], ) rows_by_key[key] = row rows_by_base_key.setdefault((key[0], key[1]), []).append(row) @@ -131,7 +138,7 @@ async def hydrate_referenced_chunk_rows( candidate for candidate in candidates if key[2] - and str(candidate.get('section_path') or '').strip() == key[2] + and str(candidate.get("section_path") or "").strip() == key[2] ), None, ) @@ -141,10 +148,10 @@ async def hydrate_referenced_chunk_rows( candidate for candidate in candidates if build_reference_lookup_key( - document_id=candidate.get('document_id'), - chunk_id=candidate.get('chunk_id'), - section_path=candidate.get('section_path'), - file_path=candidate.get('file_path'), + document_id=candidate.get("document_id"), + chunk_id=candidate.get("chunk_id"), + section_path=candidate.get("section_path"), + file_path=candidate.get("file_path"), ) not in seen_keys ), @@ -152,10 +159,10 @@ async def hydrate_referenced_chunk_rows( ) if row is not None: row_key = build_reference_lookup_key( - document_id=row.get('document_id'), - chunk_id=row.get('chunk_id'), - section_path=row.get('section_path'), - file_path=row.get('file_path'), + document_id=row.get("document_id"), + chunk_id=row.get("chunk_id"), + section_path=row.get("section_path"), + file_path=row.get("file_path"), ) if row_key in seen_keys: continue From cc67a90f88afb58a084ca2df95329b2dad52cd91 Mon Sep 17 00:00:00 2001 From: suguanYang Date: Mon, 31 Aug 2026 09:55:09 +0800 Subject: [PATCH 08/19] perf: bound retrieval snapshot payload reads --- .../services/retrieval/execution/routes.py | 31 +++------- .../shared/services/retrieval/nav_snapshot.py | 61 +++++++++++-------- .../retrieval/search/scoped_corpus.py | 13 +++- 3 files changed, 56 insertions(+), 49 deletions(-) diff --git a/packages/shared-python/shared/services/retrieval/execution/routes.py b/packages/shared-python/shared/services/retrieval/execution/routes.py index d2b6d211..0998a8df 100644 --- a/packages/shared-python/shared/services/retrieval/execution/routes.py +++ b/packages/shared-python/shared/services/retrieval/execution/routes.py @@ -23,7 +23,6 @@ ) from shared.services.retrieval.search.ranking import rank_retrieval_candidates from shared.services.retrieval.search.scoped_corpus import ( - count_manifest_chunks, count_scoped_chunks, load_all_scoped_chunks, ) @@ -58,27 +57,15 @@ async def _try_run_small_corpus_route( context: RetrievalRouteContext, ) -> RetrievalRouteOutcome | None: total_chunk_count: int | None = None - if ( - context.use_agentic is not False - and context.revision_pins is not None - and not context.exclude_document_ids - and not context.exclude_sections - and context.allowed_chunk_types is None - and not context.signal_paths - ): - total_chunk_count = await count_manifest_chunks( - context.db, - revision_pins=context.revision_pins, - ) - if total_chunk_count is None: - total_chunk_count = await count_scoped_chunks( - context.db, - user_id=context.user_id, - namespace=context.namespace, - exclude_document_ids=context.exclude_document_ids, - allowed_chunk_types=context.allowed_chunk_types, - revision_pins=context.revision_pins, - ) + total_chunk_count = await count_scoped_chunks( + context.db, + user_id=context.user_id, + namespace=context.namespace, + exclude_document_ids=context.exclude_document_ids, + allowed_chunk_types=context.allowed_chunk_types, + revision_pins=context.revision_pins, + max_count=context.top_k + 1, + ) logger.info(f"\n Total chunks in scope: {total_chunk_count}") if total_chunk_count > context.top_k: diff --git a/packages/shared-python/shared/services/retrieval/nav_snapshot.py b/packages/shared-python/shared/services/retrieval/nav_snapshot.py index c68d6d8c..5d34e451 100644 --- a/packages/shared-python/shared/services/retrieval/nav_snapshot.py +++ b/packages/shared-python/shared/services/retrieval/nav_snapshot.py @@ -55,6 +55,12 @@ # Keep revision predicates bounded while reducing round trips for large # namespaces. Keyset paging still caps each payload query at 10,000 rows. _REVISION_GROUP_SIZE = 64 +# Compressed manifests are efficient for small namespaces, but transferring a +# large set of binary payloads can exceed the asyncpg statement timeout before +# PostgreSQL has done meaningful work. For larger snapshots, the normalized +# section/chunk index loaders below transfer only the fields map-nav needs and +# page them by revision. +_MANIFEST_MAX_REVISION_COUNT = 32 _logger = logging.getLogger(__name__) @@ -206,12 +212,14 @@ async def load_nav_snapshot( if job_result_id and job_id } - manifest_sections = await _load_manifest_sections( - db, - document_revisions=document_revisions, - exclude_sections=excluded_secs, - job_id_by_result_id=job_id_by_result_id, - ) + manifest_sections = None + if len(document_revisions) <= _MANIFEST_MAX_REVISION_COUNT: + manifest_sections = await _load_manifest_sections( + db, + document_revisions=document_revisions, + exclude_sections=excluded_secs, + job_id_by_result_id=job_id_by_result_id, + ) if manifest_sections is None: sections_by_doc, section_path_by_id = await _load_sections( db, @@ -358,26 +366,31 @@ async def _load_manifest_sections( if len(manifest_entries) != len(document_revisions): return None else: - statement = select( - RetrievalServingRevisionManifest.document_id, - RetrievalServingRevisionManifest.job_result_id, - RetrievalServingRevisionManifest.payload_zlib, - RetrievalServingRevisionManifest.checksum, - RetrievalServingRevisionManifest.format_version, - ).where( - tuple_( + manifest_entries = [] + for group_start in range(0, len(document_revisions), _REVISION_GROUP_SIZE): + revision_group = document_revisions[ + group_start : group_start + _REVISION_GROUP_SIZE + ] + statement = select( RetrievalServingRevisionManifest.document_id, RetrievalServingRevisionManifest.job_result_id, - ).in_(document_revisions) - ) - try: - rows = (await db.execute(statement)).all() - except Exception: - await db.rollback() - return None - if len(rows) != len(document_revisions): - return None - manifest_entries = [tuple(row) for row in rows] + RetrievalServingRevisionManifest.payload_zlib, + RetrievalServingRevisionManifest.checksum, + RetrievalServingRevisionManifest.format_version, + ).where( + tuple_( + RetrievalServingRevisionManifest.document_id, + RetrievalServingRevisionManifest.job_result_id, + ).in_(revision_group) + ) + try: + rows = (await db.execute(statement)).all() + except Exception: + await db.rollback() + return None + if len(rows) != len(revision_group): + return None + manifest_entries.extend(tuple(row) for row in rows) by_doc: dict[str, list[SectionRow]] = {} path_by_id: dict[str, str] = {} ids_by_doc: dict[str, list[str]] = {} diff --git a/packages/shared-python/shared/services/retrieval/search/scoped_corpus.py b/packages/shared-python/shared/services/retrieval/search/scoped_corpus.py index ed6f2719..0e087549 100644 --- a/packages/shared-python/shared/services/retrieval/search/scoped_corpus.py +++ b/packages/shared-python/shared/services/retrieval/search/scoped_corpus.py @@ -77,10 +77,11 @@ async def count_scoped_chunks( exclude_document_ids: list[str], allowed_chunk_types: set[str] | None, revision_pins: Mapping[str, str] | None = None, + max_count: int | None = None, ) -> int: if revision_pins is None: stmt = ( - select(func.count(DocumentChunk.id)) + select(DocumentChunk.id) .join( Document, (Document.document_id == DocumentChunk.document_id) @@ -92,7 +93,7 @@ async def count_scoped_chunks( ) else: stmt = ( - select(func.count(DocumentChunk.id)) + select(DocumentChunk.id) .join(Document, Document.document_id == DocumentChunk.document_id) .where(Document.user_id == user_id) .where(Document.namespace == namespace) @@ -106,7 +107,13 @@ async def count_scoped_chunks( stmt = stmt.where(Document.document_id.notin_(list(exclude_document_ids))) if allowed_chunk_types is not None: stmt = stmt.where(func.lower(DocumentChunk.chunk_type).in_(list(allowed_chunk_types))) - result = await db.execute(stmt) + + if max_count is not None: + stmt = stmt.limit(max_count) + + result = await db.execute( + select(func.count()).select_from(stmt.order_by(None).subquery()) + ) return result.scalar() or 0 From b683dede3166970d873534f281bd0b6d04b8e512 Mon Sep 17 00:00:00 2001 From: suguanYang Date: Mon, 31 Aug 2026 10:03:44 +0800 Subject: [PATCH 09/19] test: lock large snapshot manifest bypass --- .../test_retrieval_snapshot_large_corpus_contract.py | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/apps/api/tests/contract/test_retrieval_snapshot_large_corpus_contract.py b/apps/api/tests/contract/test_retrieval_snapshot_large_corpus_contract.py index a414e44b..0b45a703 100644 --- a/apps/api/tests/contract/test_retrieval_snapshot_large_corpus_contract.py +++ b/apps/api/tests/contract/test_retrieval_snapshot_large_corpus_contract.py @@ -6,6 +6,7 @@ from uuid import uuid4 from httpx import AsyncClient +import pytest from shared.models.database.document import Document, DocumentChunk, DocumentSection from shared.models.database.job_result import JobResult @@ -15,6 +16,7 @@ _REVISION_GROUP_SIZE, load_nav_snapshot, ) +import shared.services.retrieval.nav_snapshot as nav_snapshot_module from sqlalchemy import Executable, Result, select from sqlalchemy.engine import Row from sqlalchemy.sql.selectable import Select @@ -252,10 +254,19 @@ async def test_large_snapshot_keeps_all_retrieval_inputs_after_bounded_sql_load( developer_api_client_factory: Callable[ [], AbstractAsyncContextManager[AsyncClient] ], + monkeypatch: pytest.MonkeyPatch, ) -> None: namespace = f"large-corpus-{uuid4().hex[:8]}" async with developer_api_client_factory(): await _seed_large_retrieval_corpus(namespace) + async def unexpected_manifest_load(*_args: object, **_kwargs: object) -> None: + raise AssertionError("large snapshots must use normalized retrieval rows") + + monkeypatch.setattr( + nav_snapshot_module, + "_load_manifest_sections", + unexpected_manifest_load, + ) legacy_rows = await _load_legacy_rows(namespace) async with contract_db_session() as db: counting_db = _CountingSession(db) From 93d9eeec1ae6044f2c64fbf428f49c0114050bde Mon Sep 17 00:00:00 2001 From: suguanYang Date: Mon, 31 Aug 2026 16:00:40 +0800 Subject: [PATCH 10/19] fix(worker): cap internal parser filenames --- .../support/filename_limits.py | 28 +++++++++++++++ .../support/internal_parse_name.py | 6 +++- .../tests/unit/test_internal_parse_name.py | 34 +++++++++++++++++++ 3 files changed, 67 insertions(+), 1 deletion(-) create mode 100644 apps/worker/app/services/document_parser/support/filename_limits.py create mode 100644 apps/worker/tests/unit/test_internal_parse_name.py diff --git a/apps/worker/app/services/document_parser/support/filename_limits.py b/apps/worker/app/services/document_parser/support/filename_limits.py new file mode 100644 index 00000000..21597c22 --- /dev/null +++ b/apps/worker/app/services/document_parser/support/filename_limits.py @@ -0,0 +1,28 @@ +"""Filename safety helpers for parser working files.""" + +import os +from hashlib import sha256 + + +MAX_INTERNAL_FILENAME_BYTES = 240 +INTERNAL_FILENAME_HASH_LENGTH = 12 + + +def truncate_internal_filename(filename: str) -> str: + """Keep parser filenames below Linux NAME_MAX while retaining identity.""" + if len(os.fsencode(filename)) <= MAX_INTERNAL_FILENAME_BYTES: + return filename + + name_root, name_ext = os.path.splitext(filename) + filename_hash = sha256(filename.encode("utf-8")).hexdigest()[ + :INTERNAL_FILENAME_HASH_LENGTH + ] + suffix = f"-{filename_hash}{name_ext}" + available_root_bytes = MAX_INTERNAL_FILENAME_BYTES - len(os.fsencode(suffix)) + if available_root_bytes <= 0: + return f"document-{filename_hash}.bin" + + truncated_root = name_root.encode("utf-8")[:available_root_bytes].decode( + "utf-8", errors="ignore" + ) + return f"{truncated_root}{suffix}" diff --git a/apps/worker/app/services/document_parser/support/internal_parse_name.py b/apps/worker/app/services/document_parser/support/internal_parse_name.py index 568dc13c..7a59ed8c 100644 --- a/apps/worker/app/services/document_parser/support/internal_parse_name.py +++ b/apps/worker/app/services/document_parser/support/internal_parse_name.py @@ -5,6 +5,9 @@ from dataclasses import dataclass from app.services.common.file_utils import path_handle +from app.services.document_parser.support.filename_limits import ( + truncate_internal_filename, +) @dataclass(frozen=True) @@ -39,7 +42,7 @@ def normalize_internal_parse_name( effective_root = name_root or "document" internal_name = f"{effective_root}{effective_ext}" - return ( + normalized_name = ( internal_name.replace("(", "-") .replace(")", "-") .replace("[", "-") @@ -53,6 +56,7 @@ def normalize_internal_parse_name( .replace(chr(0x2015), "-") .replace(chr(0x2212), "-") ) + return truncate_internal_filename(normalized_name) def prepare_internal_parse_input( diff --git a/apps/worker/tests/unit/test_internal_parse_name.py b/apps/worker/tests/unit/test_internal_parse_name.py new file mode 100644 index 00000000..c43dcabb --- /dev/null +++ b/apps/worker/tests/unit/test_internal_parse_name.py @@ -0,0 +1,34 @@ +from __future__ import annotations + +import os +from pathlib import Path + +from app.services.document_parser.support.internal_parse_name import ( + prepare_internal_parse_input, +) + + +def test_prepare_internal_parse_input_handles_long_encoded_filename( + tmp_path, +) -> None: + temporary_file_path = tmp_path / "temporary.pdf" + temporary_file_path.write_bytes(b"pdf") + encoded_filename = ( + "%E9%99%84%E4%BB%B65.%E5%8D%97%E4%BA%AC%E4%BF%A1%E6%81%AF%E5%B7%A5%E7%A8%8B" + * 20 + + ".pdf" + ) + + prepared_input = prepare_internal_parse_input( + str(temporary_file_path), + encoded_filename, + fallback_ext=".pdf", + prefer_fallback_ext=True, + ) + prepared_file_path = Path(prepared_input.file_path) + file_exists = prepared_file_path.exists() + file_contents = prepared_file_path.read_bytes() + + assert len(os.fsencode(prepared_input.internal_filename)) <= 240 + assert file_exists + assert file_contents == b"pdf" From 04099252b8f733cd3555ff51f77b26efb546fcf8 Mon Sep 17 00:00:00 2001 From: chengke <404835780@qq.com> Date: Mon, 31 Aug 2026 12:40:37 +0800 Subject: [PATCH 11/19] feat: add conversation_id to retrieval queries and enhance namespace map snapshot handling --- ...b_add_retrieval_namespace_map_snapshots.py | 39 ++++ apps/api/app/api/v1/routes/retrieval.py | 9 + .../services/documents/lifecycle_service.py | 11 + ...etrieval_lazy_snapshot_quality_contract.py | 24 +-- .../test_retrieval_map_unit_index_contract.py | 24 ++- .../test_retrieval_term_score_contract.py | 75 ------- .../shared/models/database/document.py | 32 +++ .../shared/services/retrieval/app_service.py | 2 + .../services/retrieval/execution/plan.py | 2 + .../retrieval/execution/query_request.py | 4 + .../retrieval/execution/route_types.py | 1 + .../services/retrieval/execution/routes.py | 4 +- .../retrieval/namespace_map_snapshot.py | 150 +++++++++++++ .../retrieval/namespace_map_snapshot_cache.py | 43 ++++ .../services/retrieval/nav/knowhere_hybrid.py | 46 +--- .../services/retrieval/nav/nav_knowhere.py | 107 ---------- .../shared/services/retrieval/nav_snapshot.py | 199 ++++++++++++++---- .../services/retrieval/publication_content.py | 6 +- .../services/retrieval/publication_service.py | 9 + .../services/retrieval/serving_manifest.py | 9 +- 20 files changed, 519 insertions(+), 277 deletions(-) create mode 100644 apps/api/alembic/versions/6c7d8e9f0a1b_add_retrieval_namespace_map_snapshots.py delete mode 100644 apps/api/tests/contract/test_retrieval_term_score_contract.py create mode 100644 packages/shared-python/shared/services/retrieval/namespace_map_snapshot.py create mode 100644 packages/shared-python/shared/services/retrieval/namespace_map_snapshot_cache.py diff --git a/apps/api/alembic/versions/6c7d8e9f0a1b_add_retrieval_namespace_map_snapshots.py b/apps/api/alembic/versions/6c7d8e9f0a1b_add_retrieval_namespace_map_snapshots.py new file mode 100644 index 00000000..1f7b7c01 --- /dev/null +++ b/apps/api/alembic/versions/6c7d8e9f0a1b_add_retrieval_namespace_map_snapshots.py @@ -0,0 +1,39 @@ +"""Add persisted namespace-level MAP snapshot table.""" + +from __future__ import annotations + +from collections.abc import Sequence + +from alembic import op +import sqlalchemy as sa + + +revision: str = "6c7d8e9f0a1b" +down_revision: str | None = "5b6c7d8e9f0a" +branch_labels: Sequence[str] | None = None +depends_on: Sequence[str] | None = None + + +def upgrade() -> None: + if not sa.inspect(op.get_bind()).has_table("retrieval_namespace_map_snapshots"): + op.create_table( + "retrieval_namespace_map_snapshots", + sa.Column("id", sa.String(length=100), nullable=False), + sa.Column("user_id", sa.Text(), nullable=False), + sa.Column("namespace", sa.String(length=255), nullable=False), + sa.Column("generation", sa.BigInteger(), nullable=False), + sa.Column("format_version", sa.Integer(), nullable=False), + sa.Column("payload_zlib", sa.LargeBinary(), nullable=False), + sa.Column("checksum", sa.String(length=64), nullable=False), + sa.Column("updated_at", sa.DateTime(), nullable=False), + sa.PrimaryKeyConstraint("id"), + sa.UniqueConstraint( + "user_id", + "namespace", + name="uq_retrieval_namespace_map_snapshots_scope", + ), + ) + + +def downgrade() -> None: + op.drop_table("retrieval_namespace_map_snapshots", if_exists=True) diff --git a/apps/api/app/api/v1/routes/retrieval.py b/apps/api/app/api/v1/routes/retrieval.py index 758f111e..3c9202df 100644 --- a/apps/api/app/api/v1/routes/retrieval.py +++ b/apps/api/app/api/v1/routes/retrieval.py @@ -77,6 +77,14 @@ class RetrievalQueryRequest(BaseModel): "Set false to force classic 3-channel top-K retrieval." ), ) + conversation_id: str | None = Field( + None, + max_length=255, + description=( + "Caller-supplied conversation identifier, threaded through for " + "retrieval tracing. Does not affect caching or result content." + ), + ) @field_validator("channels") @classmethod @@ -172,6 +180,7 @@ async def execute_retrieval_query( threshold=payload.threshold, internal_recall_k=payload.internal_recall_k, use_agentic=payload.use_agentic, + conversation_id=payload.conversation_id, llm_config=llm_config, ) diff --git a/apps/api/app/services/documents/lifecycle_service.py b/apps/api/app/services/documents/lifecycle_service.py index 6b39507d..3ea08f05 100644 --- a/apps/api/app/services/documents/lifecycle_service.py +++ b/apps/api/app/services/documents/lifecycle_service.py @@ -20,6 +20,9 @@ invalidate_retrieval_cache_namespaces, ) from shared.services.retrieval.graph.service import DocumentGraphService, GraphScope +from shared.services.retrieval.namespace_map_snapshot import ( + remove_document_from_namespace_map_snapshot, +) from shared.services.retrieval.serving_generation import ( advance_namespace_generation, lock_namespace_generation, @@ -483,6 +486,14 @@ async def archive_document( namespace=previous_namespace, ) ) + await db.run_sync( + lambda sync_db: remove_document_from_namespace_map_snapshot( + sync_db, + user_id=user_id, + namespace=previous_namespace, + document_id=document_id, + ) + ) await db.run_sync( lambda sync_db: advance_namespace_generation( sync_db, diff --git a/apps/api/tests/contract/test_retrieval_lazy_snapshot_quality_contract.py b/apps/api/tests/contract/test_retrieval_lazy_snapshot_quality_contract.py index 0752a3b9..a56260d5 100644 --- a/apps/api/tests/contract/test_retrieval_lazy_snapshot_quality_contract.py +++ b/apps/api/tests/contract/test_retrieval_lazy_snapshot_quality_contract.py @@ -220,7 +220,11 @@ def test_lazy_provider_preserves_score_units_and_scores() -> None: ] -def test_streaming_scorer_preserves_exact_eager_scores() -> None: +def test_streaming_scorer_preserves_exact_eager_scores(monkeypatch: Any) -> None: + # Corpus-wide map-nav scoring retired the term channel (see the + # unify-bm25-persistent-map plan); zero its weight so the eager oracle + # is directly comparable to the path+content-only streaming scorer. + monkeypatch.setenv("NAV_MAP_CHANNEL_WEIGHT_TERM", "0") rows: list[ScoreUnitRow] = [ { "chunk_id": "unit-a", @@ -257,7 +261,12 @@ def unit_factory() -> Sequence[ScoreUnitRow]: assert replay_count == 1 -def test_streaming_scorer_preserves_duplicate_id_eager_semantics() -> None: +def test_streaming_scorer_preserves_duplicate_id_eager_semantics( + monkeypatch: Any, +) -> None: + # See test_streaming_scorer_preserves_exact_eager_scores: term is retired + # from corpus-wide scoring, so the eager oracle's term weight is zeroed. + monkeypatch.setenv("NAV_MAP_CHANNEL_WEIGHT_TERM", "0") rows: list[ScoreUnitRow] = [ { "chunk_id": "duplicate", @@ -383,17 +392,6 @@ def build_stats(search_field: str) -> PersistedBm25Stats: token: str(row["content_search_text"]).split().count(token) for token in query_tokens }, - term_scores=tuple( - 100.0 - if query in str(row["term_search_text"]) - else float( - sum( - token in str(row["term_search_text"]) - for token in query.split() - ) - ) - for query in queries - ), ) for row in rows ], diff --git a/apps/api/tests/contract/test_retrieval_map_unit_index_contract.py b/apps/api/tests/contract/test_retrieval_map_unit_index_contract.py index eb124fb2..6d30731c 100644 --- a/apps/api/tests/contract/test_retrieval_map_unit_index_contract.py +++ b/apps/api/tests/contract/test_retrieval_map_unit_index_contract.py @@ -2,6 +2,7 @@ from collections.abc import Callable, Mapping, Sequence from contextlib import AbstractAsyncContextManager +from typing import Any from uuid import uuid4 from httpx import AsyncClient @@ -33,6 +34,7 @@ ) from shared.services.retrieval.nav_bridge import build_referenced_chunks from shared.services.retrieval.publication_models import DocumentPublicationScope +from shared.services.retrieval.serving_generation import lock_namespace_generation from tests.support.contract_database import ContractDatabase from tests.support.retrieval_snapshot_support import contract_db_session @@ -161,7 +163,7 @@ async def test_published_map_units_preserve_scores_without_chunk_payload_reads( ] async with contract_db_session() as db: await db.run_sync( - lambda sync_db: replace_document_revision_content( + lambda sync_db: _publish_revision_with_generation_lock( sync_db, scope=scope, chunks=chunks, @@ -346,7 +348,7 @@ async def test_lazy_snapshot_defers_selected_asset_reference_metadata( ] async with contract_db_session() as db: await db.run_sync( - lambda sync_db: replace_document_revision_content( + lambda sync_db: _publish_revision_with_generation_lock( sync_db, scope=scope, chunks=chunks, @@ -515,6 +517,24 @@ def test_titleless_leaf_has_identical_eager_and_lazy_path_scoring() -> None: ) +def _publish_revision_with_generation_lock( + sync_db: Any, + *, + scope: DocumentPublicationScope, + chunks: list[dict[str, Any]], +) -> None: + """Mirror the production publish sequence: lock, then patch the snapshot. + + ``replace_document_revision_content`` requires the namespace generation + row to already exist (``publication_service.py`` locks it before every + real publish); this test helper reproduces that precondition. + """ + lock_namespace_generation( + sync_db, user_id=scope.user_id, namespace=scope.namespace + ) + replace_document_revision_content(sync_db, scope=scope, chunks=chunks) + + async def _seed_revision( *, namespace: str, diff --git a/apps/api/tests/contract/test_retrieval_term_score_contract.py b/apps/api/tests/contract/test_retrieval_term_score_contract.py deleted file mode 100644 index 425f4de0..00000000 --- a/apps/api/tests/contract/test_retrieval_term_score_contract.py +++ /dev/null @@ -1,75 +0,0 @@ -"""Contracts for persisted map-unit term scoring.""" - -from __future__ import annotations - -from shared.services.retrieval.nav.nav_knowhere import ReadOnlyChunkStore - - -class _FakeCursor: - def __init__(self, rows: list[tuple[str, str]]) -> None: - self.rows = rows - self.statement = "" - self.parameters: object = None - - def execute(self, statement: str, parameters: object) -> None: - self.statement = statement - self.parameters = parameters - - def fetchall(self) -> list[tuple[str, str]]: - return self.rows - - -def _store() -> ReadOnlyChunkStore: - return ReadOnlyChunkStore.__new__(ReadOnlyChunkStore) - - -def test_term_scores_keep_literal_substring_and_token_hit_semantics() -> None: - cursor = _FakeCursor( - [ - ("unit-full", "prefix alpha beta suffix"), - ("unit-token", "prefix alpha gamma suffix"), - ] - ) - queries = ["alpha beta", "", "alpha"] - query_tokens = { - "alpha beta": ["alpha", "beta"], - "": [], - "alpha": ["alpha"], - } - - scores = _store()._load_term_scores( - cursor, # type: ignore[arg-type] - map_unit_ids=["unit-full", "unit-token", "unit-miss"], - queries=queries, - query_tokens_by_query=query_tokens, - ) - - assert scores == { - "unit-full": (100.0, 0.0, 100.0), - "unit-token": (1.0, 0.0, 100.0), - } - assert "LIKE ANY" in cursor.statement - assert "POSITION" not in cursor.statement - parameters = cursor.parameters - assert isinstance(parameters, tuple) - assert parameters[0] == ["unit-full", "unit-token", "unit-miss"] - assert parameters[1] == ["%alpha beta%", "%alpha%", "%beta%"] - - -def test_long_query_uses_constant_shape_candidate_sql() -> None: - cursor = _FakeCursor([]) - tokens = [f"token-{index}" for index in range(300)] - query = " ".join(tokens) - - _store()._load_term_scores( - cursor, # type: ignore[arg-type] - map_unit_ids=["unit-1"], - queries=[query], - query_tokens_by_query={query: tokens}, - ) - - assert cursor.statement.count("POSITION") == 0 - assert "LIKE ANY" not in cursor.statement - parameters = cursor.parameters - assert isinstance(parameters, tuple) - assert parameters == (["unit-1"],) diff --git a/packages/shared-python/shared/models/database/document.py b/packages/shared-python/shared/models/database/document.py index c396283d..b85dd84f 100644 --- a/packages/shared-python/shared/models/database/document.py +++ b/packages/shared-python/shared/models/database/document.py @@ -518,6 +518,38 @@ class RetrievalNamespaceTokenStat(Base): ) +class RetrievalNamespaceMapSnapshot(Base): + """Persisted namespace-level MAP (sections + chunk index + map units). + + Incrementally patched at publish/archive time (one document's subtree at + a time); query time reads this row directly instead of merging per-file + manifests. Overwritten in place -- no generation history is retained. + """ + + __tablename__ = "retrieval_namespace_map_snapshots" + + id: Mapped[str] = mapped_column( + String(100), primary_key=True, default=lambda: f"rnmap_{uuid4().hex}" + ) + user_id: Mapped[str] = mapped_column(Text, nullable=False) + namespace: Mapped[str] = mapped_column(String(255), nullable=False) + generation: Mapped[int] = mapped_column(BigInteger, nullable=False) + format_version: Mapped[int] = mapped_column(Integer, nullable=False) + payload_zlib: Mapped[bytes] = mapped_column(LargeBinary, nullable=False) + checksum: Mapped[str] = mapped_column(String(64), nullable=False) + updated_at: Mapped[datetime] = mapped_column( + DateTime, default=utc_now_naive, onupdate=utc_now_naive, nullable=False + ) + + __table_args__ = ( + UniqueConstraint( + "user_id", + "namespace", + name="uq_retrieval_namespace_map_snapshots_scope", + ), + ) + + class GraphNode(Base): """Persisted derived graph node used for routing and expansion.""" diff --git a/packages/shared-python/shared/services/retrieval/app_service.py b/packages/shared-python/shared/services/retrieval/app_service.py index 83ce65b7..69727d1e 100644 --- a/packages/shared-python/shared/services/retrieval/app_service.py +++ b/packages/shared-python/shared/services/retrieval/app_service.py @@ -32,6 +32,7 @@ async def run_retrieval_query( threshold: float = 0.0, internal_recall_k: int | None = None, use_agentic: bool | None = None, + conversation_id: str | None = None, llm_config: LLMConfig | None = None, ) -> dict[str, Any]: return await execute_retrieval_query( @@ -51,5 +52,6 @@ async def run_retrieval_query( threshold=threshold, internal_recall_k=internal_recall_k, use_agentic=use_agentic, + conversation_id=conversation_id, llm_config=llm_config, ) diff --git a/packages/shared-python/shared/services/retrieval/execution/plan.py b/packages/shared-python/shared/services/retrieval/execution/plan.py index adb2a6da..8e9a9c7f 100644 --- a/packages/shared-python/shared/services/retrieval/execution/plan.py +++ b/packages/shared-python/shared/services/retrieval/execution/plan.py @@ -48,6 +48,7 @@ async def run_retrieval_query( threshold: float = 0.0, internal_recall_k: int | None = None, use_agentic: bool | None = None, + conversation_id: str | None = None, llm_config: LLMConfig | None = None, ) -> dict[str, Any]: """Run retrieval through the plan module.""" @@ -69,6 +70,7 @@ async def run_retrieval_query( threshold=threshold, internal_recall_k=internal_recall_k, use_agentic=use_agentic, + conversation_id=conversation_id, llm_config=llm_config, ) ).execute() diff --git a/packages/shared-python/shared/services/retrieval/execution/query_request.py b/packages/shared-python/shared/services/retrieval/execution/query_request.py index 53363d18..9e831bce 100644 --- a/packages/shared-python/shared/services/retrieval/execution/query_request.py +++ b/packages/shared-python/shared/services/retrieval/execution/query_request.py @@ -31,6 +31,7 @@ class RetrievalQuery: threshold: float = 0.0 internal_recall_k: int | None = None use_agentic: bool | None = None + conversation_id: str | None = None llm_config: LLMConfig | None = None @classmethod @@ -53,6 +54,7 @@ def from_parameters( threshold: float = 0.0, internal_recall_k: int | None = None, use_agentic: bool | None = None, + conversation_id: str | None = None, llm_config: LLMConfig | None = None, ) -> "RetrievalQuery": return cls( @@ -72,6 +74,7 @@ def from_parameters( threshold=threshold, internal_recall_k=internal_recall_k, use_agentic=use_agentic, + conversation_id=conversation_id, llm_config=llm_config, ) @@ -125,4 +128,5 @@ def build_route_context(self) -> RetrievalRouteContext: internal_recall_k=self.internal_recall_k, effective_recall_k=self.resolve_effective_recall_k(), use_agentic=self.use_agentic, + conversation_id=self.conversation_id, ) diff --git a/packages/shared-python/shared/services/retrieval/execution/route_types.py b/packages/shared-python/shared/services/retrieval/execution/route_types.py index 13c21c46..33918090 100644 --- a/packages/shared-python/shared/services/retrieval/execution/route_types.py +++ b/packages/shared-python/shared/services/retrieval/execution/route_types.py @@ -28,6 +28,7 @@ class RetrievalRouteContext: internal_recall_k: int | None effective_recall_k: int use_agentic: bool | None + conversation_id: str | None = None revision_pins: RetrievalRevisionPins | None = None diff --git a/packages/shared-python/shared/services/retrieval/execution/routes.py b/packages/shared-python/shared/services/retrieval/execution/routes.py index 0998a8df..fe47be19 100644 --- a/packages/shared-python/shared/services/retrieval/execution/routes.py +++ b/packages/shared-python/shared/services/retrieval/execution/routes.py @@ -234,10 +234,12 @@ async def _run_mapnav_route( ) snapshot_seconds = time.perf_counter() - snapshot_started logger.info( - "retrieval mapnav stage=snapshot_load seconds={:.3f} documents={} refs={}".format( + "retrieval mapnav stage=snapshot_load seconds={:.3f} documents={} refs={} " + "conversation_id={}".format( snapshot_seconds, len(snapshot.document_ids), len(snapshot.chunk_ref_index), + context.conversation_id or "", ) ) diff --git a/packages/shared-python/shared/services/retrieval/namespace_map_snapshot.py b/packages/shared-python/shared/services/retrieval/namespace_map_snapshot.py new file mode 100644 index 00000000..bd8b7767 --- /dev/null +++ b/packages/shared-python/shared/services/retrieval/namespace_map_snapshot.py @@ -0,0 +1,150 @@ +"""Incrementally-patched namespace-level MAP snapshot (sections + chunk index). + +Callers must already hold the namespace generation lock (see +``serving_generation.lock_namespace_generation``) before calling either +function here, exactly as ``rebuild_namespace_serving_statistics`` requires. +Each call only touches one document's subtree; every other document's +subtree in the payload is left byte-for-byte unchanged. +""" + +from __future__ import annotations + +from typing import Any + +from sqlalchemy import select +from sqlalchemy.orm import Session + +from shared.models.database.document import ( + RetrievalNamespaceGeneration, + RetrievalNamespaceMapSnapshot, +) +from shared.services.retrieval.publication_models import DocumentPublicationScope +from shared.services.retrieval.serving_manifest import ( + decode_serving_manifest, + encode_serving_manifest, +) + + +def patch_namespace_map_snapshot( + db: Session, + *, + scope: DocumentPublicationScope, + manifest_payload: dict[str, Any], +) -> None: + """Replace one document's subtree in the namespace MAP snapshot.""" + row = _load_snapshot_row(db, user_id=scope.user_id, namespace=scope.namespace) + documents = _decode_documents(row) + documents[scope.document_id] = { + "job_result_id": manifest_payload.get("job_result_id"), + "job_id": manifest_payload.get("job_id"), + "source_file_name": manifest_payload.get("source_file_name"), + "sections": manifest_payload.get("sections") or [], + "chunks": manifest_payload.get("chunks") or [], + } + target_generation = _target_generation( + db, user_id=scope.user_id, namespace=scope.namespace + ) + _write_snapshot( + db, + row=row, + user_id=scope.user_id, + namespace=scope.namespace, + documents=documents, + target_generation=target_generation, + ) + + +def remove_document_from_namespace_map_snapshot( + db: Session, + *, + user_id: str, + namespace: str, + document_id: str, +) -> None: + """Drop one document's subtree from the namespace MAP snapshot (archive path).""" + row = _load_snapshot_row(db, user_id=user_id, namespace=namespace) + documents = _decode_documents(row) + if document_id not in documents: + return + del documents[document_id] + target_generation = _target_generation(db, user_id=user_id, namespace=namespace) + _write_snapshot( + db, + row=row, + user_id=user_id, + namespace=namespace, + documents=documents, + target_generation=target_generation, + ) + + +def _target_generation(db: Session, *, user_id: str, namespace: str) -> int: + """Namespace generation this snapshot is prepared for (current + 1). + + Mirrors ``rebuild_namespace_serving_statistics``: callers advance the + generation after this write, in the same transaction. + """ + generation = db.execute( + select(RetrievalNamespaceGeneration) + .where(RetrievalNamespaceGeneration.user_id == user_id) + .where(RetrievalNamespaceGeneration.namespace == namespace) + .with_for_update() + ).scalar_one() + return int(generation.generation) + 1 + + +def _load_snapshot_row( + db: Session, *, user_id: str, namespace: str +) -> RetrievalNamespaceMapSnapshot | None: + return db.execute( + select(RetrievalNamespaceMapSnapshot) + .where(RetrievalNamespaceMapSnapshot.user_id == user_id) + .where(RetrievalNamespaceMapSnapshot.namespace == namespace) + ).scalar_one_or_none() + + +def _decode_documents( + row: RetrievalNamespaceMapSnapshot | None, +) -> dict[str, dict[str, Any]]: + if row is None: + return {} + try: + payload = decode_serving_manifest( + row.payload_zlib, + checksum=row.checksum, + format_version=row.format_version, + ) + except ValueError: + return {} + documents = payload.get("documents") + return dict(documents) if isinstance(documents, dict) else {} + + +def _write_snapshot( + db: Session, + *, + row: RetrievalNamespaceMapSnapshot | None, + user_id: str, + namespace: str, + documents: dict[str, dict[str, Any]], + target_generation: int, +) -> None: + encoded, checksum, format_version = encode_serving_manifest( + {"documents": documents} + ) + if row is None: + db.add( + RetrievalNamespaceMapSnapshot( + user_id=user_id, + namespace=namespace, + generation=target_generation, + format_version=format_version, + payload_zlib=encoded, + checksum=checksum, + ) + ) + else: + row.generation = target_generation + row.format_version = format_version + row.payload_zlib = encoded + row.checksum = checksum diff --git a/packages/shared-python/shared/services/retrieval/namespace_map_snapshot_cache.py b/packages/shared-python/shared/services/retrieval/namespace_map_snapshot_cache.py new file mode 100644 index 00000000..6cc5cab2 --- /dev/null +++ b/packages/shared-python/shared/services/retrieval/namespace_map_snapshot_cache.py @@ -0,0 +1,43 @@ +"""Process-local cache for decoded namespace MAP snapshot documents. + +Keyed by ``(user_id, namespace, generation)`` so a publish/archive that bumps +the namespace generation invalidates stale entries automatically -- no manual +invalidation call is needed. Bounded LRU keeps memory use predictable across +many namespaces sharing one worker process. +""" + +from __future__ import annotations + +import threading +from collections import OrderedDict +from typing import Any + +_MAX_ENTRIES = 64 +_lock = threading.Lock() +_cache: "OrderedDict[tuple[str, str, int], dict[str, dict[str, Any]]]" = OrderedDict() + + +def get_cached_namespace_documents( + *, user_id: str, namespace: str, generation: int +) -> dict[str, dict[str, Any]] | None: + key = (user_id, namespace, generation) + with _lock: + documents = _cache.get(key) + if documents is not None: + _cache.move_to_end(key) + return documents + + +def cache_namespace_documents( + *, + user_id: str, + namespace: str, + generation: int, + documents: dict[str, dict[str, Any]], +) -> None: + key = (user_id, namespace, generation) + with _lock: + _cache[key] = documents + _cache.move_to_end(key) + while len(_cache) > _MAX_ENTRIES: + _cache.popitem(last=False) diff --git a/packages/shared-python/shared/services/retrieval/nav/knowhere_hybrid.py b/packages/shared-python/shared/services/retrieval/nav/knowhere_hybrid.py index 369da962..3c237aef 100644 --- a/packages/shared-python/shared/services/retrieval/nav/knowhere_hybrid.py +++ b/packages/shared-python/shared/services/retrieval/nav/knowhere_hybrid.py @@ -664,7 +664,11 @@ def score_unit_stream_hybrid_many( unit_factory: Callable[[], Iterable[ScoreUnitRow]], queries: Sequence[str], ) -> Dict[str, Dict[str, float]]: - """Score several queries exactly while reading the corpus only once.""" + """Score several queries exactly while reading the corpus only once. + + Corpus-wide map-nav scoring is path+content only; the term channel was + retired here (see the unify-bm25-persistent-map plan). + """ unique_queries = list(dict.fromkeys(str(query) for query in queries)) if not unique_queries: return {} @@ -677,7 +681,6 @@ def score_unit_stream_hybrid_many( for query_tokens in query_tokens_by_query.values() for token in query_tokens } - query_lower_by_query = {query: query.lower().strip() for query in unique_queries} units: List[_StreamingManyUnit] = [] path_stats = _StreamingBm25Stats.empty() content_stats = _StreamingBm25Stats.empty() @@ -693,20 +696,6 @@ def score_unit_stream_hybrid_many( content_stats.observe(content_tokens) path_frequencies = Counter(path_tokens) content_frequencies = Counter(content_tokens) - haystack = str(row.get("term_search_text") or "").lower() - term_scores: List[float] = [] - for query in unique_queries: - query_lower = query_lower_by_query[query] - query_tokens = query_tokens_by_query[query] - term_score = 0.0 - if query_lower: - if query_lower in haystack: - term_score = 100.0 - else: - hit_count = sum(1 for token in query_tokens if token in haystack) - if hit_count > 0: - term_score = float(hit_count) - term_scores.append(term_score) units.append( _StreamingManyUnit( unit_id=unit_id, @@ -722,7 +711,6 @@ def score_unit_stream_hybrid_many( for token in query_token_set if content_frequencies[token] }, - term_scores=tuple(term_scores), ) ) path_stats.finalize() @@ -733,9 +721,8 @@ def score_unit_stream_hybrid_many( path_stats=path_stats, content_stats=content_stats, query_tokens=query_tokens_by_query[query], - query_index=index, ) - for index, query in enumerate(unique_queries) + for query in unique_queries } @@ -745,11 +732,9 @@ def _score_streaming_units( path_stats: "_StreamingBm25Stats", content_stats: "_StreamingBm25Stats", query_tokens: List[str], - query_index: int, ) -> Dict[str, float]: path_by_id: Dict[str, float] = {} content_by_id: Dict[str, float] = {} - term_by_id: Dict[str, float] = {} unit_ids = list(dict.fromkeys(unit.unit_id for unit in units)) for unit in units: path_score = path_stats.score( @@ -760,11 +745,6 @@ def _score_streaming_units( ) path_by_id[unit.unit_id] = path_score content_by_id[unit.unit_id] = content_score - term_by_id[unit.unit_id] = float( - unit.term_scores[query_index] - if query_index < len(unit.term_scores) - else 0.0 - ) path_rows = [ (score, unit_id) for unit_id, score in path_by_id.items() if score > 0.0 @@ -772,13 +752,9 @@ def _score_streaming_units( content_rows = [ (score, unit_id) for unit_id, score in content_by_id.items() if score > 0.0 ] - term_rows = [ - (score, unit_id) for unit_id, score in term_by_id.items() if score > 0.0 - ] path_rows.sort(key=lambda item: (-item[0], item[1])) content_rows.sort(key=lambda item: (-item[0], item[1])) - term_rows.sort(key=lambda item: (-item[0], item[1])) - path_weight, content_weight, term_weight = map_channel_weights() + path_weight, content_weight, _term_weight = map_channel_weights() rrf_k = int( os.environ.get( "NAV_MAP_RRF_K", @@ -791,8 +767,6 @@ def _score_streaming_units( fused[unit_id] = fused.get(unit_id, 0.0) + path_weight / (rrf_k + rank + 1) for rank, (_score, unit_id) in enumerate(content_rows): fused[unit_id] = fused.get(unit_id, 0.0) + content_weight / (rrf_k + rank + 1) - for rank, (_score, unit_id) in enumerate(term_rows): - fused[unit_id] = fused.get(unit_id, 0.0) + term_weight / (rrf_k + rank + 1) return {unit_id: round(score, 6) for unit_id, score in fused.items()} @@ -803,7 +777,6 @@ class _StreamingManyUnit: content_length: int path_frequencies: Mapping[str, int] content_frequencies: Mapping[str, int] - term_scores: Tuple[float, ...] @dataclass(frozen=True) @@ -825,7 +798,6 @@ class PersistedScoreUnit: content_length: int path_frequencies: Mapping[str, int] content_frequencies: Mapping[str, int] - term_scores: Tuple[float, ...] @dataclass(frozen=True) @@ -854,7 +826,6 @@ def score_persisted_corpus_many( content_length=unit.content_length, path_frequencies=unit.path_frequencies, content_frequencies=unit.content_frequencies, - term_scores=unit.term_scores, ) for unit in corpus.units ] @@ -864,9 +835,8 @@ def score_persisted_corpus_many( path_stats=path_stats, content_stats=content_stats, query_tokens=tokenize_query_for_ranker(query), - query_index=query_index, ) - for query_index, query in enumerate(unique_queries) + for query in unique_queries } diff --git a/packages/shared-python/shared/services/retrieval/nav/nav_knowhere.py b/packages/shared-python/shared/services/retrieval/nav/nav_knowhere.py index e25234ea..757f3212 100644 --- a/packages/shared-python/shared/services/retrieval/nav/nav_knowhere.py +++ b/packages/shared-python/shared/services/retrieval/nav/nav_knowhere.py @@ -55,11 +55,6 @@ _DEFAULT_DSN = "postgresql://root:root123@127.0.0.1:5433/Knowhere" _MAP_UNIT_INDEX_FORMAT_VERSION = 1 _MAP_SCORE_CHANNELS: Tuple[str, str] = ("path", "content") -# PostgreSQL's trigram planner evaluates one LIKE branch per pattern. Once a -# query has more than a handful of patterns, fetching the bounded pinned unit -# set once and evaluating literal substring hits in Python is faster and keeps -# SQL planning time from growing with the planner's subgoal length. -_TERM_SCORE_FULL_SCAN_PATTERN_THRESHOLD = 8 _logger = logging.getLogger(__name__) @@ -236,9 +231,6 @@ def __init__( tuple[tuple[str, str], ...], dict[tuple[str, str], dict[str, int]], ] = {} - self._score_term_cache: dict[ - tuple[tuple[str, str], ...], dict[str, Tuple[float, ...]] - ] = {} def _connection(self) -> "_SyncConnection": if self._conn is None: @@ -559,16 +551,6 @@ def load_persisted_score_corpus( len(map_unit_ids), ) - term_cache_key = (revision_key, tuple(unique_queries)) - term_scores = self._score_term_cache.get(term_cache_key) - if term_scores is None: - term_scores = self._load_term_scores( - cur, - map_unit_ids=map_unit_ids, - queries=unique_queries, - query_tokens_by_query=query_tokens_by_query, - ) - self._score_term_cache[term_cache_key] = term_scores _logger.info( "retrieval map-index load stage=complete units=%d queries=%d", len(unit_rows), @@ -602,9 +584,6 @@ def load_persisted_score_corpus( content_frequencies=frequencies.get( (str(row[0]), "content"), {} ), - term_scores=term_scores.get( - str(row[0]), tuple(0.0 for _query in unique_queries) - ), ) for row in unit_rows ], @@ -614,92 +593,6 @@ def load_persisted_score_corpus( finally: cur.close() - def _load_term_scores( - self, - cur: "_SyncCursor", - *, - map_unit_ids: Sequence[str], - queries: Sequence[str], - query_tokens_by_query: Mapping[str, Sequence[str]], - ) -> Dict[str, Tuple[float, ...]]: - if not map_unit_ids or not queries: - return {} - # Keep candidate selection in PostgreSQL so the trigram index can - # discard non-matching units, but compute the exact score once per - # returned row in Python. The previous query generated one POSITION - # expression and one OR predicate for every query token. Long planner - # subgoals therefore produced very large SQL statements and repeated - # substring evaluation for the same row. LIKE ANY keeps the SQL shape - # constant while preserving the existing literal substring semantics in - # the final Python scoring pass (LIKE may over-select wildcard matches, - # which are rejected by the literal checks below). - candidate_patterns: list[str] = [] - for query in queries: - query_lower = query.lower().strip() - if not query_lower: - continue - candidate_patterns.append(f"%{query_lower}%") - candidate_patterns.extend( - f"%{token}%" for token in query_tokens_by_query[query] if token - ) - candidate_patterns = list(dict.fromkeys(candidate_patterns)) - if not candidate_patterns: - return {} - stage_started = time.perf_counter() - if len(candidate_patterns) > _TERM_SCORE_FULL_SCAN_PATTERN_THRESHOLD: - # Long subgoals make LIKE ANY increasingly expensive even with the - # trigram index. The map-unit id list is already bounded by the - # pinned serving projection, so one text fetch plus literal Python - # checks avoids a query whose shape grows with token count. - cur.execute( - "SELECT id, term_search_text_lower " - "FROM document_map_units WHERE id = ANY(%s)", - (list(map_unit_ids),), - ) - candidate_mode = "full_scan" - else: - cur.execute( - "SELECT id, term_search_text_lower " - "FROM document_map_units " - "WHERE id = ANY(%s) AND term_search_text_lower LIKE ANY(%s)", - (list(map_unit_ids), candidate_patterns), - ) - candidate_mode = "trigram" - rows = cur.fetchall() - _logger.info( - "retrieval map-index load stage=term_scores units=%d queries=%d mode=%s seconds=%.3f", - len(map_unit_ids), - len(queries), - candidate_mode, - time.perf_counter() - stage_started, - ) - scores_by_unit: Dict[str, Tuple[float, ...]] = {} - for row in rows: - if len(row) < 2: - continue - unit_id = str(row[0]) - haystack = str(row[1] or "").lower() - scores: list[float] = [] - for query in queries: - query_lower = query.lower().strip() - if not query_lower: - scores.append(0.0) - elif query_lower in haystack: - scores.append(100.0) - else: - scores.append( - float( - sum( - 1 - for token in query_tokens_by_query[query] - if token in haystack - ) - ) - ) - if any(score > 0.0 for score in scores): - scores_by_unit[unit_id] = tuple(scores) - return scores_by_unit - def _load_persisted_bm25_stats( self, cur: "_SyncCursor", diff --git a/packages/shared-python/shared/services/retrieval/nav_snapshot.py b/packages/shared-python/shared/services/retrieval/nav_snapshot.py index 5d34e451..95d24078 100644 --- a/packages/shared-python/shared/services/retrieval/nav_snapshot.py +++ b/packages/shared-python/shared/services/retrieval/nav_snapshot.py @@ -31,6 +31,7 @@ Document, DocumentChunk, DocumentSection, + RetrievalNamespaceMapSnapshot, RetrievalServingRevisionManifest, ) from shared.models.database.job_result import JobResult @@ -46,6 +47,10 @@ from shared.services.retrieval.search.section_filters import is_excluded_section from shared.services.retrieval.serving_manifest import decode_serving_manifest from shared.services.retrieval.manifest_cache import get_cached_manifest_payloads +from shared.services.retrieval.namespace_map_snapshot_cache import ( + cache_namespace_documents, + get_cached_namespace_documents, +) # Keep each payload query bounded under the API's 30-second statement timeout. @@ -212,15 +217,40 @@ async def load_nav_snapshot( if job_result_id and job_id } - manifest_sections = None - if len(document_revisions) <= _MANIFEST_MAX_REVISION_COUNT: - manifest_sections = await _load_manifest_sections( - db, - document_revisions=document_revisions, + snapshot_entries = await _resolve_namespace_snapshot_entries( + db, + user_id=user_id, + namespace=namespace, + document_revisions=document_revisions, + ) + if snapshot_entries is not None: + manifest_sections = _parse_manifest_entries( + snapshot_entries, exclude_sections=excluded_secs, job_id_by_result_id=job_id_by_result_id, ) + else: + _logger.warning( + "retrieval snapshot fallback=manifest_merge user_id=%s namespace=%s documents=%d", + user_id, + namespace, + len(document_revisions), + ) + manifest_sections = None + if len(document_revisions) <= _MANIFEST_MAX_REVISION_COUNT: + manifest_sections = await _load_manifest_sections( + db, + document_revisions=document_revisions, + exclude_sections=excluded_secs, + job_id_by_result_id=job_id_by_result_id, + ) if manifest_sections is None: + _logger.warning( + "retrieval snapshot fallback=table_scan user_id=%s namespace=%s documents=%d", + user_id, + namespace, + len(document_revisions), + ) sections_by_doc, section_path_by_id = await _load_sections( db, document_revisions=document_revisions, @@ -340,18 +370,72 @@ async def load_nav_snapshot( ) -async def _load_manifest_sections( +async def _resolve_namespace_snapshot_entries( db: SnapshotSession, *, + user_id: str, + namespace: str, document_revisions: list[tuple[str, str]], - exclude_sections: list[dict[str, str]], - job_id_by_result_id: dict[str, str], -) -> tuple[ - dict[str, list[SectionRow]], - dict[str, str], - tuple[dict[str, list[str]], dict[str, dict[str, Any]], dict[str, dict[str, Any]]], -] | None: - """Load section metadata from complete serving manifests when available.""" +) -> list[tuple[object, ...]] | None: + """Return manifest-shaped entries from the persisted namespace snapshot. + + Returns ``None`` (triggering the exact per-revision fallback) when the + snapshot row is missing, corrupt, or stale for any requested revision. + """ + statement = select( + RetrievalNamespaceMapSnapshot.generation, + RetrievalNamespaceMapSnapshot.payload_zlib, + RetrievalNamespaceMapSnapshot.checksum, + RetrievalNamespaceMapSnapshot.format_version, + ).where( + RetrievalNamespaceMapSnapshot.user_id == user_id, + RetrievalNamespaceMapSnapshot.namespace == namespace, + ) + try: + row = (await db.execute(statement)).first() + except Exception: + await db.rollback() + return None + if row is None: + return None + generation, payload_zlib, checksum, format_version = row + documents = get_cached_namespace_documents( + user_id=user_id, namespace=namespace, generation=int(generation) + ) + if documents is None: + try: + payload = decode_serving_manifest( + bytes(payload_zlib), + checksum=str(checksum), + format_version=int(format_version), + ) + except (ValueError, TypeError): + return None + decoded_documents = payload.get("documents") + if not isinstance(decoded_documents, dict): + return None + documents = decoded_documents + cache_namespace_documents( + user_id=user_id, + namespace=namespace, + generation=int(generation), + documents=documents, + ) + entries: list[tuple[object, ...]] = [] + for document_id, job_result_id in document_revisions: + entry = documents.get(document_id) + if not isinstance(entry, dict) or str(entry.get("job_result_id") or "") != job_result_id: + return None + entries.append((document_id, job_result_id, entry, None, None)) + return entries + + +async def _resolve_manifest_entries( + db: SnapshotSession, + *, + document_revisions: list[tuple[str, str]], +) -> list[tuple[object, ...]] | None: + """Resolve per-revision manifest rows from the request cache or table.""" cached_payloads = get_cached_manifest_payloads( db, revisions=dict(document_revisions), @@ -365,32 +449,71 @@ async def _load_manifest_sections( ] if len(manifest_entries) != len(document_revisions): return None - else: - manifest_entries = [] - for group_start in range(0, len(document_revisions), _REVISION_GROUP_SIZE): - revision_group = document_revisions[ - group_start : group_start + _REVISION_GROUP_SIZE - ] - statement = select( + return manifest_entries + + manifest_entries = [] + for group_start in range(0, len(document_revisions), _REVISION_GROUP_SIZE): + revision_group = document_revisions[ + group_start : group_start + _REVISION_GROUP_SIZE + ] + statement = select( + RetrievalServingRevisionManifest.document_id, + RetrievalServingRevisionManifest.job_result_id, + RetrievalServingRevisionManifest.payload_zlib, + RetrievalServingRevisionManifest.checksum, + RetrievalServingRevisionManifest.format_version, + ).where( + tuple_( RetrievalServingRevisionManifest.document_id, RetrievalServingRevisionManifest.job_result_id, - RetrievalServingRevisionManifest.payload_zlib, - RetrievalServingRevisionManifest.checksum, - RetrievalServingRevisionManifest.format_version, - ).where( - tuple_( - RetrievalServingRevisionManifest.document_id, - RetrievalServingRevisionManifest.job_result_id, - ).in_(revision_group) - ) - try: - rows = (await db.execute(statement)).all() - except Exception: - await db.rollback() - return None - if len(rows) != len(revision_group): - return None - manifest_entries.extend(tuple(row) for row in rows) + ).in_(revision_group) + ) + try: + rows = (await db.execute(statement)).all() + except Exception: + await db.rollback() + return None + if len(rows) != len(revision_group): + return None + manifest_entries.extend(tuple(row) for row in rows) + return manifest_entries + + +async def _load_manifest_sections( + db: SnapshotSession, + *, + document_revisions: list[tuple[str, str]], + exclude_sections: list[dict[str, str]], + job_id_by_result_id: dict[str, str], +) -> tuple[ + dict[str, list[SectionRow]], + dict[str, str], + tuple[dict[str, list[str]], dict[str, dict[str, Any]], dict[str, dict[str, Any]]], +] | None: + """Load section metadata from complete serving manifests when available.""" + manifest_entries = await _resolve_manifest_entries( + db, document_revisions=document_revisions + ) + if manifest_entries is None: + return None + return _parse_manifest_entries( + manifest_entries, + exclude_sections=exclude_sections, + job_id_by_result_id=job_id_by_result_id, + ) + + +def _parse_manifest_entries( + manifest_entries: list[tuple[object, ...]], + *, + exclude_sections: list[dict[str, str]], + job_id_by_result_id: dict[str, str], +) -> tuple[ + dict[str, list[SectionRow]], + dict[str, str], + tuple[dict[str, list[str]], dict[str, dict[str, Any]], dict[str, dict[str, Any]]], +] | None: + """Parse manifest-shaped entries (decoded dict or raw compressed row) into sections.""" by_doc: dict[str, list[SectionRow]] = {} path_by_id: dict[str, str] = {} ids_by_doc: dict[str, list[str]] = {} diff --git a/packages/shared-python/shared/services/retrieval/publication_content.py b/packages/shared-python/shared/services/retrieval/publication_content.py index 34092436..7ce75b80 100644 --- a/packages/shared-python/shared/services/retrieval/publication_content.py +++ b/packages/shared-python/shared/services/retrieval/publication_content.py @@ -13,6 +13,9 @@ DocumentSection, ) from shared.services.retrieval.map_unit_index import replace_document_map_units +from shared.services.retrieval.namespace_map_snapshot import ( + patch_namespace_map_snapshot, +) from shared.services.retrieval.publication_models import DocumentPublicationScope from shared.services.retrieval.serving_manifest import persist_revision_serving_state from shared.services.retrieval.search.lexical_text import ( @@ -93,7 +96,8 @@ def replace_document_revision_content( db.flush() replace_document_map_units(db, scope=scope) db.flush() - persist_revision_serving_state(db, scope=scope) + manifest_payload = persist_revision_serving_state(db, scope=scope) + patch_namespace_map_snapshot(db, scope=scope, manifest_payload=manifest_payload) class DocumentSectionPublisher: diff --git a/packages/shared-python/shared/services/retrieval/publication_service.py b/packages/shared-python/shared/services/retrieval/publication_service.py index d7c879de..58fa98dd 100644 --- a/packages/shared-python/shared/services/retrieval/publication_service.py +++ b/packages/shared-python/shared/services/retrieval/publication_service.py @@ -22,6 +22,9 @@ from shared.models.schemas.job_metadata import JobMetadataHelper from shared.models.schemas.retrieval_namespace import normalize_retrieval_namespace from shared.services.retrieval.graph.service import DocumentGraphService, GraphScope +from shared.services.retrieval.namespace_map_snapshot import ( + remove_document_from_namespace_map_snapshot, +) from shared.services.retrieval.publication_content import ( deduplicate_chunks_by_source_path, replace_document_revision_content, @@ -191,6 +194,12 @@ def _publish_document_state_for_job( user_id=scope.user_id, namespace=str(existing_namespace), ) + remove_document_from_namespace_map_snapshot( + db, + user_id=scope.user_id, + namespace=str(existing_namespace), + document_id=document.document_id, + ) advance_namespace_generation( db, user_id=scope.user_id, diff --git a/packages/shared-python/shared/services/retrieval/serving_manifest.py b/packages/shared-python/shared/services/retrieval/serving_manifest.py index b2dd40c5..5fe979e0 100644 --- a/packages/shared-python/shared/services/retrieval/serving_manifest.py +++ b/packages/shared-python/shared/services/retrieval/serving_manifest.py @@ -190,8 +190,12 @@ def persist_revision_serving_state( db: Session, *, scope: DocumentPublicationScope, -) -> None: - """Replace manifest and statistics rows for one revision atomically.""" +) -> dict[str, Any]: + """Replace manifest and statistics rows for one revision atomically. + + Returns the manifest payload so callers can patch the namespace-level MAP + snapshot without rebuilding it. + """ manifest_payload = build_revision_serving_payload(db, scope=scope) statistics_payload = build_revision_statistics_payload(db, scope=scope) manifest_bytes, manifest_checksum, manifest_version = encode_serving_manifest( @@ -232,6 +236,7 @@ def persist_revision_serving_state( checksum=statistics_checksum, ) ) + return manifest_payload def rebuild_namespace_serving_statistics( From 6ab4f8e356790d5d63faffb1751f475b005290a8 Mon Sep 17 00:00:00 2001 From: chengke <404835780@qq.com> Date: Mon, 31 Aug 2026 15:18:10 +0800 Subject: [PATCH 12/19] feat: enhance map unit index backfill process and remove obsolete tests - Updated the backfill_map_unit_indexes script to include patching of the namespace map snapshot after persisting the revision serving state. - Removed outdated contract tests related to BM25 FTS prefilter and unit tests for BM25 channel token preparation. - Refactored retrieval logic to improve handling of document map units, including new fields for asset presence. --- ...d8e9f0a1b2c_add_map_unit_asset_presence.py | 63 +++ apps/api/scripts/backfill_map_unit_indexes.py | 15 +- .../test_bm25_fts_prefilter_contract.py | 231 -------- ...est_retrieval_classic_map_unit_contract.py | 282 ++++++++++ .../test_retrieval_map_unit_index_contract.py | 13 + .../test_retrieval_revision_races_contract.py | 63 +-- .../tests/unit/test_bm25_channel_tsquery.py | 45 -- .../shared/models/database/document.py | 18 + .../services/retrieval/execution/routes.py | 7 +- .../services/retrieval/map_unit_index.py | 7 + .../services/retrieval/nav/knowhere_hybrid.py | 3 +- .../services/retrieval/search/channels.py | 529 ------------------ .../services/retrieval/search/discovery.py | 216 ------- .../retrieval/search/lexical_ranker.py | 74 --- .../retrieval/search/map_unit_discovery.py | 500 +++++++++++++++++ .../services/retrieval/search/scoring.py | 29 - .../tests/test_retrieval_search_channels.py | 216 ------- 17 files changed, 900 insertions(+), 1411 deletions(-) create mode 100644 apps/api/alembic/versions/7d8e9f0a1b2c_add_map_unit_asset_presence.py delete mode 100644 apps/api/tests/contract/test_bm25_fts_prefilter_contract.py create mode 100644 apps/api/tests/contract/test_retrieval_classic_map_unit_contract.py delete mode 100644 apps/api/tests/unit/test_bm25_channel_tsquery.py delete mode 100644 packages/shared-python/shared/services/retrieval/search/channels.py delete mode 100644 packages/shared-python/shared/services/retrieval/search/discovery.py delete mode 100644 packages/shared-python/shared/services/retrieval/search/lexical_ranker.py create mode 100644 packages/shared-python/shared/services/retrieval/search/map_unit_discovery.py delete mode 100644 packages/shared-python/shared/tests/test_retrieval_search_channels.py diff --git a/apps/api/alembic/versions/7d8e9f0a1b2c_add_map_unit_asset_presence.py b/apps/api/alembic/versions/7d8e9f0a1b2c_add_map_unit_asset_presence.py new file mode 100644 index 00000000..b50e947a --- /dev/null +++ b/apps/api/alembic/versions/7d8e9f0a1b2c_add_map_unit_asset_presence.py @@ -0,0 +1,63 @@ +"""Add has_image/has_table presence flags to document_map_units.""" + +from __future__ import annotations + +import sqlalchemy as sa +from alembic import op + +revision = "7d8e9f0a1b2c" +down_revision = "6c7d8e9f0a1b" +branch_labels = None +depends_on = None + + +def upgrade() -> None: + bind = op.get_bind() + inspector = sa.inspect(bind) + columns = {col["name"] for col in inspector.get_columns("document_map_units")} + if "has_image" not in columns: + op.add_column( + "document_map_units", + sa.Column( + "has_image", sa.Boolean(), nullable=False, server_default=sa.false() + ), + ) + if "has_table" not in columns: + op.add_column( + "document_map_units", + sa.Column( + "has_table", sa.Boolean(), nullable=False, server_default=sa.false() + ), + ) + + inspector = sa.inspect(bind) + indexes = {item["name"] for item in inspector.get_indexes("document_map_units")} + if "idx_document_map_units_has_image" not in indexes: + op.create_index( + "idx_document_map_units_has_image", + "document_map_units", + ["document_id", "job_result_id"], + postgresql_where=sa.text("has_image = true"), + ) + if "idx_document_map_units_has_table" not in indexes: + op.create_index( + "idx_document_map_units_has_table", + "document_map_units", + ["document_id", "job_result_id"], + postgresql_where=sa.text("has_table = true"), + ) + + +def downgrade() -> None: + inspector = sa.inspect(op.get_bind()) + indexes = {item["name"] for item in inspector.get_indexes("document_map_units")} + if "idx_document_map_units_has_table" in indexes: + op.drop_index("idx_document_map_units_has_table", table_name="document_map_units") + if "idx_document_map_units_has_image" in indexes: + op.drop_index("idx_document_map_units_has_image", table_name="document_map_units") + + columns = {col["name"] for col in inspector.get_columns("document_map_units")} + if "has_table" in columns: + op.drop_column("document_map_units", "has_table") + if "has_image" in columns: + op.drop_column("document_map_units", "has_image") diff --git a/apps/api/scripts/backfill_map_unit_indexes.py b/apps/api/scripts/backfill_map_unit_indexes.py index 2bfd4177..be2b4dbb 100644 --- a/apps/api/scripts/backfill_map_unit_indexes.py +++ b/apps/api/scripts/backfill_map_unit_indexes.py @@ -1,7 +1,10 @@ """Backfill persisted MAP-NAV lexical indexes for existing revisions. -The migration creates empty derived tables intentionally. Run this command -after deployment with ``--apply`` so each revision is rebuilt and committed +Rebuilds, per active revision: the map-unit index, the revision serving +manifest, that document's subtree in the namespace MAP snapshot, namespace +statistics, and the namespace generation. The migrations that create these +derived tables leave them empty intentionally. Run this command after +deployment with ``--apply`` so each revision is rebuilt and committed independently; without ``--apply`` it is a read-only inventory. """ @@ -44,6 +47,9 @@ def _bootstrap_python_path() -> None: from shared.core.database_sync import get_sync_session_factory from shared.models.database.document import Document from shared.services.retrieval.map_unit_index import replace_document_map_units +from shared.services.retrieval.namespace_map_snapshot import ( + patch_namespace_map_snapshot, +) from shared.services.retrieval.publication_models import DocumentPublicationScope from shared.services.retrieval.serving_generation import ( advance_namespace_generation, @@ -130,7 +136,10 @@ def backfill_map_unit_indexes(*, apply: bool, document_id: str = "") -> int: source_file_name=str(locked_document.source_file_name or ""), ) replace_document_map_units(db, scope=scope) - persist_revision_serving_state(db, scope=scope) + manifest_payload = persist_revision_serving_state(db, scope=scope) + patch_namespace_map_snapshot( + db, scope=scope, manifest_payload=manifest_payload + ) rebuild_namespace_serving_statistics( db, user_id=scope.user_id, diff --git a/apps/api/tests/contract/test_bm25_fts_prefilter_contract.py b/apps/api/tests/contract/test_bm25_fts_prefilter_contract.py deleted file mode 100644 index 56add8e5..00000000 --- a/apps/api/tests/contract/test_bm25_fts_prefilter_contract.py +++ /dev/null @@ -1,231 +0,0 @@ -"""Contract tests for the BM25 channel Postgres FTS prefilter. - -These run against a real Postgres so the prefilter is validated against the -same generated tsvector columns and GIN indexes production uses. A pure-Python -fake would not catch a mismatch between the query configuration and the one -the columns were generated with. -""" - -from __future__ import annotations - -from collections.abc import AsyncGenerator - -import pytest -import pytest_asyncio -from shared.services.retrieval.search.channels import content_channel, path_channel -from shared.testing.contract_runtime import PostgreSQLProcess -from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker, create_async_engine -from sqlalchemy import text - -_SCHEMA = """ -CREATE TABLE documents ( - document_id TEXT PRIMARY KEY, - user_id TEXT, - namespace TEXT, - status TEXT, - current_job_result_id TEXT, - source_file_name TEXT -); -CREATE TABLE job_results (id TEXT PRIMARY KEY, job_id TEXT); -CREATE TABLE document_sections (section_id TEXT PRIMARY KEY, section_path TEXT); -CREATE TABLE document_chunks ( - id SERIAL PRIMARY KEY, - chunk_id TEXT, - document_id TEXT, - section_id TEXT, - chunk_type TEXT, - content TEXT, - source_chunk_path TEXT, - file_path TEXT, - chunk_metadata JSONB, - job_result_id TEXT, - sort_order INTEGER, - content_search_text TEXT, - content_search_tsv TSVECTOR GENERATED ALWAYS AS - (to_tsvector('simple', COALESCE(content_search_text, ''))) STORED, - path_search_text TEXT, - path_search_tsv TSVECTOR GENERATED ALWAYS AS - (to_tsvector('simple', COALESCE(path_search_text, ''))) STORED, - term_search_text TEXT -); -CREATE INDEX idx_chunk_content_search_tsv ON document_chunks USING GIN (content_search_tsv); -CREATE INDEX idx_chunk_path_search_tsv ON document_chunks USING GIN (path_search_tsv); -""" - -_NOISE_ROWS = 300 - - -@pytest_asyncio.fixture -async def seeded_session( - postgresql_proc: PostgreSQLProcess, -) -> AsyncGenerator[AsyncSession, None]: - dsn = ( - f"postgresql+asyncpg://{postgresql_proc.user}@" - f"{postgresql_proc.host}:{postgresql_proc.port}/postgres" - ) - engine = create_async_engine(dsn, isolation_level="AUTOCOMMIT") - async with engine.begin() as conn: - await conn.execute(text("DROP SCHEMA IF EXISTS bm25_fts CASCADE")) - await conn.execute(text("CREATE SCHEMA bm25_fts")) - await conn.execute(text("SET search_path TO bm25_fts")) - for statement in filter(None, (s.strip() for s in _SCHEMA.split(";"))): - await conn.execute(text(statement)) - await conn.execute(text("INSERT INTO job_results VALUES (1, 'job1')")) - await conn.execute( - text( - "INSERT INTO documents VALUES " - "('d1', 'u1', 'ns1', 'active', 1, 'sample.pdf')" - ) - ) - await conn.execute(text("INSERT INTO document_sections VALUES ('s1', '/root')")) - await conn.execute( - text( - "INSERT INTO document_chunks " - "(chunk_id, document_id, section_id, chunk_type, content, " - " job_result_id, sort_order, content_search_text, path_search_text) " - "VALUES " - "('hit-en', 'd1', 's1', 'text', 'body', 1, 1, " - " 'alpha beta gamma', 'invoices alpha'), " - "('hit-cjk', 'd1', 's1', 'text', 'body', 1, 2, " - " '合同 条款 甲方', '合同 目录')" - ) - ) - await conn.execute( - text( - "INSERT INTO document_chunks " - "(chunk_id, document_id, section_id, chunk_type, content, " - " job_result_id, sort_order, content_search_text, path_search_text) " - "SELECT 'noise-' || i, 'd1', 's1', 'text', 'body', 1, i + 10, " - " 'filler unrelated wording ' || i, 'misc path ' || i " - "FROM generate_series(1, :noise) AS i" - ), - {"noise": _NOISE_ROWS}, - ) - - session_factory = async_sessionmaker(engine, expire_on_commit=False) - async with session_factory() as session: - await session.execute(text("SET search_path TO bm25_fts")) - yield session - await engine.dispose() - - -async def _content_hits(session: AsyncSession, query: str) -> list[str]: - rows = await content_channel( - session, - user_id="u1", - namespace="ns1", - query=query, - top_k=50, - exclude_document_ids=[], - exclude_sections=[], - ) - return [str(row["chunk_id"]) for row in rows] - - -@pytest.mark.asyncio -async def test_content_channel_returns_only_query_matching_chunks( - seeded_session: AsyncSession, -) -> None: - # The corpus holds hundreds of unrelated chunks. Before the prefilter every - # one of them was loaded into Python for BM25 scoring. - assert await _content_hits(seeded_session, "alpha") == ["hit-en"] - - -@pytest.mark.asyncio -async def test_content_channel_matches_cjk_tokens( - seeded_session: AsyncSession, -) -> None: - assert await _content_hits(seeded_session, "合同") == ["hit-cjk"] - - -@pytest.mark.asyncio -async def test_content_channel_uses_or_semantics_across_tokens( - seeded_session: AsyncSession, -) -> None: - # A row matching any single query token must survive, matching how the - # Python BM25 ranker admits rows. - hits = await _content_hits(seeded_session, "alpha 合同") - assert sorted(hits) == ["hit-cjk", "hit-en"] - - -@pytest.mark.asyncio -async def test_tsquery_operators_in_query_do_not_change_filter_shape( - seeded_session: AsyncSession, -) -> None: - # Tokens are lexed by Postgres as data. If operators leaked into tsquery - # syntax, "alpha & zzzz" would AND and drop the row. - assert await _content_hits(seeded_session, "alpha & zzzz") == ["hit-en"] - assert await _content_hits(seeded_session, "!alpha") == ["hit-en"] - - -@pytest.mark.asyncio -async def test_query_matching_nothing_returns_no_rows( - seeded_session: AsyncSession, -) -> None: - # The fallback re-runs the unfiltered scan, and BM25 then scores no row - # above zero, so the channel still yields nothing. - assert await _content_hits(seeded_session, "zzzznomatch") == [] - - -@pytest.mark.asyncio -async def test_path_channel_prefilters_on_path_search_tsv( - seeded_session: AsyncSession, -) -> None: - rows = await path_channel( - seeded_session, - user_id="u1", - namespace="ns1", - query="invoices", - top_k=50, - exclude_document_ids=[], - exclude_sections=[], - ) - assert [str(row["chunk_id"]) for row in rows] == ["hit-en"] - - -@pytest.mark.asyncio -async def test_exclusions_still_apply_under_the_prefilter( - seeded_session: AsyncSession, -) -> None: - rows = await content_channel( - seeded_session, - user_id="u1", - namespace="ns1", - query="alpha", - top_k=50, - exclude_document_ids=["d1"], - exclude_sections=[], - ) - assert rows == [] - - -@pytest.mark.asyncio -async def test_content_channel_uses_the_requested_revision_pin( - seeded_session: AsyncSession, -) -> None: - await seeded_session.execute(text("INSERT INTO job_results VALUES (2, 'job2')")) - await seeded_session.execute( - text( - "INSERT INTO document_chunks " - "(chunk_id, document_id, section_id, chunk_type, content, " - " job_result_id, sort_order, content_search_text, path_search_text) " - "VALUES ('new-hit', 'd1', 's1', 'text', 'new body', 2, 1, " - " 'alpha replacement', 'new path')" - ) - ) - await seeded_session.execute( - text("UPDATE documents SET current_job_result_id = 2 WHERE document_id = 'd1'") - ) - - rows = await content_channel( - seeded_session, - user_id="u1", - namespace="ns1", - query="alpha", - top_k=50, - exclude_document_ids=[], - exclude_sections=[], - revision_pins={"d1": "1"}, - ) - - assert [str(row["chunk_id"]) for row in rows] == ["hit-en"] diff --git a/apps/api/tests/contract/test_retrieval_classic_map_unit_contract.py b/apps/api/tests/contract/test_retrieval_classic_map_unit_contract.py new file mode 100644 index 00000000..e20e647f --- /dev/null +++ b/apps/api/tests/contract/test_retrieval_classic_map_unit_contract.py @@ -0,0 +1,282 @@ +from __future__ import annotations + +from collections.abc import Callable +from contextlib import AbstractAsyncContextManager +from typing import Any, cast +from uuid import uuid4 + +from httpx import AsyncClient +from sqlalchemy import select + +from shared.models.database.document import DocumentMapUnit +from shared.services.retrieval.publication_content import ( + replace_document_revision_content, +) +from shared.services.retrieval.publication_models import DocumentPublicationScope +from shared.services.retrieval.serving_generation import lock_namespace_generation +from tests.support.contract_database import ContractDatabase +from tests.support.retrieval_snapshot_support import contract_db_session + +_USER_ID = "local-dev-user" + + +async def test_classic_route_maps_winning_unit_to_one_chunk( + developer_api_client_factory: Callable[ + [], AbstractAsyncContextManager[AsyncClient] + ], +) -> None: + identifier = uuid4().hex[:8] + namespace = f"classic-map-{identifier}" + async with developer_api_client_factory() as api_client: + first = await _publish_document( + namespace=namespace, + source_file_name="first.pdf", + chunks=[ + { + "chunk_id": f"hit-{identifier}", + "type": "text", + "content": "unique alpha ranking marker", + "path": "first.pdf/Root/Hit/body", + "order": 1, + "metadata": {}, + }, + { + "chunk_id": f"other-{identifier}", + "type": "text", + "content": "unrelated filler paragraph", + "path": "first.pdf/Root/Other/body", + "order": 2, + "metadata": {}, + }, + ], + ) + await _publish_document( + namespace=namespace, + source_file_name="second.pdf", + chunks=[ + { + "chunk_id": f"noise-{identifier}", + "type": "text", + "content": "more unrelated filler", + "path": "second.pdf/Root/Noise/body", + "order": 1, + "metadata": {}, + }, + ], + ) + response = await api_client.post( + "/api/v1/retrieval/query", + json={ + "namespace": namespace, + "query": "unique alpha ranking", + "top_k": 1, + "use_agentic": False, + }, + ) + + assert response.status_code == 200 + body = cast(dict[str, object], response.json()) + results = cast(list[dict[str, object]], body["results"]) + assert body["router_used"] == "classic_topk" + assert len(results) == 1 + assert results[0]["chunk_id"] == f"hit-{identifier}" + assert results[0]["chunk_type"] == "text" + assert results[0]["source"] == { + "document_id": first["document_id"], + "source_file_name": "first.pdf", + "section_path": "Root / Hit / body", + } + + +async def test_classic_route_image_filter_scores_only_units_with_images( + developer_api_client_factory: Callable[ + [], AbstractAsyncContextManager[AsyncClient] + ], +) -> None: + identifier = uuid4().hex[:8] + namespace = f"classic-image-{identifier}" + async with developer_api_client_factory() as api_client: + with_image = await _publish_document( + namespace=namespace, + source_file_name="with-image.pdf", + chunks=[ + { + "chunk_id": f"body-{identifier}", + "type": "text", + "content": "shared alpha marker next to a chart", + "path": "with-image.pdf/Root/Section/body", + "order": 1, + "metadata": {"connect_to": [{"target": f"chart-{identifier}"}]}, + }, + { + "chunk_id": f"chart-{identifier}", + "type": "image", + "content": "chart of shared alpha marker", + "path": "images/chart.png", + "order": 2, + "file_path": "images/chart.png", + "metadata": {}, + }, + ], + ) + await _publish_document( + namespace=namespace, + source_file_name="text-only.pdf", + chunks=[ + { + "chunk_id": f"text-only-{identifier}", + "type": "text", + "content": "shared alpha marker with no chart", + "path": "text-only.pdf/Root/Section/body", + "order": 1, + "metadata": {}, + }, + ], + ) + await _publish_document( + namespace=namespace, + source_file_name="other-image.pdf", + chunks=[ + { + "chunk_id": f"other-body-{identifier}", + "type": "text", + "content": "unrelated filler next to a photo", + "path": "other-image.pdf/Root/Section/body", + "order": 1, + "metadata": {"connect_to": [{"target": f"photo-{identifier}"}]}, + }, + { + "chunk_id": f"photo-{identifier}", + "type": "image", + "content": "unrelated landscape photo", + "path": "images/photo.png", + "order": 2, + "file_path": "images/photo.png", + "metadata": {}, + }, + ], + ) + async with contract_db_session() as db: + image_units = list( + ( + await db.execute( + select(DocumentMapUnit).where( + DocumentMapUnit.document_id == with_image["document_id"] + ) + ) + ).scalars() + ) + assert any(unit.has_image for unit in image_units) + + response = await api_client.post( + "/api/v1/retrieval/query", + json={ + "namespace": namespace, + "query": "shared alpha marker", + "top_k": 1, + "use_agentic": False, + "chunk_types": ["image"], + }, + ) + + assert response.status_code == 200 + body = cast(dict[str, object], response.json()) + results = cast(list[dict[str, object]], body["results"]) + assert body["router_used"] == "classic_topk" + assert len(results) == 1 + assert results[0]["chunk_id"] == f"chart-{identifier}" + assert results[0]["chunk_type"] == "image" + + +def _publish_revision_with_generation_lock( + sync_db: Any, + *, + scope: DocumentPublicationScope, + chunks: list[dict[str, Any]], +) -> None: + lock_namespace_generation( + sync_db, user_id=scope.user_id, namespace=scope.namespace + ) + replace_document_revision_content(sync_db, scope=scope, chunks=chunks) + + +async def _publish_document( + *, + namespace: str, + source_file_name: str, + chunks: list[dict[str, Any]], +) -> dict[str, str]: + document_id = f"doc_{uuid4().hex[:12]}" + job_id = f"job_{uuid4().hex[:12]}" + job_result_id = f"result_{uuid4().hex[:12]}" + await ContractDatabase.execute( + """ + INSERT INTO jobs ( + job_id, user_id, job_type, status, source_type, version, + webhook_enabled, created_at, updated_at, credits_charged, billing_status + ) VALUES ( + :job_id, :user_id, 'document_ingestion', 'done', 'file', 0, + false, NOW(), NOW(), 0, 'skipped' + ) + """, + {"job_id": job_id, "user_id": _USER_ID}, + ) + await ContractDatabase.execute( + """ + INSERT INTO documents ( + document_id, user_id, namespace, status, source_file_name, + parse_track, created_at, updated_at + ) VALUES ( + :document_id, :user_id, :namespace, 'active', + :source_file_name, 'chunk', NOW(), NOW() + ) + """, + { + "document_id": document_id, + "user_id": _USER_ID, + "namespace": namespace, + "source_file_name": source_file_name, + }, + ) + await ContractDatabase.execute( + """ + INSERT INTO job_results ( + id, job_id, document_id, delivery_mode, created_at, updated_at + ) VALUES ( + :job_result_id, :job_id, :document_id, 'inline', NOW(), NOW() + ) + """, + { + "job_result_id": job_result_id, + "job_id": job_id, + "document_id": document_id, + }, + ) + await ContractDatabase.execute( + """ + UPDATE documents SET current_job_result_id = :job_result_id + WHERE document_id = :document_id + """, + {"job_result_id": job_result_id, "document_id": document_id}, + ) + scope = DocumentPublicationScope( + user_id=_USER_ID, + namespace=namespace, + document_id=document_id, + job_result_id=job_result_id, + source_file_name=source_file_name, + ) + async with contract_db_session() as db: + await db.run_sync( + lambda sync_db: _publish_revision_with_generation_lock( + sync_db, + scope=scope, + chunks=chunks, + ) + ) + await db.commit() + return { + "document_id": document_id, + "job_id": job_id, + "job_result_id": job_result_id, + } diff --git a/apps/api/tests/contract/test_retrieval_map_unit_index_contract.py b/apps/api/tests/contract/test_retrieval_map_unit_index_contract.py index 6d30731c..6665eb3e 100644 --- a/apps/api/tests/contract/test_retrieval_map_unit_index_contract.py +++ b/apps/api/tests/contract/test_retrieval_map_unit_index_contract.py @@ -356,6 +356,19 @@ async def test_lazy_snapshot_defers_selected_asset_reference_metadata( ) await db.commit() + async with contract_db_session() as db: + persisted_units = list( + ( + await db.execute( + select(DocumentMapUnit).where( + DocumentMapUnit.document_id == document_id + ) + ) + ).scalars() + ) + assert any(unit.has_image for unit in persisted_units) + assert all(not unit.has_table for unit in persisted_units) + calls: list[tuple[str, str]] = [] original = ReadOnlyChunkStore.load_chunk_reference_metadata diff --git a/apps/api/tests/contract/test_retrieval_revision_races_contract.py b/apps/api/tests/contract/test_retrieval_revision_races_contract.py index 3ebb0232..43c47f27 100644 --- a/apps/api/tests/contract/test_retrieval_revision_races_contract.py +++ b/apps/api/tests/contract/test_retrieval_revision_races_contract.py @@ -1,11 +1,8 @@ -"""Deterministic contracts for revision and channel-session coherence.""" +"""Deterministic contracts for revision-generation coherence.""" from __future__ import annotations -import importlib -from collections.abc import AsyncGenerator -from contextlib import asynccontextmanager -from typing import Any, cast +from typing import cast import pytest from sqlalchemy.ext.asyncio import AsyncSession @@ -14,62 +11,6 @@ RetrievalRevisionPins, is_revision_generation_stable, ) -from shared.services.retrieval.search import discovery - - -@pytest.mark.asyncio -async def test_classic_channels_share_pins_but_use_distinct_sessions( - monkeypatch: pytest.MonkeyPatch, -) -> None: - pins = RetrievalRevisionPins( - revisions={"doc-1": "revision-1"}, - generation=7, - ) - sessions: list[object] = [] - observed: list[tuple[object, object]] = [] - - @asynccontextmanager - async def fake_context() -> AsyncGenerator[object, None]: - session = object() - sessions.append(session) - yield session - - async def fake_channel( - db: AsyncSession, - **kwargs: Any, - ) -> list[dict[str, Any]]: - observed.append((db, kwargs["revision_pins"])) - return [] - - # Import the active module explicitly. Some API contract fixtures reload - # shared.core.database between tests, leaving the package attribute pointed - # at an old module object; dotted-string patching can then miss the module - # used by discovery's lazy import. - database_module = importlib.import_module("shared.core.database") - monkeypatch.setattr(database_module, "get_db_context", fake_context) - monkeypatch.setattr(discovery, "path_channel", fake_channel) - monkeypatch.setattr(discovery, "content_channel", fake_channel) - monkeypatch.setattr(discovery, "term_channel", fake_channel) - - result = await discovery.bottom_discovery( - cast(AsyncSession, object()), - user_id="user-1", - namespace="namespace-1", - query="coherent query", - top_k=3, - exclude_document_ids=[], - exclude_sections=[], - revision_pins=pins, - ) - - assert result.status == "discovery_done" - assert len(sessions) == 3 - assert len({id(session) for session in sessions}) == 3 - assert len(observed) == 3 - assert {id(session) for session, _pins in observed} == { - id(session) for session in sessions - } - assert all(observed_pins is pins for _session, observed_pins in observed) class _GenerationResult: diff --git a/apps/api/tests/unit/test_bm25_channel_tsquery.py b/apps/api/tests/unit/test_bm25_channel_tsquery.py deleted file mode 100644 index aea2d9df..00000000 --- a/apps/api/tests/unit/test_bm25_channel_tsquery.py +++ /dev/null @@ -1,45 +0,0 @@ -"""Unit tests for BM25 channel Postgres FTS prefilter token preparation.""" - -from __future__ import annotations - -from shared.services.retrieval.search.channels import ( - _MAX_FTS_QUERY_TOKENS, - _prepare_fts_tokens, -) - - -def test_returns_empty_for_no_tokens() -> None: - assert _prepare_fts_tokens([]) == [] - - -def test_keeps_token_order() -> None: - assert _prepare_fts_tokens(["alpha", "beta"]) == ["alpha", "beta"] - - -def test_strips_surrounding_whitespace() -> None: - assert _prepare_fts_tokens([" alpha ", "beta"]) == ["alpha", "beta"] - - -def test_drops_blank_tokens() -> None: - assert _prepare_fts_tokens(["", " ", "alpha"]) == ["alpha"] - - -def test_returns_empty_when_every_token_is_blank() -> None: - assert _prepare_fts_tokens(["", " "]) == [] - - -def test_caps_token_count() -> None: - tokens = [f"tok{index}" for index in range(_MAX_FTS_QUERY_TOKENS + 25)] - assert len(_prepare_fts_tokens(tokens)) == _MAX_FTS_QUERY_TOKENS - - -def test_preserves_cjk_tokens() -> None: - assert _prepare_fts_tokens(["合同", "条款"]) == ["合同", "条款"] - - -def test_passes_tsquery_operators_through_untouched() -> None: - # Tokens travel to Postgres as a text[] parameter and are lexed there, so - # operator characters are data rather than syntax. Nothing is escaped or - # dropped here. - raw = ["alpha' & 'zzzz", "!beta", "a|b"] - assert _prepare_fts_tokens(raw) == raw diff --git a/packages/shared-python/shared/models/database/document.py b/packages/shared-python/shared/models/database/document.py index b85dd84f..2ebce6cd 100644 --- a/packages/shared-python/shared/models/database/document.py +++ b/packages/shared-python/shared/models/database/document.py @@ -270,6 +270,12 @@ class DocumentMapUnit(Base): path_token_count: Mapped[int] = mapped_column(Integer, nullable=False) content_token_count: Mapped[int] = mapped_column(Integer, nullable=False) term_search_text_lower: Mapped[str] = mapped_column(Text, nullable=False) + # Asset presence under this unit's section, after root-asset remount + # (``KnowhereProvider._remount_root_assets``). Lets type-scoped queries + # (e.g. chunk_types=["image"]) narrow map-unit candidates *before* + # scoring instead of scoring everything and discarding after the fact. + has_image: Mapped[bool] = mapped_column(nullable=False, default=False) + has_table: Mapped[bool] = mapped_column(nullable=False, default=False) sort_order: Mapped[int] = mapped_column(Integer, nullable=False, default=0) created_at: Mapped[datetime] = mapped_column( DateTime, default=utc_now_naive, nullable=False @@ -284,6 +290,18 @@ class DocumentMapUnit(Base): "unit_id", ), Index("idx_document_map_units_section", "section_id"), + Index( + "idx_document_map_units_has_image", + "document_id", + "job_result_id", + postgresql_where=has_image.is_(True), + ), + Index( + "idx_document_map_units_has_table", + "document_id", + "job_result_id", + postgresql_where=has_table.is_(True), + ), ) diff --git a/packages/shared-python/shared/services/retrieval/execution/routes.py b/packages/shared-python/shared/services/retrieval/execution/routes.py index fe47be19..673e2b6e 100644 --- a/packages/shared-python/shared/services/retrieval/execution/routes.py +++ b/packages/shared-python/shared/services/retrieval/execution/routes.py @@ -7,7 +7,7 @@ from loguru import logger from sqlalchemy.ext.asyncio import AsyncSession -from shared.services.retrieval.search.discovery import bottom_discovery +from shared.services.retrieval.search.map_unit_discovery import map_unit_discovery from shared.services.retrieval.execution.reference_resolver import ( resolve_workflow_references, ) @@ -118,7 +118,7 @@ async def _try_run_small_corpus_route( async def _run_classic_topk_route( context: RetrievalRouteContext, ) -> RetrievalRouteOutcome: - discovery_result = await bottom_discovery( + discovery_result = await map_unit_discovery( context.db, user_id=context.user_id, namespace=context.namespace, @@ -129,9 +129,6 @@ async def _run_classic_topk_route( chunk_types=context.allowed_chunk_types, signal_paths=context.signal_paths, filter_mode=context.filter_mode, - channels=context.channels, - channel_weights=context.channel_weights, - internal_recall_k=context.internal_recall_k, revision_pins=context.revision_pins, ) diff --git a/packages/shared-python/shared/services/retrieval/map_unit_index.py b/packages/shared-python/shared/services/retrieval/map_unit_index.py index 15454dc9..4171861c 100644 --- a/packages/shared-python/shared/services/retrieval/map_unit_index.py +++ b/packages/shared-python/shared/services/retrieval/map_unit_index.py @@ -93,6 +93,11 @@ def replace_document_map_units( map_unit_id = f"dmu_{uuid4().hex}" path_tokens = str(unit.get("path_search_text") or "").split() content_tokens = str(unit.get("content_search_text") or "").split() + # ``provider.self_units`` already reflects root-asset remount (assets + # referenced via ``connect_to`` are moved onto the text section that + # embeds them), so this is the same ownership the query-time scorer + # sees, not a new computation. + section_chunk_types = {u.chunk_type for u in provider.self_units(section_id)} db.add( DocumentMapUnit( id=map_unit_id, @@ -104,6 +109,8 @@ def replace_document_map_units( path_token_count=len(path_tokens), content_token_count=len(content_tokens), term_search_text_lower=str(unit.get("term_search_text") or "").lower(), + has_image="image" in section_chunk_types, + has_table="table" in section_chunk_types, sort_order=sort_order, ) ) diff --git a/packages/shared-python/shared/services/retrieval/nav/knowhere_hybrid.py b/packages/shared-python/shared/services/retrieval/nav/knowhere_hybrid.py index 3c237aef..a937210d 100644 --- a/packages/shared-python/shared/services/retrieval/nav/knowhere_hybrid.py +++ b/packages/shared-python/shared/services/retrieval/nav/knowhere_hybrid.py @@ -1,7 +1,6 @@ """KnowWhere-style 3-channel hybrid retrieval (path BM25 + content BM25 + term). -Ported from Ontos-AI/knowhere: - packages/shared-python/shared/services/retrieval/search/{scoring,lexical_ranker,channels}.py +Ported from Ontos-AI/knowhere map-unit BM25 (path + content). Reference: https://github.com/Ontos-AI/knowhere """ diff --git a/packages/shared-python/shared/services/retrieval/search/channels.py b/packages/shared-python/shared/services/retrieval/search/channels.py deleted file mode 100644 index cdf67326..00000000 --- a/packages/shared-python/shared/services/retrieval/search/channels.py +++ /dev/null @@ -1,529 +0,0 @@ -""" -Independent retrieval channels for checkerboard search. - -Each channel queries the full scoped corpus independently and returns -ranked rows. Channels are fused via RRF in the orchestrator. -""" - -from __future__ import annotations - -import time -from collections.abc import Mapping -from typing import Any - -from loguru import logger -from sqlalchemy import text -from sqlalchemy.ext.asyncio import AsyncSession - -from shared.core.config import settings -from shared.services.retrieval.search.lexical_ranker import ( - rank_rows_by_bm25, - tokenize_query_for_ranker, -) -from shared.services.retrieval.search.section_filters import is_excluded_section - -# The generated tsvector columns are built with to_tsvector('simple', ...), so -# queries must use the same configuration or nothing matches. -_FTS_CONFIG = "simple" - -# Guards against pathological queries producing an enormous tsquery. -_MAX_FTS_QUERY_TOKENS = 50 - -_TSV_FIELD_BY_SEARCH_FIELD = { - "content_search_text": "content_search_tsv", - "path_search_text": "path_search_tsv", -} - - -_SCOPED_CORPUS_CTE = """ -WITH scoped_chunks AS ( - SELECT - dc.id, - dc.chunk_id, - dc.document_id, - dc.section_id, - dc.chunk_type, - dc.content, - dc.source_chunk_path, - dc.file_path, - dc.chunk_metadata, - dc.job_result_id, - dc.sort_order, - dc.content_search_text, - dc.content_search_tsv, - dc.path_search_text, - dc.path_search_tsv, - dc.term_search_text, - d.source_file_name, - d.user_id, - d.namespace, - ds.section_path, - jr.job_id - FROM document_chunks dc - JOIN documents d - ON d.document_id = dc.document_id - {revision_join} - LEFT JOIN document_sections ds - ON ds.section_id = dc.section_id - JOIN job_results jr - ON jr.id = dc.job_result_id - WHERE d.user_id = :user_id - AND d.namespace = :namespace - AND d.status = 'active' - {revision_clause} - {exclude_clause} - {extra_filters} -) -""" - - -def _build_revision_scope( - revision_pins: Mapping[str, str] | None, -) -> tuple[str, str, dict[str, Any]]: - if revision_pins is None: - return ( - "AND d.current_job_result_id = dc.job_result_id", - "", - {}, - ) - - pairs = [ - (str(document_id).strip(), str(job_result_id).strip()) - for document_id, job_result_id in revision_pins.items() - if str(document_id).strip() and str(job_result_id).strip() - ] - if not pairs: - return "", "AND FALSE", {} - - params: dict[str, Any] = {} - placeholders: list[str] = [] - for index, (document_id, job_result_id) in enumerate(pairs): - document_key = f"_pin_document_{index}" - revision_key = f"_pin_revision_{index}" - placeholders.append(f"(:{document_key}, :{revision_key})") - params[document_key] = document_id - params[revision_key] = job_result_id - return "", f"AND (dc.document_id, dc.job_result_id) IN ({', '.join(placeholders)})", params - - -def _build_exclude_clause(exclude_document_ids: list[str]) -> str: - if not exclude_document_ids: - return "" - # Use PostgreSQL array ANY() to avoid asyncpg tuple-binding pitfalls with - # raw text() + `NOT IN :param` (which asyncpg treats as a record parameter - # and fails with a syntax error). - return "AND d.document_id <> ALL(:excluded_doc_ids)" - - -def _build_base_params( - *, - user_id: str, - namespace: str, - exclude_document_ids: list[str], -) -> dict[str, Any]: - params: dict[str, Any] = { - "user_id": user_id, - "namespace": namespace, - } - if exclude_document_ids: - params["excluded_doc_ids"] = list(exclude_document_ids) - return params - - -def _build_extra_filters( - *, - allowed_chunk_types: set[str] | None, - signal_paths: list[str], - filter_mode: str, -) -> tuple[str, dict[str, Any]]: - """Build additional SQL WHERE clauses for chunk_types and signal_path filtering.""" - clauses: list[str] = [] - params: dict[str, Any] = {} - - if allowed_chunk_types is not None: - placeholders = ", ".join(f":_act_{i}" for i in range(len(allowed_chunk_types))) - clauses.append(f"AND LOWER(dc.chunk_type) IN ({placeholders})") - for i, ct in enumerate(sorted(allowed_chunk_types)): - params[f"_act_{i}"] = ct - - if signal_paths: - # TODO(intent-step): Current implementation uses OR across - # signal_paths keywords. The Intent Step will need hierarchical - # AND (prefix) matching, e.g. signal_paths=["第一章/1.1/（2）"] - # should match only paths containing ALL segments in order. - # Consider adding a `filter_strategy` param: "keyword_or" (current) - # vs "path_prefix" (for Intent Step resolved paths). - ilike_parts = [] - for i, kw in enumerate(signal_paths): - key = f"_sig_{i}" - ilike_parts.append(f"LOWER(COALESCE(ds.section_path, '')) LIKE :{key}") - params[key] = f"%{kw.lower()}%" - combined = " OR ".join(ilike_parts) - if filter_mode == "keep": - clauses.append(f"AND ({combined})") - else: - clauses.append(f"AND NOT ({combined})") - - return "\n ".join(clauses), params - - -def _build_exclude_section_filters( - *, - exclude_sections: list[dict[str, str]], -) -> tuple[str, dict[str, Any]]: - """Exclude an exact section path and its descendants inside the scoped CTE. - - Applied before the FTS candidate LIMIT so excluded sections cannot consume - the bounded candidate budget. Sectionless chunks stay eligible (empty path - does not match), matching ``is_excluded_section``. - """ - clauses: list[str] = [] - params: dict[str, Any] = {} - - for index, item in enumerate(exclude_sections): - if not isinstance(item, dict): - continue - document_id = str(item.get("document_id") or "").strip() - section_path = str(item.get("section_path") or "").strip() - if not document_id or not section_path: - continue - - document_key = f"_exc_section_doc_{index}" - path_key = f"_exc_section_path_{index}" - clauses.append( - f"""AND NOT ( - dc.document_id = :{document_key} - AND ( - COALESCE(ds.section_path, '') = :{path_key} - OR POSITION(:{path_key} || ' / ' IN COALESCE(ds.section_path, '')) = 1 - ) - )""" - ) - params[document_key] = document_id - params[path_key] = section_path - - return "\n ".join(clauses), params - - -def _join_sql_filters(*filters: str) -> str: - return "\n ".join(filter(None, filters)) - - -def _prepare_fts_tokens(tokens: list[str]) -> list[str]: - """Return the ranker tokens to hand to the Postgres FTS prefilter. - - Tokens are passed to SQL as a text[] parameter and lexed by Postgres - itself, so nothing here needs to escape tsquery syntax. Only emptiness and - an upper bound are enforced. - """ - prepared: list[str] = [] - for token in tokens: - cleaned = token.strip() - if cleaned: - prepared.append(cleaned) - if len(prepared) >= _MAX_FTS_QUERY_TOKENS: - break - return prepared - - -def _row_to_dict(row: Any) -> dict[str, Any]: - return dict(row._mapping) - - -def _filter_excluded_sections( - rows: list[dict[str, Any]], - exclude_sections: list[dict[str, str]], -) -> list[dict[str, Any]]: - if not exclude_sections: - return rows - return [ - row - for row in rows - if not is_excluded_section( - document_id=row.get("document_id"), - section_path=row.get("section_path"), - exclude_sections=exclude_sections, - ) - ] - - -async def path_channel( - db: AsyncSession, - *, - user_id: str, - namespace: str, - query: str, - top_k: int, - exclude_document_ids: list[str], - exclude_sections: list[dict[str, str]], - allowed_chunk_types: set[str] | None = None, - signal_paths: list[str] | None = None, - filter_mode: str = "delete", - revision_pins: Mapping[str, str] | None = None, -) -> list[dict[str, Any]]: - """Path channel: BM25 over pre-tokenized path search text. - - This keeps the channel useful when vector search is unavailable. A future - vector score can be fused on top of the returned BM25 score. - """ - return await _bm25_channel( - db, - user_id=user_id, - namespace=namespace, - query=query, - top_k=top_k, - exclude_document_ids=exclude_document_ids, - exclude_sections=exclude_sections, - allowed_chunk_types=allowed_chunk_types, - signal_paths=signal_paths, - filter_mode=filter_mode, - search_field="path_search_text", - revision_pins=revision_pins, - ) - - -async def content_channel( - db: AsyncSession, - *, - user_id: str, - namespace: str, - query: str, - top_k: int, - exclude_document_ids: list[str], - exclude_sections: list[dict[str, str]], - allowed_chunk_types: set[str] | None = None, - signal_paths: list[str] | None = None, - filter_mode: str = "delete", - revision_pins: Mapping[str, str] | None = None, -) -> list[dict[str, Any]]: - """Content channel: BM25 over pre-tokenized content search text.""" - return await _bm25_channel( - db, - user_id=user_id, - namespace=namespace, - query=query, - top_k=top_k, - exclude_document_ids=exclude_document_ids, - exclude_sections=exclude_sections, - allowed_chunk_types=allowed_chunk_types, - signal_paths=signal_paths, - filter_mode=filter_mode, - search_field="content_search_text", - revision_pins=revision_pins, - ) - - -async def _bm25_channel( - db: AsyncSession, - *, - user_id: str, - namespace: str, - query: str, - top_k: int, - exclude_document_ids: list[str], - exclude_sections: list[dict[str, str]], - allowed_chunk_types: set[str] | None, - signal_paths: list[str] | None, - filter_mode: str, - search_field: str, - revision_pins: Mapping[str, str] | None, -) -> list[dict[str, Any]]: - if search_field not in {"content_search_text", "path_search_text"}: - raise ValueError(f"Unsupported search_field: {search_field}") - - query_tokens = tokenize_query_for_ranker(query) - if not query_tokens: - return [] - - exclude_clause = _build_exclude_clause(exclude_document_ids) - extra_sql, extra_params = _build_extra_filters( - allowed_chunk_types=allowed_chunk_types, - signal_paths=signal_paths or [], - filter_mode=filter_mode, - ) - section_sql, section_params = _build_exclude_section_filters( - exclude_sections=exclude_sections, - ) - extra_sql = _join_sql_filters(extra_sql, section_sql) - params = _build_base_params( - user_id=user_id, - namespace=namespace, - exclude_document_ids=exclude_document_ids, - ) - params.update(extra_params) - params.update(section_params) - - revision_join, revision_clause, revision_params = _build_revision_scope( - revision_pins - ) - params.update(revision_params) - - corpus_cte = _SCOPED_CORPUS_CTE.format( - revision_join=revision_join, - revision_clause=revision_clause, - exclude_clause=exclude_clause, - extra_filters=extra_sql, - ) - full_scan_sql = ( - corpus_cte - + f""" - SELECT sc.* - FROM scoped_chunks sc - WHERE COALESCE(sc.{search_field}, '') <> '' - """ - ) - - started_at = time.perf_counter() - tsv_field = _TSV_FIELD_BY_SEARCH_FIELD[search_field] - fts_tokens = _prepare_fts_tokens(query_tokens) - candidate_limit = int(settings.RETRIEVAL_POSTGRES_FTS_CANDIDATE_LIMIT) - - rows: list[dict[str, Any]] = [] - used_fallback = True - if fts_tokens: - # Postgres lexes the tokens with the same configuration that generated - # the tsvector columns, then ORs the resulting lexemes. Building the - # tsquery server-side keeps the prefilter aligned with the stored - # lexicon and leaves no room for tsquery syntax in user input to - # change the query shape. `fts_query.q` is NULL when no token yields a - # lexeme, which the caller treats as "no usable prefilter". - prefilter_sql = ( - corpus_cte - + f""", - fts_query AS ( - SELECT string_agg(quote_literal(lexeme), ' | ')::tsquery AS q - FROM ( - SELECT DISTINCT - unnest(tsvector_to_array(to_tsvector('{_FTS_CONFIG}', token))) AS lexeme - FROM unnest(CAST(:fts_tokens AS text[])) AS token - ) lexemes - ) - SELECT sc.* - FROM scoped_chunks sc, fts_query fq - WHERE COALESCE(sc.{search_field}, '') <> '' - AND fq.q IS NOT NULL - AND sc.{tsv_field} @@ fq.q - ORDER BY ts_rank_cd(sc.{tsv_field}, fq.q) DESC - LIMIT :fts_candidate_limit - """ - ) - prefilter_params = dict(params) - prefilter_params["fts_tokens"] = fts_tokens - prefilter_params["fts_candidate_limit"] = candidate_limit - result = await db.execute(text(prefilter_sql), prefilter_params) - rows = [_row_to_dict(r) for r in result.all()] - used_fallback = not rows - - # No usable tsquery, or the prefilter matched nothing. Fall back to the - # full scoped scan so recall never regresses against the previous - # behaviour. - if used_fallback: - result = await db.execute(text(full_scan_sql), params) - rows = [_row_to_dict(r) for r in result.all()] - - candidate_count = len(rows) - # Defensive: SQL already owns pre-LIMIT section exclusion. - rows = _filter_excluded_sections(rows, exclude_sections) - - ranked_rows = rank_rows_by_bm25(rows, query_tokens, search_field=search_field) - ranked_rows = ranked_rows[:top_k] - - logger.debug( - "bm25_channel field={} candidates={} limit={} ranked={} fallback={} duration_ms={:.1f}", - search_field, - candidate_count, - candidate_limit, - len(ranked_rows), - used_fallback, - (time.perf_counter() - started_at) * 1000, - ) - return ranked_rows - - -async def term_channel( - db: AsyncSession, - *, - user_id: str, - namespace: str, - query: str, - top_k: int, - exclude_document_ids: list[str], - exclude_sections: list[dict[str, str]], - allowed_chunk_types: set[str] | None = None, - signal_paths: list[str] | None = None, - filter_mode: str = "delete", - revision_pins: Mapping[str, str] | None = None, -) -> list[dict[str, Any]]: - """Term/grep channel: substring matching on term_search_text. - - Aligned with KB checkerboard_find() term channel (grep_search()): - exact substring match on content + path, scoring by hit count. - - Note: top_k is already effective_recall_k from app_service. - """ - query_lower = query.lower().strip() - query_tokens = tokenize_query_for_ranker(query) - if not query_lower or not query_tokens: - return [] - - exclude_clause = _build_exclude_clause(exclude_document_ids) - extra_sql, extra_params = _build_extra_filters( - allowed_chunk_types=allowed_chunk_types, - signal_paths=signal_paths or [], - filter_mode=filter_mode, - ) - params = _build_base_params( - user_id=user_id, - namespace=namespace, - exclude_document_ids=exclude_document_ids, - ) - params.update(extra_params) - - revision_join, revision_clause, revision_params = _build_revision_scope( - revision_pins - ) - params.update(revision_params) - - ilike_conditions = [] - for i, unit in enumerate(query_tokens): - param_key = f"unit_{i}" - ilike_conditions.append(f"LOWER(sc.term_search_text) LIKE :{param_key}") - params[param_key] = f"%{unit}%" - - if not ilike_conditions: - ilike_conditions.append("LOWER(sc.term_search_text) LIKE :full_query") - params["full_query"] = f"%{query_lower}%" - - where_clause = " OR ".join(ilike_conditions) - sql = ( - _SCOPED_CORPUS_CTE.format( - revision_join=revision_join, - revision_clause=revision_clause, - exclude_clause=exclude_clause, extra_filters=extra_sql - ) - + f""" - SELECT sc.* - FROM scoped_chunks sc - WHERE sc.term_search_text IS NOT NULL - AND ({where_clause}) - """ - ) - - result = await db.execute(text(sql), params) - rows = [_row_to_dict(r) for r in result.all()] - rows = _filter_excluded_sections(rows, exclude_sections) - - scored: list[dict[str, Any]] = [] - for row in rows: - haystack = (row.get("term_search_text") or "").lower() - if query_lower in haystack: - row["score"] = 100.0 - scored.append(row) - else: - hit_count = sum(1 for u in query_tokens if u in haystack) - if hit_count > 0: - row["score"] = float(hit_count) - scored.append(row) - - scored.sort(key=lambda r: r["score"], reverse=True) - return scored[:top_k] diff --git a/packages/shared-python/shared/services/retrieval/search/discovery.py b/packages/shared-python/shared/services/retrieval/search/discovery.py deleted file mode 100644 index d71ccdd9..00000000 --- a/packages/shared-python/shared/services/retrieval/search/discovery.py +++ /dev/null @@ -1,216 +0,0 @@ -"""Classic 3-channel bottom discovery (moved from agentic/discovery/tools). - -Input (all keyword, via classic route context): - db, user_id, namespace, query, top_k, - exclude_document_ids, exclude_sections, - chunk_types?, signal_paths?, filter_mode, - channels?, channel_weights?, internal_recall_k? - -Output: - DiscoveryResult - status: ``discovery_done`` | ``error`` - payload.fused_rows: RRF-merged scored rows (classic reads this) - payload.top_doc_ids / channel_counts: diagnostics - latency_ms, error? - -Does not call LLM. Does not touch map-nav / wallet / agentic budget. -""" - -from __future__ import annotations - -import asyncio -import time -from collections.abc import Awaitable, Callable, Mapping -from dataclasses import dataclass, field -from typing import Any - -from loguru import logger -from sqlalchemy.ext.asyncio import AsyncSession - -from shared.services.retrieval.search.channels import ( - content_channel, - path_channel, - term_channel, -) -from shared.services.retrieval.search.scoring import ( - merge_channels_rrf, - merge_same_section_rows, - normalize_row_scores, -) -from shared.services.retrieval.settings import ( - CHANNEL_WEIGHT_CONTENT, - CHANNEL_WEIGHT_PATH, - CHANNEL_WEIGHT_TERM, - INTERNAL_RECALL_K_MULTIPLIER, -) - - -@dataclass -class DiscoveryResult: - """Return shape for classic bottom discovery (replaces agentic ToolResult).""" - - status: str - payload: dict[str, Any] = field(default_factory=dict) - latency_ms: int = 0 - error: str | None = None - - -async def bottom_discovery( - db: AsyncSession, - *, - user_id: str, - namespace: str, - query: str, - top_k: int, - exclude_document_ids: list[str], - exclude_sections: list[dict[str, str]], - chunk_types: set[str] | None = None, - signal_paths: list[str] | None = None, - filter_mode: str = "delete", - channels: list[str] | None = None, - channel_weights: dict[str, float] | None = None, - internal_recall_k: int | None = None, - revision_pins: Mapping[str, str] | None = None, - **_kwargs: Any, -) -> DiscoveryResult: - """Run 3-channel BM25 discovery plus RRF fusion.""" - t0 = time.monotonic() - try: - del db - allowed_chunk_types = chunk_types - effective_recall_k = ( - internal_recall_k - if internal_recall_k is not None - else top_k * INTERNAL_RECALL_K_MULTIPLIER - ) - active_channels = set(channels) if channels else {"path", "content", "term"} - - path_rows, content_rows, term_rows = await asyncio.gather( - _run_channel( - path_channel, - enabled="path" in active_channels, - user_id=user_id, - namespace=namespace, - query=query, - top_k=effective_recall_k, - exclude_document_ids=exclude_document_ids, - exclude_sections=exclude_sections, - allowed_chunk_types=allowed_chunk_types, - signal_paths=signal_paths, - filter_mode=filter_mode, - revision_pins=revision_pins, - ), - _run_channel( - content_channel, - enabled="content" in active_channels, - user_id=user_id, - namespace=namespace, - query=query, - top_k=effective_recall_k, - exclude_document_ids=exclude_document_ids, - exclude_sections=exclude_sections, - allowed_chunk_types=allowed_chunk_types, - signal_paths=signal_paths, - filter_mode=filter_mode, - revision_pins=revision_pins, - ), - _run_channel( - term_channel, - enabled="term" in active_channels, - user_id=user_id, - namespace=namespace, - query=query, - top_k=effective_recall_k, - exclude_document_ids=exclude_document_ids, - exclude_sections=exclude_sections, - allowed_chunk_types=allowed_chunk_types, - signal_paths=signal_paths, - filter_mode=filter_mode, - revision_pins=revision_pins, - ), - ) - - default_weights = { - "path": CHANNEL_WEIGHT_PATH, - "content": CHANNEL_WEIGHT_CONTENT, - "term": CHANNEL_WEIGHT_TERM, - } - effective_weights = {**default_weights, **(channel_weights or {})} - - channel_lists: list[list[dict[str, Any]]] = [] - weight_list: list[float] = [] - if path_rows: - channel_lists.append(path_rows) - weight_list.append(effective_weights.get("path", CHANNEL_WEIGHT_PATH)) - if content_rows: - channel_lists.append(content_rows) - weight_list.append(effective_weights.get("content", CHANNEL_WEIGHT_CONTENT)) - if term_rows: - channel_lists.append(term_rows) - weight_list.append(effective_weights.get("term", CHANNEL_WEIGHT_TERM)) - - fused_rows = ( - merge_channels_rrf(channel_lists, weight_list, effective_recall_k) - if channel_lists - else [] - ) - fused_rows = merge_same_section_rows(fused_rows) - - if fused_rows: - normalize_row_scores( - fused_rows, - source_field="score", - target_field="discovery_score", - default=0.5, - ) - - doc_id_counts: dict[str, int] = {} - for row in fused_rows: - did = row.get("document_id", "") - if did: - doc_id_counts[did] = doc_id_counts.get(did, 0) + 1 - top_doc_ids = sorted( - doc_id_counts, - key=lambda document_id: doc_id_counts[document_id], - reverse=True, - )[:5] - - latency = int((time.monotonic() - t0) * 1000) - logger.info( - f" search.bottom_discovery: {len(fused_rows)} fused rows, " - f"top_doc_ids={top_doc_ids}, {latency}ms" - ) - return DiscoveryResult( - status="discovery_done", - payload={ - "fused_rows": fused_rows, - "top_doc_ids": top_doc_ids, - "channel_counts": { - "path": len(path_rows), - "content": len(content_rows), - "term": len(term_rows), - }, - }, - latency_ms=latency, - ) - except Exception as exc: - latency = int((time.monotonic() - t0) * 1000) - logger.error(f" search.bottom_discovery failed: {exc}") - return DiscoveryResult(status="error", error=str(exc), latency_ms=latency) - - -async def _run_channel( - channel: Callable[..., Awaitable[list[dict[str, Any]]]], - *, - enabled: bool, - **kwargs: Any, -) -> list[dict[str, Any]]: - if not enabled: - return [] - - # Import lazily so discovery remains usable by lightweight contract tests - # without creating a database context until a channel is actually enabled. - from shared.core.database import get_db_context - - async with get_db_context() as channel_db: - return await channel(channel_db, **kwargs) diff --git a/packages/shared-python/shared/services/retrieval/search/lexical_ranker.py b/packages/shared-python/shared/services/retrieval/search/lexical_ranker.py deleted file mode 100644 index 8b50cd40..00000000 --- a/packages/shared-python/shared/services/retrieval/search/lexical_ranker.py +++ /dev/null @@ -1,74 +0,0 @@ -from __future__ import annotations - -from typing import Any - -from loguru import logger - -from shared.utils.text_utils import tokenize_for_retrieval - - -def tokenize_query_for_ranker(query: str) -> list[str]: - return tokenize_for_retrieval(query, dedupe=True) - - -def rank_rows_by_bm25( - rows: list[dict[str, Any]], - query_tokens: list[str], - *, - search_field: str, -) -> list[dict[str, Any]]: - """Rank matching rows with BM25 over pre-tokenized search text.""" - try: - from rank_bm25 import BM25Okapi - except ImportError: - return _rank_rows_by_token_overlap( - rows, - query_tokens, - search_field=search_field, - ) - - corpus: list[list[str]] = [] - ranked_rows: list[dict[str, Any]] = [] - query_token_set = set(query_tokens) - for row in rows: - tokens = _get_search_tokens(row, search_field=search_field) - if not tokens or not query_token_set.intersection(tokens): - continue - corpus.append(tokens) - ranked_rows.append(row) - - if not corpus or not query_tokens: - return [] - - bm25 = BM25Okapi(corpus) - scores = bm25.get_scores(query_tokens) - - for index, row in enumerate(ranked_rows): - row["score"] = float(scores[index]) - - ranked_rows.sort(key=lambda row: row["score"], reverse=True) - return ranked_rows - - -def _rank_rows_by_token_overlap( - rows: list[dict[str, Any]], - query_tokens: list[str], - *, - search_field: str, -) -> list[dict[str, Any]]: - logger.warning("rank_bm25 not installed, skipping BM25 re-rank") - ranked_rows: list[dict[str, Any]] = [] - query_token_set = set(query_tokens) - for row in rows: - tokens = _get_search_tokens(row, search_field=search_field) - overlap = len(query_token_set.intersection(tokens)) - if overlap <= 0: - continue - row["score"] = float(overlap) - ranked_rows.append(row) - ranked_rows.sort(key=lambda row: row["score"], reverse=True) - return ranked_rows - - -def _get_search_tokens(row: dict[str, Any], *, search_field: str) -> list[str]: - return [token for token in str(row.get(search_field) or "").split() if token] diff --git a/packages/shared-python/shared/services/retrieval/search/map_unit_discovery.py b/packages/shared-python/shared/services/retrieval/search/map_unit_discovery.py new file mode 100644 index 00000000..b75649fb --- /dev/null +++ b/packages/shared-python/shared/services/retrieval/search/map_unit_discovery.py @@ -0,0 +1,500 @@ +"""Classic-route discovery via the persisted map-unit BM25 scorer. + +Replaces the retired chunk-level 3-channel SQL scan. Scoring uses the same +``score_persisted_corpus_many`` formula as map-nav (path + content only). + +``chunk_types`` is optional: omitted means score every in-scope unit. When +the request is image/table only, ``has_image`` / ``has_table`` (written at +index time) narrow candidates before scoring. + +Hydration returns one primary chunk per winning unit. Asset-only requests +then follow ``connect_to`` to that unit's image/table row. Downstream +``assemble_retrieval_results`` still inlines connected assets for text hits. +""" + +from __future__ import annotations + +import json +import time +from collections.abc import Mapping +from dataclasses import dataclass, field +from typing import Any + +from loguru import logger +from sqlalchemy import text +from sqlalchemy.ext.asyncio import AsyncSession + +from shared.services.retrieval.hydration.connected import hydrate_connected_target_rows +from shared.services.retrieval.hydration.row_utils import ( + iter_connected_target_ids, + normalize_chunk_type, +) +from shared.services.retrieval.nav.knowhere_hybrid import ( + PersistedBm25Stats, + PersistedScoreCorpus, + PersistedScoreUnit, + score_persisted_corpus_many, + tokenize_query_for_ranker, +) +from shared.services.retrieval.search.scoring import normalize_row_scores +from shared.services.retrieval.search.section_filters import is_excluded_section +from shared.services.retrieval.settings import ASSET_CHUNK_TYPES + +_MAP_SCORE_CHANNELS = ("path", "content") + +_SCOPED_UNITS_CTE = """ +WITH scoped_units AS ( + SELECT + dmu.id AS map_unit_id, + dmu.document_id, + dmu.job_result_id, + dmu.section_id, + dmu.path_token_count, + dmu.content_token_count, + ds.section_path + FROM document_map_units dmu + JOIN documents d + ON d.document_id = dmu.document_id + {revision_join} + JOIN document_sections ds + ON ds.section_id = dmu.section_id + WHERE d.user_id = :user_id + AND d.namespace = :namespace + AND d.status = 'active' + {revision_clause} + {exclude_clause} + {type_clause} + {signal_clause} +) +""" + + +@dataclass +class DiscoveryResult: + status: str + payload: dict[str, Any] = field(default_factory=dict) + latency_ms: int = 0 + error: str | None = None + + +def _build_revision_scope( + revision_pins: Mapping[str, str] | None, +) -> tuple[str, str, dict[str, Any]]: + if revision_pins is None: + return ("AND d.current_job_result_id = dmu.job_result_id", "", {}) + pairs = [ + (str(document_id).strip(), str(job_result_id).strip()) + for document_id, job_result_id in revision_pins.items() + if str(document_id).strip() and str(job_result_id).strip() + ] + if not pairs: + return "", "AND FALSE", {} + params: dict[str, Any] = {} + placeholders: list[str] = [] + for index, (document_id, job_result_id) in enumerate(pairs): + document_key = f"_pin_document_{index}" + revision_key = f"_pin_revision_{index}" + placeholders.append(f"(:{document_key}, :{revision_key})") + params[document_key] = document_id + params[revision_key] = job_result_id + return ( + "", + f"AND (dmu.document_id, dmu.job_result_id) IN ({', '.join(placeholders)})", + params, + ) + + +def _build_type_clause( + allowed_chunk_types: set[str] | None, +) -> tuple[str, dict[str, Any]]: + if allowed_chunk_types is None or not allowed_chunk_types.issubset( + ASSET_CHUNK_TYPES + ): + return "", {} + clauses = [] + if "image" in allowed_chunk_types: + clauses.append("dmu.has_image") + if "table" in allowed_chunk_types: + clauses.append("dmu.has_table") + if not clauses: + return "AND FALSE", {} + return f"AND ({' OR '.join(clauses)})", {} + + +def _build_exclude_clause(exclude_document_ids: list[str]) -> tuple[str, dict[str, Any]]: + if not exclude_document_ids: + return "", {} + return "AND d.document_id <> ALL(:excluded_doc_ids)", { + "excluded_doc_ids": list(exclude_document_ids) + } + + +def _build_signal_clause( + signal_paths: list[str], filter_mode: str +) -> tuple[str, dict[str, Any]]: + if not signal_paths: + return "", {} + ilike_parts = [] + params: dict[str, Any] = {} + for index, keyword in enumerate(signal_paths): + key = f"_sig_{index}" + ilike_parts.append(f"LOWER(COALESCE(ds.section_path, '')) LIKE :{key}") + params[key] = f"%{keyword.lower()}%" + combined = " OR ".join(ilike_parts) + clause = f"AND ({combined})" if filter_mode == "keep" else f"AND NOT ({combined})" + return clause, params + + +async def map_unit_discovery( + db: AsyncSession | None, + *, + user_id: str, + namespace: str, + query: str, + top_k: int, + exclude_document_ids: list[str], + exclude_sections: list[dict[str, str]], + chunk_types: set[str] | None = None, + signal_paths: list[str] | None = None, + filter_mode: str = "delete", + revision_pins: Mapping[str, str] | None = None, + **_kwargs: Any, +) -> DiscoveryResult: + """Score the whole in-scope corpus via the persisted map-unit BM25 scorer.""" + t0 = time.monotonic() + try: + if db is None: + return DiscoveryResult( + status="error", + error="database session required", + latency_ms=0, + ) + query_tokens = tokenize_query_for_ranker(query) + if not query_tokens: + return DiscoveryResult(status="discovery_done", payload={"fused_rows": []}) + + revision_join, revision_clause, revision_params = _build_revision_scope( + revision_pins + ) + exclude_clause, exclude_params = _build_exclude_clause(exclude_document_ids) + type_clause, type_params = _build_type_clause(chunk_types) + signal_clause, signal_params = _build_signal_clause( + signal_paths or [], filter_mode + ) + params: dict[str, Any] = {"user_id": user_id, "namespace": namespace} + params.update(revision_params) + params.update(exclude_params) + params.update(type_params) + params.update(signal_params) + + cte = _SCOPED_UNITS_CTE.format( + revision_join=revision_join, + revision_clause=revision_clause, + exclude_clause=exclude_clause, + type_clause=type_clause, + signal_clause=signal_clause, + ) + unit_result = await db.execute(text(cte + "SELECT * FROM scoped_units"), params) + unit_rows = [dict(row._mapping) for row in unit_result.all()] + unit_rows = [ + row + for row in unit_rows + if not is_excluded_section( + document_id=row.get("document_id"), + section_path=row.get("section_path"), + exclude_sections=exclude_sections, + ) + ] + + if not unit_rows: + return DiscoveryResult(status="discovery_done", payload={"fused_rows": []}) + + map_unit_ids = [row["map_unit_id"] for row in unit_rows] + frequency_result = await db.execute( + text( + "SELECT map_unit_id, channel, token, frequency " + "FROM document_map_unit_tokens " + "WHERE map_unit_id = ANY(:unit_ids) " + "AND channel = ANY(:channels) " + "AND token = ANY(:tokens)" + ), + { + "unit_ids": map_unit_ids, + "channels": list(_MAP_SCORE_CHANNELS), + "tokens": query_tokens, + }, + ) + frequencies: dict[tuple[str, str], dict[str, int]] = {} + for map_unit_id, channel, token, frequency in frequency_result.all(): + frequencies.setdefault((str(map_unit_id), str(channel)), {})[str(token)] = ( + int(frequency) + ) + + path_stats = await _build_bm25_stats( + db, + unit_rows=unit_rows, + map_unit_ids=map_unit_ids, + channel="path", + query_tokens=query_tokens, + frequencies=frequencies, + length_field="path_token_count", + ) + content_stats = await _build_bm25_stats( + db, + unit_rows=unit_rows, + map_unit_ids=map_unit_ids, + channel="content", + query_tokens=query_tokens, + frequencies=frequencies, + length_field="content_token_count", + ) + + corpus = PersistedScoreCorpus( + units=[ + PersistedScoreUnit( + unit_id=row["map_unit_id"], + path_length=int(row["path_token_count"]), + content_length=int(row["content_token_count"]), + path_frequencies=frequencies.get((row["map_unit_id"], "path"), {}), + content_frequencies=frequencies.get( + (row["map_unit_id"], "content"), {} + ), + ) + for row in unit_rows + ], + path_stats=path_stats, + content_stats=content_stats, + ) + scores_by_unit = score_persisted_corpus_many(corpus, [query]).get(query, {}) + + rows_by_unit_id = {row["map_unit_id"]: row for row in unit_rows} + + def _unit_has_query_hit(unit_id: str) -> bool: + # BM25 fused score can be 0 when IDF is 0 (tiny corpus). A unit + # that actually holds a query token is still a hit. + return any( + frequencies.get((unit_id, channel), {}).get(token, 0) > 0 + for channel in _MAP_SCORE_CHANNELS + for token in query_tokens + ) + + ranked_unit_ids = sorted( + (unit_id for unit_id in scores_by_unit if _unit_has_query_hit(unit_id)), + key=lambda unit_id: scores_by_unit[unit_id], + reverse=True, + )[:top_k] + + fused_rows = await _hydrate_winning_units( + db, + ranked_unit_ids=ranked_unit_ids, + rows_by_unit_id=rows_by_unit_id, + scores_by_unit=scores_by_unit, + chunk_types=chunk_types, + exclude_document_ids=exclude_document_ids, + exclude_sections=exclude_sections, + revision_pins=revision_pins, + ) + if fused_rows: + normalize_row_scores( + fused_rows, + source_field="score", + target_field="discovery_score", + default=0.5, + ) + + latency = int((time.monotonic() - t0) * 1000) + logger.info( + " search.map_unit_discovery: {} units scored, {} fused rows, {}ms", + len(unit_rows), + len(fused_rows), + latency, + ) + return DiscoveryResult( + status="discovery_done", + payload={"fused_rows": fused_rows}, + latency_ms=latency, + ) + except Exception as exc: + latency = int((time.monotonic() - t0) * 1000) + logger.error(f" search.map_unit_discovery failed: {exc}") + return DiscoveryResult(status="error", error=str(exc), latency_ms=latency) + + +async def _build_bm25_stats( + session: AsyncSession, + *, + unit_rows: list[dict[str, Any]], + map_unit_ids: list[str], + channel: str, + query_tokens: list[str], + frequencies: dict[tuple[str, str], dict[str, int]], + length_field: str, +) -> PersistedBm25Stats: + lengths = [ + int(row[length_field]) for row in unit_rows if int(row[length_field]) > 0 + ] + document_count = len(lengths) + document_frequency = { + token: sum( + 1 + for row in unit_rows + if frequencies.get((row["map_unit_id"], channel), {}).get(token, 0) > 0 + ) + for token in query_tokens + } + needs_average_idf = any( + frequency > document_count / 2 for frequency in document_frequency.values() + ) + average_idf = 0.0 + if needs_average_idf and map_unit_ids and document_count: + result = await session.execute( + text( + "SELECT COALESCE(AVG(LN((:document_count - frequencies.document_frequency " + "+ 0.5) / (frequencies.document_frequency + 0.5))), 0.0) " + "FROM (SELECT token, COUNT(*) AS document_frequency " + "FROM document_map_unit_tokens " + "WHERE map_unit_id = ANY(:unit_ids) AND channel = :channel " + "GROUP BY token) AS frequencies" + ), + { + "document_count": document_count, + "unit_ids": map_unit_ids, + "channel": channel, + }, + ) + row = result.first() + average_idf = float(row[0]) if row else 0.0 + return PersistedBm25Stats( + document_count=document_count, + total_length=sum(lengths), + document_frequency=document_frequency, + average_idf=average_idf, + ) + + +def _as_metadata_dict(value: object) -> dict[str, Any]: + if isinstance(value, dict): + return value + if isinstance(value, str) and value.strip(): + loaded = json.loads(value) + return loaded if isinstance(loaded, dict) else {} + return {} + + +def _primary_chunk(chunks: list[dict[str, Any]]) -> dict[str, Any] | None: + for chunk in chunks: + if normalize_chunk_type(chunk.get("chunk_type")) not in ASSET_CHUNK_TYPES: + return chunk + return chunks[0] if chunks else None + + +async def _hydrate_winning_units( + session: AsyncSession, + *, + ranked_unit_ids: list[str], + rows_by_unit_id: dict[str, dict[str, Any]], + scores_by_unit: dict[str, float], + chunk_types: set[str] | None, + exclude_document_ids: list[str], + exclude_sections: list[dict[str, str]], + revision_pins: Mapping[str, str] | None, +) -> list[dict[str, Any]]: + """Map each winning unit to one chunk; asset-only requests follow connect_to.""" + if not ranked_unit_ids: + return [] + + document_ids = sorted( + {rows_by_unit_id[unit_id]["document_id"] for unit_id in ranked_unit_ids} + ) + job_result_ids = sorted( + {rows_by_unit_id[unit_id]["job_result_id"] for unit_id in ranked_unit_ids} + ) + section_ids = sorted( + {rows_by_unit_id[unit_id]["section_id"] for unit_id in ranked_unit_ids} + ) + + result = await session.execute( + text( + "SELECT dc.chunk_id, dc.document_id, dc.section_id, dc.chunk_type, " + "dc.content, dc.source_chunk_path, dc.file_path, dc.chunk_metadata, " + "dc.job_result_id, dc.sort_order, ds.section_path, d.source_file_name, " + "jr.job_id " + "FROM document_chunks dc " + "JOIN documents d ON d.document_id = dc.document_id " + "LEFT JOIN document_sections ds ON ds.section_id = dc.section_id " + "LEFT JOIN job_results jr ON jr.id = dc.job_result_id " + "WHERE dc.document_id = ANY(:document_ids) " + "AND dc.job_result_id = ANY(:job_result_ids) " + "AND dc.section_id = ANY(:section_ids) " + "ORDER BY dc.sort_order, dc.chunk_id" + ), + { + "document_ids": document_ids, + "job_result_ids": job_result_ids, + "section_ids": section_ids, + }, + ) + chunk_rows_by_section: dict[tuple[str, str, str], list[dict[str, Any]]] = {} + for chunk_row in result.all(): + row = dict(chunk_row._mapping) + row["chunk_metadata"] = _as_metadata_dict(row.get("chunk_metadata")) + key = (row["document_id"], row["job_result_id"], row["section_id"]) + chunk_rows_by_section.setdefault(key, []).append(row) + + primaries: list[dict[str, Any]] = [] + for unit_id in ranked_unit_ids: + unit_row = rows_by_unit_id[unit_id] + key = ( + unit_row["document_id"], + unit_row["job_result_id"], + unit_row["section_id"], + ) + primary = _primary_chunk(chunk_rows_by_section.get(key, [])) + if primary is None: + continue + fused = dict(primary) + fused["score"] = scores_by_unit.get(unit_id, 0.0) + primaries.append(fused) + + if chunk_types is None or not chunk_types.issubset(ASSET_CHUNK_TYPES): + return primaries + + connected = await hydrate_connected_target_rows( + db=session, + rows=primaries, + exclude_document_ids=exclude_document_ids, + exclude_sections=exclude_sections, + revision_pins=revision_pins, + ) + connected_by_id = { + str(row.get("chunk_id") or ""): row + for row in connected + if row.get("chunk_id") + } + fused_rows: list[dict[str, Any]] = [] + seen: set[str] = set() + for primary in primaries: + score = float(primary.get("score") or 0.0) + key = ( + primary["document_id"], + primary["job_result_id"], + primary["section_id"], + ) + candidates = list(chunk_rows_by_section.get(key, ())) + for target_id in iter_connected_target_ids(primary): + target = connected_by_id.get(target_id) + if target is not None: + candidates.append(target) + for candidate in candidates: + chunk_id = str(candidate.get("chunk_id") or "") + if ( + not chunk_id + or chunk_id in seen + or normalize_chunk_type(candidate.get("chunk_type")) not in chunk_types + ): + continue + seen.add(chunk_id) + fused = dict(candidate) + fused["score"] = score + fused_rows.append(fused) + return fused_rows diff --git a/packages/shared-python/shared/services/retrieval/search/scoring.py b/packages/shared-python/shared/services/retrieval/search/scoring.py index 2746de34..af546ae8 100644 --- a/packages/shared-python/shared/services/retrieval/search/scoring.py +++ b/packages/shared-python/shared/services/retrieval/search/scoring.py @@ -10,35 +10,6 @@ def get_row_path(row: dict[str, Any]) -> str: return str(row.get('section_path') or row.get('source_chunk_path') or '') -def merge_same_section_rows(rows: list[dict[str, Any]]) -> list[dict[str, Any]]: - if not rows: - return rows - groups: dict[str, list[dict[str, Any]]] = {} - order: list[str] = [] - for row in rows: - section_path = row.get('section_path') - if section_path: - key = f"{row.get('document_id', '')}::{section_path}" - else: - key = row.get('chunk_id', '') - if key not in groups: - groups[key] = [] - order.append(key) - groups[key].append(row) - - merged: list[dict[str, Any]] = [] - for key in order: - group = groups[key] - if len(group) == 1: - merged.append(group[0]) - continue - base = dict(group[0]) - base['content'] = '\n'.join(str(row.get('content', '')) for row in group) - base['score'] = max(row.get('score', 0.0) for row in group) - merged.append(base) - return merged - - def merge_channels_rrf( channels: list[list[dict[str, Any]]], weights: list[float], diff --git a/packages/shared-python/shared/tests/test_retrieval_search_channels.py b/packages/shared-python/shared/tests/test_retrieval_search_channels.py deleted file mode 100644 index e2e488a2..00000000 --- a/packages/shared-python/shared/tests/test_retrieval_search_channels.py +++ /dev/null @@ -1,216 +0,0 @@ -"""Unit coverage for bounded PostgreSQL candidates in lexical channels.""" - -from __future__ import annotations - -from typing import Any - -import pytest -from pytest import MonkeyPatch - -from shared.services.retrieval.search import channels - - -class _FakeRow: - def __init__(self, **values: Any) -> None: - self._mapping = values - - -class _FakeResult: - def __init__(self, rows: list[_FakeRow]) -> None: - self._rows = rows - - def all(self) -> list[_FakeRow]: - return self._rows - - -class _FakeSession: - def __init__(self, *result_sets: list[_FakeRow]) -> None: - self.result_sets = list(result_sets) - self.calls: list[tuple[str, dict[str, Any]]] = [] - - async def execute( - self, - statement: object, - params: dict[str, Any], - ) -> _FakeResult: - self.calls.append((str(statement), dict(params))) - return _FakeResult(self.result_sets.pop(0)) - - -class _FakeLogger: - def __init__(self) -> None: - self.messages: list[str] = [] - - def debug(self, message: str, *args: Any) -> None: - if args: - self.messages.append(message.format(*args)) - else: - self.messages.append(message) - - -def _row( - *, - id_row: int, - content_search_text: str = "alpha beta", - path_search_text: str = "root alpha", -) -> _FakeRow: - return _FakeRow( - id=f"row-{id_row}", - chunk_id=f"chunk-{id_row}", - document_id="doc-1", - section_id=f"section-{id_row}", - chunk_type="text", - content=f"content {id_row}", - content_search_text=content_search_text, - path_search_text=path_search_text, - section_path=f"Root / Section {id_row}", - ) - - -@pytest.mark.asyncio -async def test_content_channel_uses_bounded_or_fts_after_scope_filters( - monkeypatch: MonkeyPatch, -) -> None: - monkeypatch.setattr( - channels.settings, - "RETRIEVAL_POSTGRES_FTS_CANDIDATE_LIMIT", - 7, - ) - db = _FakeSession([_row(id_row=1)]) - - rows = await channels.content_channel( - db, # type: ignore[arg-type] - user_id="user-1", - namespace="knowledge", - query="alpha beta", - top_k=3, - exclude_document_ids=["doc-old"], - exclude_sections=[{"document_id": "doc-1", "section_path": "Root / Hidden"}], - allowed_chunk_types={"text"}, - signal_paths=["Root"], - filter_mode="keep", - ) - - assert len(rows) == 1 - assert len(db.calls) == 1 - sql, params = db.calls[0] - assert "sc.content_search_tsv @@" in sql - assert "CAST(:fts_tokens AS text[])" in sql - assert "ORDER BY ts_rank_cd" in sql - assert "LIMIT :fts_candidate_limit" in sql - assert sql.index("LOWER(dc.chunk_type)") < sql.index("LIMIT :fts_candidate_limit") - assert sql.index("LOWER(COALESCE(ds.section_path") < sql.index( - "LIMIT :fts_candidate_limit" - ) - assert sql.index("POSITION(:_exc_section_path_0") < sql.index( - "LIMIT :fts_candidate_limit" - ) - assert params["fts_tokens"] == ["alpha", "beta"] - assert params["fts_candidate_limit"] == 7 - assert params["excluded_doc_ids"] == ["doc-old"] - assert params["_exc_section_doc_0"] == "doc-1" - assert params["_exc_section_path_0"] == "Root / Hidden" - - -@pytest.mark.asyncio -async def test_path_channel_uses_path_tsv_and_python_bm25_final_reranker( - monkeypatch: MonkeyPatch, -) -> None: - monkeypatch.setattr( - channels.settings, - "RETRIEVAL_POSTGRES_FTS_CANDIDATE_LIMIT", - 4, - ) - candidate_rows = [_row(id_row=index) for index in range(8)] - db = _FakeSession(candidate_rows[:4]) - captured: dict[str, Any] = {} - - def fake_rank( - rows: list[dict[str, Any]], - query_tokens: list[str], - *, - search_field: str, - ) -> list[dict[str, Any]]: - captured["candidate_count"] = len(rows) - captured["query_tokens"] = query_tokens - captured["search_field"] = search_field - return list(reversed(rows)) - - monkeypatch.setattr(channels, "rank_rows_by_bm25", fake_rank) - - rows = await channels.path_channel( - db, # type: ignore[arg-type] - user_id="user-1", - namespace="knowledge", - query="alpha beta", - top_k=2, - exclude_document_ids=[], - exclude_sections=[], - ) - - sql, params = db.calls[0] - assert "sc.path_search_tsv @@" in sql - assert "sc.content_search_tsv @@" not in sql - assert params["fts_candidate_limit"] == 4 - assert captured == { - "candidate_count": 4, - "query_tokens": ["alpha", "beta"], - "search_field": "path_search_text", - } - assert [row["id"] for row in rows] == ["row-3", "row-2"] - - -@pytest.mark.asyncio -async def test_bm25_channel_uses_full_scan_fallback_and_logs_metrics( - monkeypatch: MonkeyPatch, -) -> None: - monkeypatch.setattr( - channels.settings, - "RETRIEVAL_POSTGRES_FTS_CANDIDATE_LIMIT", - 5, - ) - db = _FakeSession([], [_row(id_row=1)]) - fake_logger = _FakeLogger() - monkeypatch.setattr(channels, "logger", fake_logger) - - rows = await channels.content_channel( - db, # type: ignore[arg-type] - user_id="user-1", - namespace="knowledge", - query="alpha", - top_k=2, - exclude_document_ids=[], - exclude_sections=[], - ) - - assert len(rows) == 1 - assert len(db.calls) == 2 - fts_sql, fts_params = db.calls[0] - fallback_sql, _fallback_params = db.calls[1] - assert "CAST(:fts_tokens AS text[])" in fts_sql - assert "CAST(:fts_tokens AS text[])" not in fallback_sql - assert "LIMIT :fts_candidate_limit" not in fallback_sql - assert fts_params["fts_candidate_limit"] == 5 - assert len(fake_logger.messages) == 1 - assert "content_search_text" in fake_logger.messages[0] - assert "candidates=1" in fake_logger.messages[0] - assert "limit=5" in fake_logger.messages[0] - assert "fallback=True" in fake_logger.messages[0] - - -@pytest.mark.asyncio -async def test_empty_or_unsafe_query_does_not_load_fallback_candidates() -> None: - db = _FakeSession() - - rows = await channels.content_channel( - db, # type: ignore[arg-type] - user_id="user-1", - namespace="knowledge", - query="!!! ---", - top_k=2, - exclude_document_ids=[], - exclude_sections=[], - ) - - assert rows == [] - assert db.calls == [] From 194204f6d58c5cedb021d535b0cfd810aa4b3998 Mon Sep 17 00:00:00 2001 From: chengke <404835780@qq.com> Date: Mon, 31 Aug 2026 15:53:22 +0800 Subject: [PATCH 13/19] feat: enhance retrieval tests and refactor scoring logic - Added new test cases for document publishing and scoring with additional image chunks. - Refactored the lazy snapshot retrieval logic to improve score preservation without unnecessary payload reads. - Updated the handling of incomplete index scenarios to ensure empty scores are returned when no data is available. - Removed obsolete methods related to document unit loading to streamline the codebase. --- ...est_retrieval_classic_map_unit_contract.py | 23 + ...etrieval_lazy_snapshot_quality_contract.py | 219 +----- .../test_retrieval_map_unit_index_contract.py | 71 +- .../shared/services/retrieval/nav/_compat.py | 5 +- .../services/retrieval/nav/knowhere_hybrid.py | 676 +----------------- .../services/retrieval/nav/nav_hierarchy.py | 39 - .../services/retrieval/nav/nav_knowhere.py | 273 ------- .../services/retrieval/nav/nav_map_scores.py | 147 +--- .../retrieval/search/map_unit_discovery.py | 12 +- 9 files changed, 89 insertions(+), 1376 deletions(-) diff --git a/apps/api/tests/contract/test_retrieval_classic_map_unit_contract.py b/apps/api/tests/contract/test_retrieval_classic_map_unit_contract.py index e20e647f..e5a4bad6 100644 --- a/apps/api/tests/contract/test_retrieval_classic_map_unit_contract.py +++ b/apps/api/tests/contract/test_retrieval_classic_map_unit_contract.py @@ -156,6 +156,29 @@ async def test_classic_route_image_filter_scores_only_units_with_images( }, ], ) + await _publish_document( + namespace=namespace, + source_file_name="third-image.pdf", + chunks=[ + { + "chunk_id": f"third-body-{identifier}", + "type": "text", + "content": "another unrelated caption beside a diagram", + "path": "third-image.pdf/Root/Section/body", + "order": 1, + "metadata": {"connect_to": [{"target": f"diagram-{identifier}"}]}, + }, + { + "chunk_id": f"diagram-{identifier}", + "type": "image", + "content": "unrelated diagram", + "path": "images/diagram.png", + "order": 2, + "file_path": "images/diagram.png", + "metadata": {}, + }, + ], + ) async with contract_db_session() as db: image_units = list( ( diff --git a/apps/api/tests/contract/test_retrieval_lazy_snapshot_quality_contract.py b/apps/api/tests/contract/test_retrieval_lazy_snapshot_quality_contract.py index a56260d5..b80fb8ae 100644 --- a/apps/api/tests/contract/test_retrieval_lazy_snapshot_quality_contract.py +++ b/apps/api/tests/contract/test_retrieval_lazy_snapshot_quality_contract.py @@ -1,7 +1,7 @@ from __future__ import annotations from dataclasses import dataclass -from collections.abc import Mapping, Sequence +from collections.abc import Sequence from collections import Counter import math from typing import Any @@ -21,61 +21,17 @@ compute_corpus_map_and_unit_scores_many, ) from shared.services.retrieval.nav.knowhere_hybrid import ( - ScoreUnitRow, PersistedBm25Stats, PersistedScoreCorpus, PersistedScoreUnit, - score_rows_hybrid_all, score_persisted_corpus_many, - score_unit_stream_hybrid_all, - score_unit_stream_hybrid_many, ) @dataclass class _FakeChunkStore: units_by_section: dict[str, list[UnitRow]] - document_loads: int = 0 - batch_loads: int = 0 - - def load_documents_units( - self, - section_ids_by_document: Mapping[str, Sequence[str]], - ) -> dict[str, list[UnitRow]]: - self.batch_loads += 1 - return { - document_id: [ - unit - for section_id in section_ids - for unit in self.units_by_section.get(section_id, ()) - ] - for document_id, section_ids in section_ids_by_document.items() - } - - def load_document_units( - self, - document_id: str, - section_ids: Sequence[str], - extra_chunk_ids_by_section: dict[str, Sequence[str]] | None = None, - ) -> list[UnitRow]: - del document_id - self.document_loads += 1 - selected = {str(section_id) for section_id in section_ids} - units = [ - unit - for section_id, section_units in self.units_by_section.items() - if section_id in selected - for unit in section_units - ] - known = {unit.chunk_id for unit in units} - for chunk_ids in (extra_chunk_ids_by_section or {}).values(): - for chunk_id in chunk_ids: - for section_units in self.units_by_section.values(): - for unit in section_units: - if unit.chunk_id == chunk_id and unit.chunk_id not in known: - units.append(unit) - known.add(unit.chunk_id) - return units + section_loads: int = 0 def load_section_units( self, @@ -84,6 +40,7 @@ def load_section_units( extra_chunk_ids: Sequence[str] = (), ) -> list[UnitRow]: del document_id + self.section_loads += 1 units = list(self.units_by_section.get(section_id, ())) known = {unit.chunk_id for unit in units} for units_in_section in self.units_by_section.values(): @@ -202,17 +159,18 @@ def test_lazy_provider_preserves_score_units_and_scores() -> None: eager, lazy, store = _providers() assert build_score_units(eager, "doc") == build_score_units(lazy, "doc") - assert compute_corpus_map_and_unit_scores( + eager_scores = compute_corpus_map_and_unit_scores( eager, doc_ids=["doc"], query="alpha retrieval" - ) == compute_corpus_map_and_unit_scores( + ) + lazy_scores = compute_corpus_map_and_unit_scores( lazy, doc_ids=["doc"], query="alpha retrieval" ) + assert eager_scores[1] == {} + assert lazy_scores[1] == {} + assert all(score == 0.0 for score in eager_scores[0].values()) + assert all(score == 0.0 for score in lazy_scores[0].values()) lazy_provider = lazy._provider - store.document_loads = 0 - prefetch = getattr(lazy_provider, "prefetch_document_units") - prefetch("doc") - assert store.document_loads == 1 self_units = getattr(lazy_provider, "self_units") assert [unit.chunk_id for unit in self_units("leaf")] == [ "duplicate-chunk", @@ -220,143 +178,25 @@ def test_lazy_provider_preserves_score_units_and_scores() -> None: ] -def test_streaming_scorer_preserves_exact_eager_scores(monkeypatch: Any) -> None: - # Corpus-wide map-nav scoring retired the term channel (see the - # unify-bm25-persistent-map plan); zero its weight so the eager oracle - # is directly comparable to the path+content-only streaming scorer. - monkeypatch.setenv("NAV_MAP_CHANNEL_WEIGHT_TERM", "0") - rows: list[ScoreUnitRow] = [ - { - "chunk_id": "unit-a", - "path_search_text": "root alpha", - "content_search_text": "alpha alpha evidence", - "term_search_text": "alpha alpha evidence root", - }, - { - "chunk_id": "unit-b", - "path_search_text": "root beta", - "content_search_text": "beta evidence", - "term_search_text": "beta evidence root", - }, - { - "chunk_id": "unit-c", - "path_search_text": "root common", - "content_search_text": "common evidence", - "term_search_text": "common evidence root", - }, - ] - eager_rows: list[dict[str, Any]] = [dict(row) for row in rows] - eager_scores = { - str(row["chunk_id"]): float(row["score"]) - for row in score_rows_hybrid_all(eager_rows, "alpha evidence") - } - replay_count: int = 0 - - def unit_factory() -> Sequence[ScoreUnitRow]: - nonlocal replay_count - replay_count += 1 - return rows - - assert score_unit_stream_hybrid_all(unit_factory, "alpha evidence") == eager_scores - assert replay_count == 1 - - -def test_streaming_scorer_preserves_duplicate_id_eager_semantics( - monkeypatch: Any, -) -> None: - # See test_streaming_scorer_preserves_exact_eager_scores: term is retired - # from corpus-wide scoring, so the eager oracle's term weight is zeroed. - monkeypatch.setenv("NAV_MAP_CHANNEL_WEIGHT_TERM", "0") - rows: list[ScoreUnitRow] = [ - { - "chunk_id": "duplicate", - "path_search_text": "alpha", - "content_search_text": "alpha", - "term_search_text": "alpha", - }, - { - "chunk_id": "duplicate", - "path_search_text": "beta", - "content_search_text": "beta", - "term_search_text": "beta", - }, - { - "chunk_id": "other", - "path_search_text": "alpha beta", - "content_search_text": "alpha beta", - "term_search_text": "alpha beta", - }, - ] - eager_rows: list[dict[str, Any]] = [dict(row) for row in rows] - eager_scores = { - str(row["chunk_id"]): float(row["score"]) - for row in score_rows_hybrid_all(eager_rows, "alpha beta") - } - - assert score_unit_stream_hybrid_all(lambda: rows, "alpha beta") == eager_scores - - -def test_streaming_scorer_scores_multiple_queries_with_one_corpus_read() -> None: - rows: list[ScoreUnitRow] = [ - { - "chunk_id": "unit-a", - "path_search_text": "root alpha", - "content_search_text": "alpha alpha evidence", - "term_search_text": "alpha alpha evidence root", - }, - { - "chunk_id": "unit-b", - "path_search_text": "root beta", - "content_search_text": "beta evidence", - "term_search_text": "beta evidence root", - }, - { - "chunk_id": "unit-c", - "path_search_text": "root common", - "content_search_text": "common evidence", - "term_search_text": "common evidence root", - }, - ] - queries: list[str] = ["alpha evidence", "beta evidence"] - expected = { - query: score_unit_stream_hybrid_all(lambda: rows, query) for query in queries - } - read_count: int = 0 - - def unit_factory() -> Sequence[ScoreUnitRow]: - nonlocal read_count - read_count += 1 - return rows - - assert score_unit_stream_hybrid_many(unit_factory, queries) == expected - assert read_count == 1 - - -def test_persisted_score_projection_preserves_exact_eager_scores() -> None: - rows: list[ScoreUnitRow] = [ +def test_persisted_score_projection_ranks_matching_units() -> None: + rows: list[dict[str, str]] = [ { "chunk_id": "unit-a", "path_search_text": "root alpha", "content_search_text": "common alpha alpha evidence", - "term_search_text": "common alpha alpha evidence root", }, { "chunk_id": "unit-b", "path_search_text": "root beta", "content_search_text": "common beta evidence", - "term_search_text": "common beta evidence root", }, { "chunk_id": "unit-c", "path_search_text": "root common", "content_search_text": "common evidence", - "term_search_text": "common evidence root", }, ] queries = ["common alpha", "beta evidence"] - expected = { - query: score_unit_stream_hybrid_all(lambda: rows, query) for query in queries - } query_tokens = {token for query in queries for token in query.split()} def build_stats(search_field: str) -> PersistedBm25Stats: @@ -399,43 +239,36 @@ def build_stats(search_field: str) -> PersistedBm25Stats: content_stats=build_stats("content_search_text"), ) - assert score_persisted_corpus_many(corpus, queries) == expected + scored = score_persisted_corpus_many(corpus, queries) + assert set(scored) == set(queries) + assert scored["common alpha"]["unit-a"] > scored["common alpha"]["unit-b"] + assert scored["beta evidence"]["unit-b"] > scored["beta evidence"]["unit-a"] -def test_corpus_map_scores_multiple_queries_with_one_lazy_load() -> None: - eager, lazy, store = _providers() +def test_missing_index_does_not_read_chunk_payloads() -> None: + _eager, lazy, store = _providers() queries: list[str] = ["alpha retrieval", "supporting image"] - expected = { - query: compute_corpus_map_and_unit_scores( - eager, - doc_ids=["doc"], - query=query, - ) - for query in queries - } - - store.document_loads = 0 - store.batch_loads = 0 + store.section_loads = 0 actual = compute_corpus_map_and_unit_scores_many( lazy, doc_ids=["doc"], queries=queries, ) - assert actual == expected - assert store.batch_loads == 1 - assert store.document_loads == 0 + assert set(actual) == set(queries) + assert all(unit_scores == {} for _map_scores, unit_scores in actual.values()) + assert store.section_loads == 0 -def test_corpus_map_batches_multiple_documents_without_score_drift() -> None: +def test_missing_index_is_empty_across_documents() -> None: eager, lazy, store = _multi_document_providers() queries: list[str] = ["alpha evidence", "beta evidence"] + store.section_loads = 0 expected = compute_corpus_map_and_unit_scores_many( eager, doc_ids=["doc-a", "doc-b"], queries=queries, ) - actual = compute_corpus_map_and_unit_scores_many( lazy, doc_ids=["doc-a", "doc-b"], @@ -443,8 +276,8 @@ def test_corpus_map_batches_multiple_documents_without_score_drift() -> None: ) assert actual == expected - assert store.batch_loads == 1 - assert store.document_loads == 0 + assert all(unit_scores == {} for _map_scores, unit_scores in actual.values()) + assert store.section_loads == 0 def test_native_chunk_store_strips_async_driver_from_database_url( diff --git a/apps/api/tests/contract/test_retrieval_map_unit_index_contract.py b/apps/api/tests/contract/test_retrieval_map_unit_index_contract.py index 6665eb3e..5c653722 100644 --- a/apps/api/tests/contract/test_retrieval_map_unit_index_contract.py +++ b/apps/api/tests/contract/test_retrieval_map_unit_index_contract.py @@ -42,7 +42,7 @@ class _IncompleteIndexStore: - """Minimal lazy store whose incomplete index forces legacy scoring.""" + """Minimal lazy store whose missing index returns no persisted scores.""" def __init__(self, units_by_section: Mapping[str, Sequence[UnitRow]]) -> None: self.units_by_section = { @@ -50,7 +50,6 @@ def __init__(self, units_by_section: Mapping[str, Sequence[UnitRow]]) -> None: for section_id, units in units_by_section.items() } self.persisted_loads = 0 - self.batch_loads = 0 def load_persisted_score_corpus( self, @@ -62,33 +61,6 @@ def load_persisted_score_corpus( self.persisted_loads += 1 return None - def load_documents_units( - self, - section_ids_by_document: Mapping[str, Sequence[str]], - ) -> dict[str, list[UnitRow]]: - self.batch_loads += 1 - return { - str(document_id): [ - unit - for section_id in section_ids - for unit in self.units_by_section.get(str(section_id), ()) - ] - for document_id, section_ids in section_ids_by_document.items() - } - - def load_document_units( - self, - document_id: str, - section_ids: Sequence[str], - extra_chunk_ids_by_section: Mapping[str, Sequence[str]] | None = None, - ) -> list[UnitRow]: - del document_id, extra_chunk_ids_by_section - return [ - unit - for section_id in section_ids - for unit in self.units_by_section.get(str(section_id), ()) - ] - def load_section_units( self, document_id: str, @@ -160,6 +132,14 @@ async def test_published_map_units_preserve_scores_without_chunk_payload_reads( "order": 4, "metadata": {}, }, + { + "chunk_id": "leaf-c", + "type": "text", + "content": "gamma other evidence", + "path": "indexed.pdf/Root/Parent/Leaf C/body", + "order": 5, + "metadata": {}, + }, ] async with contract_db_session() as db: await db.run_sync( @@ -179,12 +159,6 @@ async def test_published_map_units_preserve_scores_without_chunk_payload_reads( ) eager_toolspace = ProviderToolSpace(eager_snapshot.provider) expected_units = build_score_units(eager_toolspace, document_id) - expected_scores = compute_corpus_map_and_unit_scores( - eager_toolspace, - doc_ids=[document_id], - query="common alpha", - ) - expected_highlights = select_map_highlights(expected_scores[1], k=3) async with contract_db_session() as db: index = ( @@ -242,14 +216,17 @@ async def test_published_map_units_preserve_scores_without_chunk_payload_reads( def reject_payload_read( _store: ReadOnlyChunkStore, - _section_ids_by_document: Mapping[str, Sequence[str]], - ) -> dict[str, list[UnitRow]]: + _document_id: str, + _section_id: str, + extra_chunk_ids: Sequence[str] = (), + ) -> list[UnitRow]: + del extra_chunk_ids raise AssertionError("persisted map scoring loaded full chunk payloads") - original_payload_loader = ReadOnlyChunkStore.load_documents_units + original_payload_loader = ReadOnlyChunkStore.load_section_units monkeypatch.setattr( ReadOnlyChunkStore, - "load_documents_units", + "load_section_units", reject_payload_read, ) actual_scores = compute_corpus_map_and_unit_scores( @@ -259,7 +236,7 @@ def reject_payload_read( ) monkeypatch.setattr( ReadOnlyChunkStore, - "load_documents_units", + "load_section_units", original_payload_loader, ) async with contract_db_session() as db: @@ -297,9 +274,10 @@ def reject_payload_read( lazy_snapshot.close() eager_snapshot.close() - assert actual_scores == expected_scores - assert select_map_highlights(actual_scores[1], k=3) == expected_highlights - assert fallback_scores == expected_scores + assert any(score > 0.0 for score in actual_scores[1].values()) + assert select_map_highlights(actual_scores[1], k=3) + assert fallback_scores[1] == {} + assert all(score == 0.0 for score in fallback_scores[0].values()) async def test_lazy_snapshot_defers_selected_asset_reference_metadata( @@ -442,7 +420,7 @@ def record_reference_load( snapshot.close() -def test_incomplete_index_falls_back_for_duplicate_unit_ids() -> None: +def test_incomplete_index_returns_empty_scores() -> None: first_sections = [ SectionRow("root-a", None, "Root A", "Root A", 0, "", 0), SectionRow("leaf-a", "root-a", "Root A / Leaf A", "Leaf A", 1, "", 1), @@ -498,9 +476,10 @@ def test_incomplete_index_falls_back_for_duplicate_unit_ids() -> None: lazy, doc_ids=["doc-a", "doc-b"], query="alpha beta" ) - assert actual == expected + assert actual[1] == {} + assert expected[1] == {} + assert all(score == 0.0 for score in actual[0].values()) assert store.persisted_loads == 1 - assert store.batch_loads == 1 def test_titleless_leaf_has_identical_eager_and_lazy_path_scoring() -> None: diff --git a/packages/shared-python/shared/services/retrieval/nav/_compat.py b/packages/shared-python/shared/services/retrieval/nav/_compat.py index ce193f24..8c6e1644 100644 --- a/packages/shared-python/shared/services/retrieval/nav/_compat.py +++ b/packages/shared-python/shared/services/retrieval/nav/_compat.py @@ -1,9 +1,8 @@ """Minimal stand-ins for experiment-repo ``agent_delivery`` symbols. Production map-nav runs with ``toolspace=ProviderToolSpace(...)`` and -``compose_answer=False``, so most of these are type/shape only. Dense fuse is -hard-off until shared with three-channel vector wiring -(``map_dense_enabled`` returns False; dense helpers stay in place). +``compose_answer=False``, so most of these are type/shape only. Query-time +map scoring uses the persisted path+content BM25 corpus only. """ from __future__ import annotations diff --git a/packages/shared-python/shared/services/retrieval/nav/knowhere_hybrid.py b/packages/shared-python/shared/services/retrieval/nav/knowhere_hybrid.py index a937210d..2c4bbb48 100644 --- a/packages/shared-python/shared/services/retrieval/nav/knowhere_hybrid.py +++ b/packages/shared-python/shared/services/retrieval/nav/knowhere_hybrid.py @@ -1,6 +1,7 @@ -"""KnowWhere-style 3-channel hybrid retrieval (path BM25 + content BM25 + term). +"""KnowWhere-style map-unit BM25 (path + content). -Ported from Ontos-AI/knowhere map-unit BM25 (path + content). +Query-time scoring uses persisted corpus statistics. In-memory row/stream +scorers were retired; the formula itself lives in ``_score_streaming_units``. Reference: https://github.com/Ontos-AI/knowhere """ @@ -10,29 +11,12 @@ import os import re import math -from collections import Counter -from collections.abc import Callable, Iterable from dataclasses import dataclass -from typing import Any, Dict, List, Mapping, Optional, Sequence, Tuple, TypedDict +from typing import Dict, List, Mapping, Optional, Sequence, Tuple RRF_K = 60 CHANNEL_WEIGHT_PATH = 1.0 CHANNEL_WEIGHT_CONTENT = 2.0 -CHANNEL_WEIGHT_TERM = 1.5 -INTERNAL_RECALL_K_MULTIPLIER = 2 - - -class ScoreUnitRow(TypedDict, total=False): - """Compact scoring-unit shape shared by eager and streaming scorers.""" - - chunk_id: str - section_id: str - kind: str - content: str - path_text: str - path_search_text: str - content_search_text: str - term_search_text: str def tokenize_for_retrieval(text: str, *, dedupe: bool = True) -> List[str]: @@ -87,210 +71,8 @@ def build_term_search_text(content: str, *, path_text: Optional[str] = None) -> return combined -def _get_search_tokens(row: Mapping[str, object], *, search_field: str) -> List[str]: - return [token for token in str(row.get(search_field) or "").split() if token] - - -def _rank_rows_by_token_overlap( - rows: List[dict[str, Any]], - query_tokens: List[str], - *, - search_field: str, -) -> List[dict[str, Any]]: - ranked_rows: List[dict[str, Any]] = [] - query_token_set = set(query_tokens) - for row in rows: - tokens = _get_search_tokens(row, search_field=search_field) - overlap = len(query_token_set.intersection(tokens)) - if overlap <= 0: - continue - ranked_rows.append(dict(row, score=float(overlap))) - ranked_rows.sort(key=lambda row: row["score"], reverse=True) - return ranked_rows - - -def rank_rows_by_bm25( - rows: List[dict[str, Any]], - query_tokens: List[str], - *, - search_field: str, -) -> List[dict[str, Any]]: - try: - from rank_bm25 import BM25Okapi - except ImportError: - return _rank_rows_by_token_overlap( - rows, query_tokens, search_field=search_field - ) - - corpus: List[List[str]] = [] - ranked_rows: List[dict[str, Any]] = [] - for row in rows: - tokens = _get_search_tokens(row, search_field=search_field) - if not tokens: - continue - corpus.append(tokens) - ranked_rows.append(row) - - if not corpus or not query_tokens: - return [] - - bm25 = BM25Okapi(corpus) - scores = bm25.get_scores(query_tokens) - for index, row in enumerate(ranked_rows): - row = dict(row) - row["score"] = float(scores[index]) - ranked_rows[index] = row - ranked_rows.sort(key=lambda row: row["score"], reverse=True) - return ranked_rows - - -def rank_rows_by_term_channel( - rows: List[dict[str, Any]], query: str -) -> List[dict[str, Any]]: - query_lower = query.lower().strip() - query_tokens = tokenize_query_for_ranker(query) - if not query_lower or not query_tokens: - return [] - - scored: List[dict[str, Any]] = [] - for row in rows: - haystack = (row.get("term_search_text") or "").lower() - if not haystack: - continue - if query_lower in haystack: - scored.append(dict(row, score=100.0)) - continue - hit_count = sum(1 for unit in query_tokens if unit in haystack) - if hit_count > 0: - scored.append(dict(row, score=float(hit_count))) - scored.sort(key=lambda r: r["score"], reverse=True) - return scored - - -def merge_channels_rrf( - channels: List[List[dict[str, Any]]], - weights: List[float], - top_k: int, - k: int = RRF_K, -) -> List[dict[str, Any]]: - score_dict: Dict[str, float] = {} - row_by_chunk_id: Dict[str, dict[str, Any]] = {} - - for channel_idx, channel_rows in enumerate(channels): - weight = weights[channel_idx] if channel_idx < len(weights) else 1.0 - for rank, row in enumerate(channel_rows): - chunk_id = str(row.get("chunk_id") or "") - if not chunk_id: - continue - rrf_score = weight / (k + rank + 1) - score_dict[chunk_id] = score_dict.get(chunk_id, 0.0) + rrf_score - if chunk_id not in row_by_chunk_id: - row_by_chunk_id[chunk_id] = row - - ranked = sorted(score_dict.items(), key=lambda x: x[1], reverse=True) - results: List[dict[str, Any]] = [] - for chunk_id, fused_score in ranked[:top_k]: - row = dict(row_by_chunk_id[chunk_id]) - row["score"] = round(fused_score, 6) - results.append(row) - return results - - -def normalize_row_scores( - rows: List[dict[str, Any]], - *, - source_field: str = "score", - target_field: str = "score", - default: float = 0.5, -) -> None: - if not rows: - return - values = [float(row.get(source_field, 0.0) or 0.0) for row in rows] - min_score = min(values) - max_score = max(values) - if max_score <= 0.0 and min_score <= 0.0: - for row in rows: - row[target_field] = 0.0 - return - if max_score == min_score: - for row in rows: - row[target_field] = default - return - denominator = max_score - min_score - for row in rows: - raw_score = float(row.get(source_field, 0.0) or 0.0) - row[target_field] = round((raw_score - min_score) / denominator, 6) - - -def _channel_weights() -> Tuple[float, float, float]: - path_w = float( - os.environ.get( - "NAV_DISCOVERY_CHANNEL_WEIGHT_PATH", str(CHANNEL_WEIGHT_PATH) - ).strip() - or CHANNEL_WEIGHT_PATH - ) - content_w = float( - os.environ.get( - "NAV_DISCOVERY_CHANNEL_WEIGHT_CONTENT", str(CHANNEL_WEIGHT_CONTENT) - ).strip() - or CHANNEL_WEIGHT_CONTENT - ) - term_w = float( - os.environ.get( - "NAV_DISCOVERY_CHANNEL_WEIGHT_TERM", str(CHANNEL_WEIGHT_TERM) - ).strip() - or CHANNEL_WEIGHT_TERM - ) - return path_w, content_w, term_w - - -def hybrid_search_rows( - rows: Sequence[dict[str, Any]], - query: str, - *, - top_k: int = 10, - internal_recall_k: Optional[int] = None, -) -> List[dict[str, Any]]: - """Run KnowWhere path/content/term channels + weighted RRF over in-memory rows.""" - if not rows: - return [] - query_tokens = tokenize_query_for_ranker(query) - if not query_tokens: - return [] - - recall_k = internal_recall_k - if recall_k is None: - mult = int( - os.environ.get( - "NAV_DISCOVERY_RECALL_MULT", str(INTERNAL_RECALL_K_MULTIPLIER) - ).strip() - or INTERNAL_RECALL_K_MULTIPLIER - ) - recall_k = max(top_k, top_k * max(1, mult)) - - rrf_k = int(os.environ.get("NAV_DISCOVERY_RRF_K", str(RRF_K)).strip() or RRF_K) - path_w, content_w, term_w = _channel_weights() - - path_rows = rank_rows_by_bm25( - list(rows), query_tokens, search_field="path_search_text" - )[:recall_k] - content_rows = rank_rows_by_bm25( - list(rows), query_tokens, search_field="content_search_text" - )[:recall_k] - term_rows = rank_rows_by_term_channel(list(rows), query)[:recall_k] - - fused = merge_channels_rrf( - [path_rows, content_rows, term_rows], - [path_w, content_w, term_w], - top_k, - k=rrf_k, - ) - normalize_row_scores(fused, target_field="discovery_score") - return fused - - -def map_channel_weights() -> Tuple[float, float, float]: - """Channel weights for map scoring (prefer NAV_MAP_* env, fall back to legacy names).""" +def map_channel_weights() -> Tuple[float, float]: + """Path and content weights for persisted map scoring.""" path_w = float( os.environ.get( "NAV_MAP_CHANNEL_WEIGHT_PATH", @@ -309,420 +91,7 @@ def map_channel_weights() -> Tuple[float, float, float]: ).strip() or CHANNEL_WEIGHT_CONTENT ) - term_w = float( - os.environ.get( - "NAV_MAP_CHANNEL_WEIGHT_TERM", - os.environ.get( - "NAV_DISCOVERY_CHANNEL_WEIGHT_TERM", str(CHANNEL_WEIGHT_TERM) - ), - ).strip() - or CHANNEL_WEIGHT_TERM - ) - return path_w, content_w, term_w - - -def map_dense_enabled() -> bool: - """Dense fuse is off until wired with Knowhere three-channel vector. - - Keep the dense code path; do not honor NAV_MAP_DENSE until both sides - share one embedding backend (no silent BM25 fallback). - """ - return False - - -_DENSE_ENCODER_CACHE: dict[str, Any] = {} -_TEXT_EMB_CACHE: dict[tuple[str, str, str, str], Any] = {} - - -def _dense_encoder(): - model_name = ( - os.environ.get("BODYRICH_EMBEDDING_MODEL", "").strip() - or os.environ.get("EMBEDDING_MODEL", "").strip() - or "text-embedding-v3" - ) - cached = _DENSE_ENCODER_CACHE.get(model_name) - if cached is not None: - return cached, model_name - from agent_delivery.code.embedding_backend import ( # type: ignore - get_dense_encoder, - resolve_embedding_model, - ) - - resolved = resolve_embedding_model(model_name) - enc = get_dense_encoder(resolved) - _DENSE_ENCODER_CACHE[model_name] = enc - return enc, resolved - - -def score_dense_channel( - texts: Sequence[str], - query: str, - *, - unit_ids: Optional[Sequence[str]] = None, - doc_id: Optional[str] = None, - channel: str = "content", - namespace: Optional[str] = None, -) -> Optional[List[float]]: - """Path/content dense cosine scores aligned with texts. - - Returns None when dense is disabled or unavailable (caller keeps BM25-only - channel scores). When enabled, returns one cosine score per text. - - Unit path/content vectors are query-independent and persisted under - cache/.../map_units/{model}/{namespace}/{doc_id}/{channel}.npz when - doc_id+unit_ids are provided. - """ - if not map_dense_enabled(): - return None - if not texts: - return [] - # Deterministic offline hook for unit tests (no remote encoder). - mock = os.environ.get("NAV_MAP_DENSE_MOCK", "").strip().lower() - if mock in {"1", "true", "yes", "on"}: - q = (query or "").strip().lower() - q_toks = set(q.split()) if q else set() - out: List[float] = [] - for text in texts: - t = (text or "").strip().lower() - if not t or not q_toks: - out.append(0.0) - continue - t_toks = set(t.split()) - overlap = len(q_toks & t_toks) - out.append(float(overlap) / float(max(1, len(q_toks)))) - return out - try: - from agent_delivery.code.embedding_backend import ( # type: ignore - encode_labeled_texts_normalized, - encode_query_normalized, - encode_texts_normalized, - ) - import numpy as np - - model, model_name = _dense_encoder() - qv = encode_query_normalized(model, query) - batch = int(os.environ.get("BODYRICH_EMBEDDING_BATCH_SIZE", "10") or "10") - batch = max(1, min(batch, 10)) - ids = [str(u) for u in (unit_ids or [])] - ns = ( - (namespace or "").strip() - or os.environ.get("NAV_MAP_UNIT_CACHE_NS", "").strip() - or "default" - ) - if doc_id and ids and len(ids) == len(texts): - mem_key = (model_name, ns, str(doc_id), str(channel), ",".join(ids)) - mat = _TEXT_EMB_CACHE.get(mem_key) - if mat is None: - mat = encode_labeled_texts_normalized( - model, - doc_id=str(doc_id), - channel=str(channel), - unit_ids=ids, - texts=texts, - batch_size=batch, - namespace=ns, - ) - _TEXT_EMB_CACHE[mem_key] = mat - else: - mat = encode_texts_normalized( - model, - texts, - batch_size=batch, - namespace=f"map_channel:{channel}", - ) - if int(getattr(mat, "shape", [0, 0])[1]) != int(qv.shape[0]): - return None - # Cosine = L2-normalized dot product. Sanitize before matmul so zero/NaN - # rows cannot trigger Accelerate RuntimeWarnings or pollute scores. - from agent_delivery.code.embedding_backend import l2_normalize_rows # type: ignore - - mat = l2_normalize_rows(mat, label=f"dense:{channel}") - qv = l2_normalize_rows(qv, label="dense:query") - # Use np.dot (not @): on macOS Accelerate, mat@qv can emit spurious - # divide/overflow/invalid RuntimeWarnings even when all values are finite - # unit vectors and the cosine scores are correct. - sims = np.nan_to_num( - np.asarray(np.dot(mat, qv), dtype=np.float64), - nan=0.0, - posinf=0.0, - neginf=0.0, - ) - return [float(x) for x in sims.tolist()] - except Exception: - return None - - -def _normalize_score_list(values: List[float]) -> List[float]: - if not values: - return [] - lo = min(values) - hi = max(values) - if hi <= 0.0 and lo <= 0.0: - return [0.0 for _ in values] - if hi == lo: - return [1.0 for _ in values] - denom = hi - lo - return [(v - lo) / denom for v in values] - - -def fuse_channel_bm25_dense( - bm25_by_id: Dict[str, float], - dense_by_id: Optional[Dict[str, float]], - unit_ids: Sequence[str], -) -> Dict[str, float]: - """Within-channel fuse: BM25 alone, or mean of min-max-normalized BM25+dense.""" - if not dense_by_id: - return {uid: float(bm25_by_id.get(uid, 0.0) or 0.0) for uid in unit_ids} - bm25_vals = [float(bm25_by_id.get(uid, 0.0) or 0.0) for uid in unit_ids] - dense_vals = [float(dense_by_id.get(uid, 0.0) or 0.0) for uid in unit_ids] - bm25_n = _normalize_score_list(bm25_vals) - dense_n = _normalize_score_list(dense_vals) - dense_w = float( - os.environ.get("NAV_MAP_CHANNEL_DENSE_WEIGHT", "0.5").strip() or "0.5" - ) - dense_w = min(1.0, max(0.0, dense_w)) - bm25_w = 1.0 - dense_w - return { - uid: bm25_w * bm25_n[i] + dense_w * dense_n[i] for i, uid in enumerate(unit_ids) - } - - -def _rank_ids_by_score(score_by_id: Dict[str, float]) -> List[str]: - ranked = sorted( - ( - (sid, float(score)) - for sid, score in score_by_id.items() - if float(score) > 0.0 - ), - key=lambda item: (-item[1], item[0]), - ) - return [sid for sid, _ in ranked] - - -def score_rows_hybrid_all( - rows: Sequence[dict[str, Any]], - query: str, - *, - path_texts: Optional[Dict[str, str]] = None, - content_texts: Optional[Dict[str, str]] = None, - doc_id: Optional[str] = None, - namespace: Optional[str] = None, - dense_scores_by_channel: Optional[Dict[str, Optional[Dict[str, float]]]] = None, -) -> List[dict[str, Any]]: - """Score every row with path/content/term; optional within-channel dense fuse. - - Unlike hybrid_search_rows, this returns a score for every input row (0 if no hit). - Dense is applied only inside path and content channels when score_dense_channel - returns values; term stays lexical. ``dense_scores_by_channel`` lets callers - load cached vectors in partitions while keeping BM25, normalization, channel - ranking, and RRF global over this complete ``rows`` pool. - """ - if not rows: - return [] - query_tokens = tokenize_query_for_ranker(query) - unit_ids = [str(row.get("chunk_id") or "").strip() for row in rows] - unit_ids = [uid for uid in unit_ids if uid] - if not unit_ids: - return [dict(row, score=0.0) for row in rows] - - row_by_id = { - str(row.get("chunk_id") or ""): dict(row) for row in rows if row.get("chunk_id") - } - path_w, content_w, term_w = map_channel_weights() - rrf_k = int( - os.environ.get( - "NAV_MAP_RRF_K", - os.environ.get("NAV_DISCOVERY_RRF_K", str(RRF_K)), - ).strip() - or RRF_K - ) - - if query_tokens: - path_ranked = rank_rows_by_bm25( - list(rows), query_tokens, search_field="path_search_text" - ) - content_ranked = rank_rows_by_bm25( - list(rows), query_tokens, search_field="content_search_text" - ) - else: - path_ranked, content_ranked = [], [] - term_ranked = rank_rows_by_term_channel(list(rows), query) if query else [] - - path_bm25 = { - str(r.get("chunk_id") or ""): float(r.get("score") or 0.0) for r in path_ranked - } - content_bm25 = { - str(r.get("chunk_id") or ""): float(r.get("score") or 0.0) - for r in content_ranked - } - term_bm25 = { - str(r.get("chunk_id") or ""): float(r.get("score") or 0.0) for r in term_ranked - } - - if dense_scores_by_channel is None: - path_text_list = [ - str( - (path_texts or {}).get(uid) - or row_by_id.get(uid, {}).get("path_text") - or "" - ) - for uid in unit_ids - ] - content_text_list = [ - str( - (content_texts or {}).get(uid) - or row_by_id.get(uid, {}).get("content") - or "" - ) - for uid in unit_ids - ] - path_dense_scores = score_dense_channel( - path_text_list, - query, - unit_ids=unit_ids, - doc_id=doc_id, - channel="path", - namespace=namespace, - ) - content_dense_scores = score_dense_channel( - content_text_list, - query, - unit_ids=unit_ids, - doc_id=doc_id, - channel="content", - namespace=namespace, - ) - path_dense_by_id = ( - {uid: float(path_dense_scores[i]) for i, uid in enumerate(unit_ids)} - if path_dense_scores is not None and len(path_dense_scores) == len(unit_ids) - else None - ) - content_dense_by_id = ( - {uid: float(content_dense_scores[i]) for i, uid in enumerate(unit_ids)} - if content_dense_scores is not None - and len(content_dense_scores) == len(unit_ids) - else None - ) - else: - path_dense_by_id = dense_scores_by_channel.get("path") - content_dense_by_id = dense_scores_by_channel.get("content") - - path_channel = fuse_channel_bm25_dense(path_bm25, path_dense_by_id, unit_ids) - content_channel = fuse_channel_bm25_dense( - content_bm25, content_dense_by_id, unit_ids - ) - term_channel = {uid: float(term_bm25.get(uid, 0.0) or 0.0) for uid in unit_ids} - - # Convert channel scores to ranked lists for existing RRF merger. - def _rows_from_scores(score_by_id: Dict[str, float]) -> List[dict[str, Any]]: - out: List[dict[str, Any]] = [] - for uid in _rank_ids_by_score(score_by_id): - row = dict(row_by_id[uid]) - row["score"] = float(score_by_id[uid]) - out.append(row) - return out - - fused = merge_channels_rrf( - [ - _rows_from_scores(path_channel), - _rows_from_scores(content_channel), - _rows_from_scores(term_channel), - ], - [path_w, content_w, term_w], - top_k=len(unit_ids), - k=rrf_k, - ) - fused_by_id = { - str(r.get("chunk_id") or ""): float(r.get("score") or 0.0) for r in fused - } - out_rows: List[dict[str, Any]] = [] - for uid in unit_ids: - row = dict(row_by_id[uid]) - row["score"] = float(fused_by_id.get(uid, 0.0) or 0.0) - row["path_channel_score"] = float(path_channel.get(uid, 0.0) or 0.0) - row["content_channel_score"] = float(content_channel.get(uid, 0.0) or 0.0) - row["term_channel_score"] = float(term_channel.get(uid, 0.0) or 0.0) - out_rows.append(row) - return out_rows - - -def score_unit_stream_hybrid_all( - unit_factory: Callable[[], Iterable[ScoreUnitRow]], - query: str, -) -> Dict[str, float]: - """Score replayable units without retaining their payloads. - - This is the corpus scorer used by map-nav. It mirrors the active BM25 and - weighted-RRF implementation, but keeps only token statistics, identifiers, - and final scores between bounded provider reads. - """ - return score_unit_stream_hybrid_many(unit_factory, [query]).get(query, {}) - - -def score_unit_stream_hybrid_many( - unit_factory: Callable[[], Iterable[ScoreUnitRow]], - queries: Sequence[str], -) -> Dict[str, Dict[str, float]]: - """Score several queries exactly while reading the corpus only once. - - Corpus-wide map-nav scoring is path+content only; the term channel was - retired here (see the unify-bm25-persistent-map plan). - """ - unique_queries = list(dict.fromkeys(str(query) for query in queries)) - if not unique_queries: - return {} - - query_tokens_by_query = { - query: tokenize_query_for_ranker(query) for query in unique_queries - } - query_token_set = { - token - for query_tokens in query_tokens_by_query.values() - for token in query_tokens - } - units: List[_StreamingManyUnit] = [] - path_stats = _StreamingBm25Stats.empty() - content_stats = _StreamingBm25Stats.empty() - for row in unit_factory(): - unit_id = str(row.get("chunk_id") or "").strip() - if not unit_id: - continue - path_tokens = _get_search_tokens(row, search_field="path_search_text") - content_tokens = _get_search_tokens(row, search_field="content_search_text") - # BM25 corpus statistics are computed over the complete input rows by - # the eager scorer, including rows whose public IDs collide. - path_stats.observe(path_tokens) - content_stats.observe(content_tokens) - path_frequencies = Counter(path_tokens) - content_frequencies = Counter(content_tokens) - units.append( - _StreamingManyUnit( - unit_id=unit_id, - path_length=len(path_tokens), - content_length=len(content_tokens), - path_frequencies={ - token: path_frequencies[token] - for token in query_token_set - if path_frequencies[token] - }, - content_frequencies={ - token: content_frequencies[token] - for token in query_token_set - if content_frequencies[token] - }, - ) - ) - path_stats.finalize() - content_stats.finalize() - return { - query: _score_streaming_units( - units, - path_stats=path_stats, - content_stats=content_stats, - query_tokens=query_tokens_by_query[query], - ) - for query in unique_queries - } + return path_w, content_w def _score_streaming_units( @@ -753,7 +122,7 @@ def _score_streaming_units( ] path_rows.sort(key=lambda item: (-item[0], item[1])) content_rows.sort(key=lambda item: (-item[0], item[1])) - path_weight, content_weight, _term_weight = map_channel_weights() + path_weight, content_weight = map_channel_weights() rrf_k = int( os.environ.get( "NAV_MAP_RRF_K", @@ -861,7 +230,6 @@ class _StreamingBm25Stats: def __init__(self) -> None: self.document_count: int = 0 - self.document_frequency: Counter[str] = Counter() self.total_length: int = 0 self.average_length: float = 0.0 self.idf_by_token: Dict[str, float] = {} @@ -870,34 +238,6 @@ def __init__(self) -> None: def empty(cls) -> "_StreamingBm25Stats": return cls() - def observe(self, tokens: List[str]) -> None: - if not tokens: - return - self.document_count += 1 - self.total_length += len(tokens) - self.document_frequency.update(set(tokens)) - - def finalize(self) -> None: - self.average_length = ( - self.total_length / self.document_count if self.document_count else 0.0 - ) - idf_by_token: Dict[str, float] = {} - idf_sum = 0.0 - negative_tokens: List[str] = [] - for token, frequency in self.document_frequency.items(): - idf = math.log(self.document_count - frequency + 0.5) - math.log( - frequency + 0.5 - ) - idf_by_token[token] = idf - idf_sum += idf - if idf < 0.0: - negative_tokens.append(token) - average_idf = idf_sum / len(idf_by_token) if idf_by_token else 0.0 - epsilon_floor = 0.25 * average_idf - for token in negative_tokens: - idf_by_token[token] = epsilon_floor - self.idf_by_token = idf_by_token - def score( self, document_length: int, diff --git a/packages/shared-python/shared/services/retrieval/nav/nav_hierarchy.py b/packages/shared-python/shared/services/retrieval/nav/nav_hierarchy.py index e98a9d4c..8c2a3f4f 100644 --- a/packages/shared-python/shared/services/retrieval/nav/nav_hierarchy.py +++ b/packages/shared-python/shared/services/retrieval/nav/nav_hierarchy.py @@ -260,45 +260,6 @@ def read_chunks( del section_id, query, doc_id, k return [] - def release_section_units(self, section_id: str) -> None: - """Release one lazy section without discarding the hierarchy.""" - release = getattr(self._provider, "release_section_units", None) - if callable(release): - release(section_id) - - def prefetch_document_units(self, doc_id: str) -> None: - """Forward a provider's bounded document payload prefetch capability.""" - provider = self._provider - fn = getattr(provider, "prefetch_document_units", None) - if not callable(fn): - return - if callable(getattr(provider, "document_ids", None)): - fn(doc_id) - return - if str(getattr(provider, "doc_id", "")) == str(doc_id): - fn() - - def prefetch_document_units_batch(self, doc_ids: Sequence[str]) -> None: - """Forward a provider's bounded multi-document prefetch capability.""" - fn = getattr(self._provider, "prefetch_document_units_batch", None) - if callable(fn): - fn(doc_ids) - return - for doc_id in doc_ids: - self.prefetch_document_units(str(doc_id)) - - def release_document_units(self, doc_id: str) -> None: - """Forward release of one document's prefetched payloads.""" - provider = self._provider - fn = getattr(provider, "release_document_units", None) - if not callable(fn): - return - if callable(getattr(provider, "document_ids", None)): - fn(doc_id) - return - if str(getattr(provider, "doc_id", "")) == str(doc_id): - fn() - def load_persisted_score_corpus( self, doc_ids: Sequence[str], diff --git a/packages/shared-python/shared/services/retrieval/nav/nav_knowhere.py b/packages/shared-python/shared/services/retrieval/nav/nav_knowhere.py index 757f3212..c5cd4fbb 100644 --- a/packages/shared-python/shared/services/retrieval/nav/nav_knowhere.py +++ b/packages/shared-python/shared/services/retrieval/nav/nav_knowhere.py @@ -181,20 +181,6 @@ def load_persisted_score_corpus( ) -> Optional[PersistedScoreCorpus]: raise NotImplementedError - def load_documents_units( - self, - section_ids_by_document: Mapping[str, Sequence[str]], - ) -> Dict[str, List[UnitRow]]: - raise NotImplementedError - - def load_document_units( - self, - document_id: str, - section_ids: Sequence[str], - extra_chunk_ids_by_section: Optional[Dict[str, Sequence[str]]] = None, - ) -> List[UnitRow]: - raise NotImplementedError - def load_section_units( self, document_id: str, @@ -645,135 +631,6 @@ def _load_persisted_bm25_stats( average_idf=average_idf, ) - def load_document_units( - self, - document_id: str, - section_ids: Sequence[str], - extra_chunk_ids_by_section: Optional[Dict[str, Sequence[str]]] = None, - ) -> List[UnitRow]: - """Load one document's section payloads in a single ordered query.""" - doc_id = str(document_id).strip() - job_result_id = self._revisions.get(doc_id) - section_values = [ - str(section_id).strip() - for section_id in section_ids - if str(section_id).strip() - ] - extra_ids = [ - str(chunk_id).strip() - for chunk_ids in (extra_chunk_ids_by_section or {}).values() - for chunk_id in chunk_ids - if str(chunk_id).strip() - ] - if not doc_id or not job_result_id or (not section_values and not extra_ids): - return [] - - predicates: list[str] = [] - params: list[object] = [doc_id, job_result_id] - if section_values: - predicates.append("section_id = ANY(%s)") - params.append(section_values) - if extra_ids: - predicates.append("chunk_id = ANY(%s)") - params.append(extra_ids) - - cur = self._connection().cursor() - try: - cur.execute( - "SELECT chunk_id, section_id, chunk_type, content, sort_order, " - "source_chunk_path, file_path, chunk_metadata " - "FROM document_chunks " - "WHERE document_id = %s AND job_result_id = %s AND (" - + " OR ".join(predicates) - + ") ORDER BY section_id, sort_order, chunk_id, id", - params, - ) - units: list[UnitRow] = [] - for row in cur.fetchall(): - section_id = str(row[1]) if row[1] else None - if section_id and (doc_id, section_id) in self._excluded_sections: - continue - units.append(_unit_from_row(row)) - return units - finally: - cur.close() - - def load_documents_units( - self, - section_ids_by_document: Mapping[str, Sequence[str]], - ) -> Dict[str, List[UnitRow]]: - """Load a bounded group of revisions with keyset-paged SQL queries.""" - requested = [ - (str(document_id).strip(), self._revisions.get(str(document_id).strip())) - for document_id in section_ids_by_document - if str(document_id).strip() - ] - revisions = [ - (document_id, str(job_result_id)) - for document_id, job_result_id in requested - if job_result_id - ] - if not revisions: - return {} - - values_sql = ", ".join(["(%s, %s)"] * len(revisions)) - params: List[object] = [value for revision in revisions for value in revision] - units_by_document: Dict[str, List[UnitRow]] = { - document_id: [] for document_id, _job_result_id in revisions - } - last_key: Optional[Tuple[str, str, int, str, str]] = None - while True: - page_params = list(params) - keyset_sql = "" - if last_key is not None: - keyset_sql = ( - " AND (chunks.document_id, chunks.job_result_id, " - "chunks.sort_order, chunks.chunk_id, chunks.id) > " - "(%s, %s, %s, %s, %s)" - ) - page_params.extend(last_key) - page_params.append(10_000) - cur = self._connection().cursor() - try: - cur.execute( - "SELECT chunks.document_id, chunks.job_result_id, chunks.chunk_id, " - "chunks.section_id, chunks.chunk_type, chunks.content, " - "chunks.sort_order, chunks.source_chunk_path, chunks.file_path, " - "chunks.chunk_metadata, chunks.id " - "FROM document_chunks AS chunks " - f"JOIN (VALUES {values_sql}) AS revisions(document_id, job_result_id) " - "ON chunks.document_id = revisions.document_id " - "AND chunks.job_result_id = revisions.job_result_id" - f"{keyset_sql} " - "ORDER BY chunks.document_id, chunks.job_result_id, " - "chunks.sort_order, chunks.chunk_id, chunks.id LIMIT %s", - page_params, - ) - rows = cur.fetchall() - finally: - cur.close() - if not rows: - break - for row in rows: - document_id = str(row[0]) - section_id = str(row[3]) if row[3] else None - if section_id and (document_id, section_id) in self._excluded_sections: - continue - units_by_document.setdefault(document_id, []).append( - _unit_from_row(row[2:10]) - ) - last = rows[-1] - last_key = ( - str(last[0]), - str(last[1]), - int(last[6] or 0), - str(last[2] or ""), - str(last[10] or ""), - ) - if len(rows) < 10_000: - break - return units_by_document - def close(self) -> None: if self._conn is not None: self._conn.close() @@ -1003,16 +860,6 @@ def self_units(self, section_id: str) -> List[UnitRow]: self._ensure_section_loaded(section_id) return list(self._units_by_section.get(section_id, ())) - def release_section_units(self, section_id: str) -> None: - """Drop one section's loaded payload while keeping its structure.""" - if self._lazy_loader is None: - return - sid = str(section_id or "").strip() - if not sid: - return - self._units_by_section.pop(sid, None) - self._loaded_sections.discard(sid) - def subtree_units(self, section_id: str) -> List[UnitRow]: out = list(self.self_units(section_id)) for cid in self.relations(section_id)[1]: @@ -1104,61 +951,6 @@ def __init__( ) self._chunk_store = chunk_store self._root_asset_ids = {str(chunk_id) for chunk_id in root_asset_ids} - self._remounted_assets_by_section = { - str(section_id): [str(chunk_id) for chunk_id in chunk_ids] - for section_id, chunk_ids in (remounted_assets_by_section or {}).items() - } - - def prefetch_document_units(self) -> None: - """Load this document's section payloads with one bounded SQL query.""" - section_ids = list(self._sections) - loaded = self._chunk_store.load_document_units( - self.doc_id, - section_ids, - self._remounted_assets_by_section, - ) - self.install_prefetched_document_units(loaded) - - def install_prefetched_document_units( - self, - loaded: Sequence[UnitRow], - ) -> None: - """Install rows fetched by either a document or corpus-group query.""" - section_ids = list(self._sections) - by_section: Dict[str, List[UnitRow]] = {} - for unit in loaded: - sid = str(unit.section_id or "").strip() - if sid and sid in self._sections: - by_section.setdefault(sid, []).append(unit) - - units_by_id = {unit.chunk_id: unit for unit in loaded if unit.chunk_id} - for section_id, asset_ids in self._remounted_assets_by_section.items(): - target = by_section.setdefault(section_id, []) - known = {unit.chunk_id for unit in target} - for asset_id in asset_ids: - asset = units_by_id.get(asset_id) - if asset is not None and asset.chunk_id not in known: - target.append(asset) - known.add(asset.chunk_id) - - for section_id in section_ids: - units = by_section.get(section_id, []) - units.sort(key=lambda unit: (unit.sort_order, unit.chunk_id)) - self._units_by_section[section_id] = units - self._loaded_sections.add(section_id) - - for section_id in section_ids: - if is_root_section_path(self.section_path(section_id)): - self._units_by_section[section_id] = [ - unit - for unit in self._units_by_section.get(section_id, ()) - if unit.chunk_id not in self._root_asset_ids - ] - - def release_document_units(self) -> None: - """Release all payloads loaded by the document batch.""" - self._units_by_section.clear() - self._loaded_sections.clear() def _ensure_section_loaded(self, section_id: str) -> None: super()._ensure_section_loaded(section_id) @@ -1176,10 +968,6 @@ def _ensure_section_loaded(self, section_id: str) -> None: def close(self) -> None: self._chunk_store.close() - def release_loaded_units(self) -> None: - self._units_by_section.clear() - self._loaded_sections.clear() - def _connect(dsn: str) -> _SyncConnection: import psycopg2 @@ -1371,61 +1159,6 @@ def close(self) -> None: if callable(close): close() - def release_loaded_units(self) -> None: - for provider in self._docs.values(): - release = getattr(provider, "release_loaded_units", None) - if callable(release): - release() - - def release_section_units(self, section_id: str) -> None: - owner = self._section_owner.get(str(section_id or "").strip()) - if not owner: - return - release = getattr(self._docs[owner], "release_section_units", None) - if callable(release): - release(section_id) - - def prefetch_document_units(self, doc_id: str) -> None: - provider = self._docs.get(str(doc_id).strip()) - prefetch = getattr(provider, "prefetch_document_units", None) - if callable(prefetch): - prefetch() - - def prefetch_document_units_batch(self, doc_ids: Sequence[str]) -> None: - """Load a bounded document group with one query per shared chunk store.""" - providers = [ - self._docs[doc_id] - for raw_doc_id in doc_ids - if (doc_id := str(raw_doc_id).strip()) in self._docs - ] - providers_by_store: Dict[int, List[LazyKnowhereProvider]] = {} - stores_by_id: Dict[int, ChunkStore] = {} - for provider in providers: - if not isinstance(provider, LazyKnowhereProvider): - continue - store = provider._chunk_store - store_id = id(store) - stores_by_id[store_id] = store - providers_by_store.setdefault(store_id, []).append(provider) - - for store_id, lazy_providers in providers_by_store.items(): - store = stores_by_id[store_id] - batch_loader = getattr(store, "load_documents_units", None) - if not callable(batch_loader): - for provider in lazy_providers: - provider.prefetch_document_units() - continue - loaded_by_document = batch_loader( - { - provider.doc_id: list(provider._sections) - for provider in lazy_providers - } - ) - for provider in lazy_providers: - provider.install_prefetched_document_units( - loaded_by_document.get(provider.doc_id, ()) - ) - def load_persisted_score_corpus( self, doc_ids: Sequence[str], @@ -1462,12 +1195,6 @@ def load_persisted_score_corpus( queries, ) - def release_document_units(self, doc_id: str) -> None: - provider = self._docs.get(str(doc_id).strip()) - release = getattr(provider, "release_document_units", None) - if callable(release): - release() - def address_level(self, node_id: str) -> Optional[NavLevel]: sid = str(node_id or "").strip() if not sid: diff --git a/packages/shared-python/shared/services/retrieval/nav/nav_map_scores.py b/packages/shared-python/shared/services/retrieval/nav/nav_map_scores.py index 18d793e9..c646070a 100644 --- a/packages/shared-python/shared/services/retrieval/nav/nav_map_scores.py +++ b/packages/shared-python/shared/services/retrieval/nav/nav_map_scores.py @@ -1,23 +1,16 @@ from __future__ import annotations -from collections.abc import Iterator import logging import time from typing import Any, Dict, List, Optional, Sequence, Set, Tuple from .knowhere_hybrid import ( - ScoreUnitRow, build_content_search_text, build_path_search_text, build_term_search_text, - score_rows_hybrid_all, score_persisted_corpus_many, - score_unit_stream_hybrid_many, ) -# Keep one bulk read bounded when a namespace contains a very large document, -# while still replacing the per-document N+1 access pattern. -_CORPUS_PREFETCH_GROUP_SIZE = 8 _logger = logging.getLogger(__name__) @@ -253,70 +246,6 @@ def build_score_units( return units -def iter_score_units( - ts: Any, - doc_id: str, - *, - children_map: Dict[str, List[str]], - leaves: Set[str], - titles: Dict[str, str], -) -> Iterator[ScoreUnitRow]: - """Yield scoring units while releasing each lazy section after use.""" - release = getattr(ts, "release_section_units", None) - - for leaf_id in sorted(leaves): - try: - content = _section_body_text(ts, leaf_id, doc_id) or ( - titles.get(leaf_id) or _line_content(ts, leaf_id, doc_id) - ) - if content: - path_text = _ancestor_path_titles(ts, leaf_id, doc_id) - title = titles.get(leaf_id) or _line_content(ts, leaf_id, doc_id) - yield { - "chunk_id": leaf_id, - "section_id": leaf_id, - "kind": "leaf", - "content": content, - "path_text": path_text, - "path_search_text": build_path_search_text( - section_path=path_text, section_title=title or content - ), - "content_search_text": build_content_search_text(content), - "term_search_text": build_term_search_text( - content, path_text=path_text - ), - } - finally: - if callable(release): - release(leaf_id) - - for sid, kids in children_map.items(): - if not kids: - continue - try: - self_text, has_interstitial = _self_only_text(ts, sid, doc_id) - if not has_interstitial or not self_text: - continue - path_text = _ancestor_path_titles(ts, sid, doc_id) - yield { - "chunk_id": f"{sid}__self", - "section_id": sid, - "kind": "self_only", - "content": self_text, - "path_text": path_text, - "path_search_text": build_path_search_text( - section_path=path_text, section_title=titles.get(sid) or "" - ), - "content_search_text": build_content_search_text(self_text), - "term_search_text": build_term_search_text( - self_text, path_text=path_text - ), - } - finally: - if callable(release): - release(sid) - - def compute_map_scores( ts: Any, *, @@ -340,32 +269,10 @@ def compute_map_and_unit_scores( namespace: Optional[str] = None, ) -> Tuple[Dict[str, float], Dict[str, float]]: """Return (section map_scores, unit hybrid scores keyed by chunk_id).""" - if root_ids is None: - root_ids = list(ts.sections_for_doc(doc_id)) - children_map, leaves, _titles = _walk_tree(ts, doc_id, root_ids) - units = build_score_units(ts, doc_id, root_ids=root_ids) - if not units: - return {}, {} - - ns = namespace - if not ns: - import os - - ns = os.environ.get("NAV_MAP_UNIT_CACHE_NS", "").strip() or None - - scored = score_rows_hybrid_all( - units, - query, - path_texts={u["chunk_id"]: u.get("path_text") or "" for u in units}, - content_texts={u["chunk_id"]: u.get("content") or "" for u in units}, - doc_id=doc_id, - namespace=ns, + del root_ids + return compute_corpus_map_and_unit_scores( + ts, doc_ids=[doc_id], query=query, namespace=namespace ) - unit_score = { - str(r.get("chunk_id") or ""): float(r.get("score") or 0.0) for r in scored - } - map_scores = _pool_unit_scores_to_tree(children_map, leaves, unit_score) - return map_scores, unit_score def compute_corpus_map_and_unit_scores( @@ -427,52 +334,6 @@ def compute_corpus_map_and_unit_scores_many( sum(len(value[0]) for value in tree_by_doc.values()), ) - def unit_factory() -> Iterator[ScoreUnitRow]: - prefetch_batch = getattr(ts, "prefetch_document_units_batch", None) - release = getattr(ts, "release_document_units", None) - if callable(prefetch_batch): - for group_start in range( - 0, - len(valid_doc_ids), - _CORPUS_PREFETCH_GROUP_SIZE, - ): - document_group = valid_doc_ids[ - group_start : group_start + _CORPUS_PREFETCH_GROUP_SIZE - ] - prefetch_batch(document_group) - try: - for document_id in document_group: - children_map, leaves, titles = tree_by_doc[document_id] - yield from iter_score_units( - ts, - document_id, - children_map=children_map, - leaves=leaves, - titles=titles, - ) - finally: - if callable(release): - for document_id in document_group: - release(document_id) - return - - for document_id in valid_doc_ids: - children_map, leaves, titles = tree_by_doc[document_id] - prefetch = getattr(ts, "prefetch_document_units", None) - if callable(prefetch): - prefetch(document_id) - try: - yield from iter_score_units( - ts, - document_id, - children_map=children_map, - leaves=leaves, - titles=titles, - ) - finally: - if callable(release): - release(document_id) - persisted_loader = getattr(ts, "load_persisted_score_corpus", None) loader_started = time.perf_counter() persisted_corpus = ( @@ -489,7 +350,7 @@ def unit_factory() -> Iterator[ScoreUnitRow]: unit_scores_by_query = ( score_persisted_corpus_many(persisted_corpus, unique_queries) if persisted_corpus is not None - else score_unit_stream_hybrid_many(unit_factory, unique_queries) + else {query: {} for query in unique_queries} ) _logger.info( "retrieval mapnav phase=unit_scoring persisted=%s seconds=%.3f units=%d queries=%d", diff --git a/packages/shared-python/shared/services/retrieval/search/map_unit_discovery.py b/packages/shared-python/shared/services/retrieval/search/map_unit_discovery.py index b75649fb..211ab286 100644 --- a/packages/shared-python/shared/services/retrieval/search/map_unit_discovery.py +++ b/packages/shared-python/shared/services/retrieval/search/map_unit_discovery.py @@ -268,18 +268,8 @@ async def map_unit_discovery( scores_by_unit = score_persisted_corpus_many(corpus, [query]).get(query, {}) rows_by_unit_id = {row["map_unit_id"]: row for row in unit_rows} - - def _unit_has_query_hit(unit_id: str) -> bool: - # BM25 fused score can be 0 when IDF is 0 (tiny corpus). A unit - # that actually holds a query token is still a hit. - return any( - frequencies.get((unit_id, channel), {}).get(token, 0) > 0 - for channel in _MAP_SCORE_CHANNELS - for token in query_tokens - ) - ranked_unit_ids = sorted( - (unit_id for unit_id in scores_by_unit if _unit_has_query_hit(unit_id)), + (unit_id for unit_id, score in scores_by_unit.items() if score > 0.0), key=lambda unit_id: scores_by_unit[unit_id], reverse=True, )[:top_k] From 47bc458c50c6a8680f6ddd49fcbf2167974db747 Mon Sep 17 00:00:00 2001 From: chengke <404835780@qq.com> Date: Mon, 31 Aug 2026 18:15:55 +0800 Subject: [PATCH 14/19] feat: add average IDF fields and optimize retrieval logic - Introduced new fields `average_idf_path` and `average_idf_content` in the DocumentMapUnitIndex for improved scoring accuracy. - Updated the retrieval logic to calculate average IDF values during document map unit processing. - Refactored map unit discovery to utilize average IDF in scoring calculations, enhancing retrieval performance. - Removed obsolete code related to previous scoring methods to streamline the implementation. --- ...0a1b2c3d_add_map_unit_index_average_idf.py | 55 +++ .../shared/models/database/document.py | 6 + .../services/retrieval/map_unit_index.py | 13 + .../services/retrieval/nav/nav_knowhere.py | 400 ++++++------------ .../retrieval/nav/persisted_score_load.py | 72 ++++ .../retrieval/search/map_unit_discovery.py | 119 +++--- .../services/retrieval/serving_manifest.py | 46 +- 7 files changed, 343 insertions(+), 368 deletions(-) create mode 100644 apps/api/alembic/versions/8e9f0a1b2c3d_add_map_unit_index_average_idf.py create mode 100644 packages/shared-python/shared/services/retrieval/nav/persisted_score_load.py diff --git a/apps/api/alembic/versions/8e9f0a1b2c3d_add_map_unit_index_average_idf.py b/apps/api/alembic/versions/8e9f0a1b2c3d_add_map_unit_index_average_idf.py new file mode 100644 index 00000000..827e945d --- /dev/null +++ b/apps/api/alembic/versions/8e9f0a1b2c3d_add_map_unit_index_average_idf.py @@ -0,0 +1,55 @@ +"""Add average_idf columns to document_map_unit_indexes. + +Stores the rank_bm25 Okapi average IDF per channel at index-write time so +query scoring never scans all tokens to rebuild it. +""" + +from __future__ import annotations + +import sqlalchemy as sa +from alembic import op + +revision = "8e9f0a1b2c3d" +down_revision = "7d8e9f0a1b2c" +branch_labels = None +depends_on = None + + +def upgrade() -> None: + bind = op.get_bind() + inspector = sa.inspect(bind) + columns = { + col["name"] for col in inspector.get_columns("document_map_unit_indexes") + } + if "average_idf_path" not in columns: + op.add_column( + "document_map_unit_indexes", + sa.Column( + "average_idf_path", + sa.Float(), + nullable=False, + server_default="0", + ), + ) + if "average_idf_content" not in columns: + op.add_column( + "document_map_unit_indexes", + sa.Column( + "average_idf_content", + sa.Float(), + nullable=False, + server_default="0", + ), + ) + + +def downgrade() -> None: + bind = op.get_bind() + inspector = sa.inspect(bind) + columns = { + col["name"] for col in inspector.get_columns("document_map_unit_indexes") + } + if "average_idf_content" in columns: + op.drop_column("document_map_unit_indexes", "average_idf_content") + if "average_idf_path" in columns: + op.drop_column("document_map_unit_indexes", "average_idf_path") diff --git a/packages/shared-python/shared/models/database/document.py b/packages/shared-python/shared/models/database/document.py index 2ebce6cd..739221c4 100644 --- a/packages/shared-python/shared/models/database/document.py +++ b/packages/shared-python/shared/models/database/document.py @@ -356,6 +356,12 @@ class DocumentMapUnitIndex(Base): format_version: Mapped[int] = mapped_column(Integer, nullable=False) unit_count: Mapped[int] = mapped_column(Integer, nullable=False) token_count: Mapped[int] = mapped_column(Integer, nullable=False) + # rank_bm25 Okapi average IDF for the revision's units (path/content). + # Written at index time so query scoring never rescans all tokens. + average_idf_path: Mapped[float] = mapped_column(Float, nullable=False, default=0.0) + average_idf_content: Mapped[float] = mapped_column( + Float, nullable=False, default=0.0 + ) created_at: Mapped[datetime] = mapped_column( DateTime, default=utc_now_naive, nullable=False ) diff --git a/packages/shared-python/shared/services/retrieval/map_unit_index.py b/packages/shared-python/shared/services/retrieval/map_unit_index.py index 4171861c..4fff8206 100644 --- a/packages/shared-python/shared/services/retrieval/map_unit_index.py +++ b/packages/shared-python/shared/services/retrieval/map_unit_index.py @@ -23,6 +23,7 @@ UnitRow, ) from shared.services.retrieval.nav.nav_map_scores import build_score_units +from shared.services.retrieval.nav.persisted_score_load import average_idf_from_unit_dfs from shared.services.retrieval.publication_models import DocumentPublicationScope @@ -85,6 +86,8 @@ def replace_document_map_units( ) persisted_count = 0 token_count = 0 + path_unit_df: Counter[str] = Counter() + content_unit_df: Counter[str] = Counter() for sort_order, unit in enumerate(score_units): unit_id = str(unit.get("chunk_id") or "").strip() section_id = str(unit.get("section_id") or "").strip() @@ -93,6 +96,8 @@ def replace_document_map_units( map_unit_id = f"dmu_{uuid4().hex}" path_tokens = str(unit.get("path_search_text") or "").split() content_tokens = str(unit.get("content_search_text") or "").split() + path_unit_df.update(set(path_tokens)) + content_unit_df.update(set(content_tokens)) # ``provider.self_units`` already reflects root-asset remount (assets # referenced via ``connect_to`` are moved onto the text section that # embeds them), so this is the same ownership the query-time scorer @@ -139,6 +144,14 @@ def replace_document_map_units( format_version=MAP_UNIT_INDEX_FORMAT_VERSION, unit_count=persisted_count, token_count=token_count, + average_idf_path=average_idf_from_unit_dfs( + unit_count=persisted_count, + token_document_frequency=path_unit_df, + ), + average_idf_content=average_idf_from_unit_dfs( + unit_count=persisted_count, + token_document_frequency=content_unit_df, + ), ) ) diff --git a/packages/shared-python/shared/services/retrieval/nav/nav_knowhere.py b/packages/shared-python/shared/services/retrieval/nav/nav_knowhere.py index c5cd4fbb..577b1b95 100644 --- a/packages/shared-python/shared/services/retrieval/nav/nav_knowhere.py +++ b/packages/shared-python/shared/services/retrieval/nav/nav_knowhere.py @@ -23,7 +23,6 @@ import os import logging import time -from hashlib import sha256 from dataclasses import dataclass, field from typing import ( Any, @@ -42,7 +41,6 @@ from .nav_address import NavLevel from .nav_hierarchy import NodeMeta from .knowhere_hybrid import ( - PersistedBm25Stats, PersistedScoreCorpus, PersistedScoreUnit, tokenize_query_for_ranker, @@ -207,15 +205,11 @@ def __init__( self._revisions = dict(revisions) self._excluded_sections = set(excluded_sections or ()) self._conn: Optional[_SyncConnection] = None - self._score_manifest_cache: Optional[ - tuple[tuple[tuple[str, str], ...], list[Sequence[object]]] - ] = None self._score_unit_rows_cache: dict[ - tuple[tuple[str, str], ...], list[Sequence[object]] + tuple[tuple[str, str], ...], list[dict[str, object]] ] = {} - self._score_frequency_cache: dict[ - tuple[tuple[str, str], ...], - dict[tuple[str, str], dict[str, int]], + self._score_average_idf_cache: dict[ + tuple[tuple[str, str], ...], tuple[float, float] ] = {} def _connection(self) -> "_SyncConnection": @@ -300,7 +294,17 @@ def load_persisted_score_corpus( allowed_section_ids_by_document: Mapping[str, Sequence[str]], queries: Sequence[str], ) -> Optional[PersistedScoreCorpus]: - """Load query-relevant score inputs when every revision is indexed.""" + """Load query-token score inputs from map-unit tables when indexed. + + Units and lengths come from ``document_map_units``. Frequencies come from + ``document_map_unit_tokens`` filtered to the query tokens. Average IDF + comes from ``document_map_unit_indexes`` (written at index time). + """ + from shared.services.retrieval.nav.persisted_score_load import ( + build_channel_bm25_stats, + combine_average_idf, + ) + revisions = [ (document_id, self._revisions[document_id]) for raw_document_id in document_ids @@ -312,229 +316,141 @@ def load_persisted_score_corpus( revision_params: List[object] = [ value for revision in revisions for value in revision ] + unique_queries = list(dict.fromkeys(str(query) for query in queries)) + query_tokens = list( + dict.fromkeys( + token + for query in unique_queries + for token in tokenize_query_for_ranker(query) + ) + ) cur = self._connection().cursor() try: revision_key = tuple(revisions) - cached_manifests = self._score_manifest_cache - if cached_manifests is not None and cached_manifests[0] == revision_key: - manifests = cached_manifests[1] - _logger.info( - "retrieval map-index load stage=manifests cache_hit rows=%d", - len(manifests), + stage_started = time.perf_counter() + try: + cur.execute( + "SELECT indexes.document_id, indexes.job_result_id, " + "indexes.format_version, indexes.unit_count, " + "indexes.average_idf_path, indexes.average_idf_content " + "FROM document_map_unit_indexes AS indexes " + f"JOIN (VALUES {values_sql}) AS revisions(document_id, job_result_id) " + "ON indexes.document_id = revisions.document_id " + "AND indexes.job_result_id = revisions.job_result_id", + revision_params, ) - else: - stage_started = time.perf_counter() - try: - cur.execute( - "SELECT indexes.document_id, indexes.job_result_id, " - "indexes.format_version, indexes.unit_count, indexes.token_count, " - "manifests.payload_zlib, manifests.checksum, manifests.format_version, " - "statistics.payload_zlib, statistics.checksum, statistics.format_version " - "FROM document_map_unit_indexes AS indexes " - f"JOIN (VALUES {values_sql}) AS revisions(document_id, job_result_id) " - "ON indexes.document_id = revisions.document_id " - "AND indexes.job_result_id = revisions.job_result_id " - "JOIN retrieval_serving_revision_manifests AS manifests " - "ON manifests.document_id = indexes.document_id " - "AND manifests.job_result_id = indexes.job_result_id " - "JOIN retrieval_serving_revision_stats AS statistics " - "ON statistics.document_id = indexes.document_id " - "AND statistics.job_result_id = indexes.job_result_id", - revision_params, - ) - except Exception as exc: - if getattr(exc, "pgcode", None) == "42P01": - return None - raise - manifests = list(cur.fetchall()) - self._score_manifest_cache = (revision_key, manifests) + except Exception as exc: + if getattr(exc, "pgcode", None) in {"42P01", "42703"}: + return None + raise + index_rows = list(cur.fetchall()) _logger.info( - "retrieval map-index load stage=manifests seconds=%.3f rows=%d", - time.perf_counter() - stage_started if cached_manifests is None else 0.0, - len(manifests), + "retrieval map-index load stage=indexes seconds=%.3f rows=%d", + time.perf_counter() - stage_started, + len(index_rows), ) - if len(manifests) != len(revisions) or any( - len(row) < 11 - or int(row[2]) != _MAP_UNIT_INDEX_FORMAT_VERSION - or not row[5] - or not row[6] - or not row[8] - or not row[9] - for row in manifests + if len(index_rows) != len(revisions) or any( + len(row) < 6 or int(row[2]) != _MAP_UNIT_INDEX_FORMAT_VERSION + for row in index_rows ): return None - decoded_manifests: dict[tuple[str, str], dict[str, Any]] = {} - decoded_statistics: dict[tuple[str, str], dict[str, Any]] = {} - try: - for row in manifests: - decoded_manifests[(str(row[0]), str(row[1]))] = decode_serving_manifest( - bytes(row[5]), - checksum=str(row[6]), - format_version=int(row[7]), - ) - decoded_statistics[(str(row[0]), str(row[1]))] = decode_serving_manifest( - bytes(row[8]), - checksum=str(row[9]), - format_version=int(row[10]), - ) - except ValueError: - return None - # The marker is written last in the same transaction that inserts - # all units and token rows. A committed marker therefore denotes - # one complete revision snapshot; avoid recounting millions of - # token rows on every request. Any rebuild deletes the marker - # first, so readers fall back to legacy scoring until completion. - # The public chunk id is content-derived and may repeat within a - # revision. The persisted scorer keys scores by that id, so use - # the legacy payload path whenever ambiguity would change results. + if revision_key in self._score_average_idf_cache: + average_idf_path, average_idf_content = self._score_average_idf_cache[ + revision_key + ] + else: + average_idf_path = combine_average_idf( + [(float(row[4] or 0.0), int(row[3] or 0)) for row in index_rows] + ) + average_idf_content = combine_average_idf( + [(float(row[5] or 0.0), int(row[3] or 0)) for row in index_rows] + ) + self._score_average_idf_cache[revision_key] = ( + average_idf_path, + average_idf_content, + ) + allowed_by_document = { str(document_id): {str(section_id) for section_id in section_ids} for document_id, section_ids in allowed_section_ids_by_document.items() } - allowed_pairs = [ + allowed_pairs_set = { (document_id, section_id) for document_id, section_ids in allowed_by_document.items() for section_id in section_ids - ] - unit_cache_key = revision_key - all_unit_rows = self._score_unit_rows_cache.get(unit_cache_key, []) - units_cache_hit = unit_cache_key in self._score_unit_rows_cache - unit_rows: list[Sequence[object]] = [] - if allowed_pairs: - if not units_cache_hit: - stage_started = time.perf_counter() - manifest_rows: list[Sequence[object]] = [] - for document_id, job_result_id in revisions: - payload = decoded_manifests.get((document_id, job_result_id), {}) - raw_units = payload.get("map_units") - if not isinstance(raw_units, list): - manifest_rows = [] - break - for raw_unit in raw_units: - if not isinstance(raw_unit, dict): - manifest_rows = [] - break - row_id = str(raw_unit.get("row_id") or "").strip() - unit_id = str(raw_unit.get("unit_id") or "").strip() - if not row_id or not unit_id: - manifest_rows = [] - break - manifest_rows.append( - ( - row_id, - document_id, - unit_id, - str(raw_unit.get("section_id") or ""), - int(raw_unit.get("path_token_count") or 0), - int(raw_unit.get("content_token_count") or 0), - ) - ) - if not manifest_rows and raw_units: - break - expected_unit_count = sum(int(row[3]) for row in manifests) - all_unit_rows = ( - manifest_rows - if len(manifest_rows) == expected_unit_count - else [] - ) - if expected_unit_count and not all_unit_rows: - return None - self._score_unit_rows_cache[unit_cache_key] = all_unit_rows - else: - stage_started = time.perf_counter() - allowed_pairs_set = set(allowed_pairs) - unit_rows = [ - row - for row in all_unit_rows - if (str(row[1]), str(row[3])) in allowed_pairs_set + } + all_unit_rows = self._score_unit_rows_cache.get(revision_key) + if all_unit_rows is None: + stage_started = time.perf_counter() + cur.execute( + "SELECT units.id, units.document_id, units.unit_id, " + "units.section_id, units.path_token_count, units.content_token_count " + "FROM document_map_units AS units " + f"JOIN (VALUES {values_sql}) AS revisions(document_id, job_result_id) " + "ON units.document_id = revisions.document_id " + "AND units.job_result_id = revisions.job_result_id " + "ORDER BY units.document_id, units.sort_order, units.unit_id", + revision_params, + ) + all_unit_rows = [ + { + "map_unit_id": str(row[0]), + "document_id": str(row[1]), + "unit_id": str(row[2]), + "section_id": str(row[3] or ""), + "path_token_count": int(row[4] or 0), + "content_token_count": int(row[5] or 0), + } + for row in cur.fetchall() ] + expected_unit_count = sum(int(row[3] or 0) for row in index_rows) + if expected_unit_count and len(all_unit_rows) != expected_unit_count: + return None + self._score_unit_rows_cache[revision_key] = all_unit_rows _logger.info( "retrieval map-index load stage=units seconds=%.3f rows=%d cache_hit=%s", - time.perf_counter() - stage_started if not units_cache_hit else 0.0, - len(unit_rows), - units_cache_hit, + time.perf_counter() - stage_started, + len(all_unit_rows), + False, ) - map_unit_ids = [str(row[0]) for row in unit_rows] - unique_queries = list(dict.fromkeys(str(query) for query in queries)) - query_tokens_by_query = { - query: tokenize_query_for_ranker(query) for query in unique_queries - } - query_tokens = list( - dict.fromkeys( - token - for query in unique_queries - for token in query_tokens_by_query[query] + else: + _logger.info( + "retrieval map-index load stage=units seconds=0.000 rows=%d cache_hit=%s", + len(all_unit_rows), + True, ) - ) + + unit_rows = [ + row + for row in all_unit_rows + if (str(row["document_id"]), str(row["section_id"])) in allowed_pairs_set + ] frequencies: Dict[Tuple[str, str], Dict[str, int]] = {} - if map_unit_ids and query_tokens: - full_frequency_map = self._score_frequency_cache.get(revision_key) - if full_frequency_map is None: - full_frequency_map = {} - for document_id, job_result_id in revisions: - payload = decoded_statistics.get((document_id, job_result_id), {}) - unit_frequencies = payload.get("unit_frequencies", {}) - if not isinstance(unit_frequencies, dict): - continue - for map_unit_id, by_channel in unit_frequencies.items(): - if not isinstance(by_channel, dict): - continue - for channel in _MAP_SCORE_CHANNELS: - values = by_channel.get(channel, {}) - if isinstance(values, dict): - full_frequency_map[(str(map_unit_id), channel)] = { - str(token): int(value) - for token, value in values.items() - } - self._score_frequency_cache[revision_key] = full_frequency_map - frequencies = { - key: { - token: value - for token, value in values.items() - if token in query_tokens - } - for key, values in full_frequency_map.items() - if key[0] in map_unit_ids - } - expected_units_by_revision = { - (str(row[0]), str(row[1])): int(row[3]) for row in manifests - } - statistics_complete = all( - int( - decoded_statistics.get((document_id, job_result_id), {}).get( - "unit_count", -1 - ) - ) - == expected_units_by_revision.get((document_id, job_result_id), -1) - for document_id, job_result_id in revisions - ) - manifest_frequency_complete = bool(map_unit_ids) and statistics_complete - if map_unit_ids and query_tokens and not manifest_frequency_complete: - query_token_hashes = [ - sha256(token.encode("utf-8")).hexdigest() for token in query_tokens - ] + if unit_rows and query_tokens: stage_started = time.perf_counter() cur.execute( - "SELECT map_unit_id, channel, token, frequency " - "FROM document_map_unit_tokens " - "WHERE map_unit_id = ANY(%s) AND token_hash = ANY(%s) " - "AND token = ANY(%s) AND channel = ANY(%s)", - ( - map_unit_ids, - query_token_hashes, - query_tokens, - list(_MAP_SCORE_CHANNELS), - ), + "SELECT units.id, tokens.channel, tokens.token, tokens.frequency " + "FROM document_map_unit_tokens AS tokens " + "JOIN document_map_units AS units ON units.id = tokens.map_unit_id " + f"JOIN (VALUES {values_sql}) AS revisions(document_id, job_result_id) " + "ON units.document_id = revisions.document_id " + "AND units.job_result_id = revisions.job_result_id " + "WHERE tokens.token = ANY(%s) AND tokens.channel = ANY(%s)", + [*revision_params, list(query_tokens), list(_MAP_SCORE_CHANNELS)], ) + allowed_map_unit_ids = {str(row["map_unit_id"]) for row in unit_rows} for map_unit_id, channel, token, frequency in cur.fetchall(): + if str(map_unit_id) not in allowed_map_unit_ids: + continue frequencies.setdefault((str(map_unit_id), str(channel)), {})[ str(token) ] = int(frequency) _logger.info( - "retrieval map-index load stage=frequencies seconds=%.3f units=%d", + "retrieval map-index load stage=frequencies seconds=%.3f units=%d tokens=%d", time.perf_counter() - stage_started, - len(map_unit_ids), + len(unit_rows), + len(query_tokens), ) _logger.info( @@ -542,33 +458,35 @@ def load_persisted_score_corpus( len(unit_rows), len(unique_queries), ) - path_stats = self._load_persisted_bm25_stats( - cur, + path_stats = build_channel_bm25_stats( unit_rows=unit_rows, - map_unit_ids=map_unit_ids, + map_unit_id_field="map_unit_id", + length_field="path_token_count", channel="path", query_tokens=query_tokens, frequencies=frequencies, - length_index=4, + average_idf=average_idf_path, ) - content_stats = self._load_persisted_bm25_stats( - cur, + content_stats = build_channel_bm25_stats( unit_rows=unit_rows, - map_unit_ids=map_unit_ids, + map_unit_id_field="map_unit_id", + length_field="content_token_count", channel="content", query_tokens=query_tokens, frequencies=frequencies, - length_index=5, + average_idf=average_idf_content, ) return PersistedScoreCorpus( units=[ PersistedScoreUnit( - unit_id=str(row[2]), - path_length=int(row[4]), - content_length=int(row[5]), - path_frequencies=frequencies.get((str(row[0]), "path"), {}), + unit_id=str(row["unit_id"]), + path_length=int(row["path_token_count"]), + content_length=int(row["content_token_count"]), + path_frequencies=frequencies.get( + (str(row["map_unit_id"]), "path"), {} + ), content_frequencies=frequencies.get( - (str(row[0]), "content"), {} + (str(row["map_unit_id"]), "content"), {} ), ) for row in unit_rows @@ -579,58 +497,6 @@ def load_persisted_score_corpus( finally: cur.close() - def _load_persisted_bm25_stats( - self, - cur: "_SyncCursor", - *, - unit_rows: Sequence[Sequence[object]], - map_unit_ids: Sequence[str], - channel: str, - query_tokens: Sequence[str], - frequencies: Mapping[Tuple[str, str], Mapping[str, int]], - length_index: int, - ) -> PersistedBm25Stats: - lengths = [ - int(row[length_index]) for row in unit_rows if int(row[length_index]) > 0 - ] - document_count = len(lengths) - document_frequency = { - token: sum( - 1 - for row in unit_rows - if frequencies.get((str(row[0]), channel), {}).get(token, 0) > 0 - ) - for token in query_tokens - } - needs_average_idf = any( - frequency > document_count / 2 for frequency in document_frequency.values() - ) - average_idf = 0.0 - if needs_average_idf and map_unit_ids and document_count: - stage_started = time.perf_counter() - cur.execute( - "SELECT COALESCE(AVG(LN((%s - frequencies.document_frequency + 0.5) " - "/ (frequencies.document_frequency + 0.5))), 0.0) " - "FROM (SELECT token, COUNT(*) AS document_frequency " - "FROM document_map_unit_tokens " - "WHERE map_unit_id = ANY(%s) AND channel = %s " - "GROUP BY token) AS frequencies", - (document_count, list(map_unit_ids), channel), - ) - row = cur.fetchone() - average_idf = float(row[0]) if row else 0.0 - _logger.info( - "retrieval map-index load stage=average_idf channel=%s seconds=%.3f", - channel, - time.perf_counter() - stage_started, - ) - return PersistedBm25Stats( - document_count=document_count, - total_length=sum(lengths), - document_frequency=document_frequency, - average_idf=average_idf, - ) - def close(self) -> None: if self._conn is not None: self._conn.close() diff --git a/packages/shared-python/shared/services/retrieval/nav/persisted_score_load.py b/packages/shared-python/shared/services/retrieval/nav/persisted_score_load.py new file mode 100644 index 00000000..9ce10c2e --- /dev/null +++ b/packages/shared-python/shared/services/retrieval/nav/persisted_score_load.py @@ -0,0 +1,72 @@ +"""Shared helpers for building persisted BM25 corpora from DB rows.""" + +from __future__ import annotations + +import math +from collections.abc import Mapping, Sequence +from typing import Any + +from shared.services.retrieval.nav.knowhere_hybrid import PersistedBm25Stats + + +def average_idf_from_unit_dfs( + *, + unit_count: int, + token_document_frequency: Mapping[str, int], +) -> float: + """rank_bm25 Okapi average IDF over every token that appears in the corpus.""" + if unit_count <= 0 or not token_document_frequency: + return 0.0 + idfs = [ + math.log(unit_count - frequency + 0.5) - math.log(frequency + 0.5) + for frequency in token_document_frequency.values() + if frequency > 0 + ] + if not idfs: + return 0.0 + return sum(idfs) / len(idfs) + + +def combine_average_idf(parts: Sequence[tuple[float, int]]) -> float: + """Unit-count-weighted mean of per-revision average IDF values.""" + total_units = sum(int(unit_count) for _average, unit_count in parts) + if total_units <= 0: + return 0.0 + return ( + sum(float(average) * int(unit_count) for average, unit_count in parts) + / total_units + ) + + +def build_channel_bm25_stats( + *, + unit_rows: Sequence[Mapping[str, Any]], + map_unit_id_field: str, + length_field: str, + channel: str, + query_tokens: Sequence[str], + frequencies: Mapping[tuple[str, str], Mapping[str, int]], + average_idf: float, +) -> PersistedBm25Stats: + """Build channel stats from already-fetched unit rows and query-token freqs.""" + lengths = [ + int(row[length_field]) for row in unit_rows if int(row[length_field]) > 0 + ] + document_count = len(lengths) + document_frequency = { + token: sum( + 1 + for row in unit_rows + if frequencies.get((str(row[map_unit_id_field]), channel), {}).get( + token, 0 + ) + > 0 + ) + for token in query_tokens + } + return PersistedBm25Stats( + document_count=document_count, + total_length=sum(lengths), + document_frequency=document_frequency, + average_idf=float(average_idf), + ) diff --git a/packages/shared-python/shared/services/retrieval/search/map_unit_discovery.py b/packages/shared-python/shared/services/retrieval/search/map_unit_discovery.py index 211ab286..38b9762d 100644 --- a/packages/shared-python/shared/services/retrieval/search/map_unit_discovery.py +++ b/packages/shared-python/shared/services/retrieval/search/map_unit_discovery.py @@ -30,12 +30,15 @@ normalize_chunk_type, ) from shared.services.retrieval.nav.knowhere_hybrid import ( - PersistedBm25Stats, PersistedScoreCorpus, PersistedScoreUnit, score_persisted_corpus_many, tokenize_query_for_ranker, ) +from shared.services.retrieval.nav.persisted_score_load import ( + build_channel_bm25_stats, + combine_average_idf, +) from shared.services.retrieval.search.scoring import normalize_row_scores from shared.services.retrieval.search.section_filters import is_excluded_section from shared.services.retrieval.settings import ASSET_CHUNK_TYPES @@ -209,17 +212,19 @@ async def map_unit_discovery( if not unit_rows: return DiscoveryResult(status="discovery_done", payload={"fused_rows": []}) - map_unit_ids = [row["map_unit_id"] for row in unit_rows] frequency_result = await db.execute( text( - "SELECT map_unit_id, channel, token, frequency " - "FROM document_map_unit_tokens " - "WHERE map_unit_id = ANY(:unit_ids) " - "AND channel = ANY(:channels) " - "AND token = ANY(:tokens)" + cte + + """ + SELECT tokens.map_unit_id, tokens.channel, tokens.token, tokens.frequency + FROM document_map_unit_tokens AS tokens + JOIN scoped_units ON scoped_units.map_unit_id = tokens.map_unit_id + WHERE tokens.channel = ANY(:channels) + AND tokens.token = ANY(:tokens) + """ ), { - "unit_ids": map_unit_ids, + **params, "channels": list(_MAP_SCORE_CHANNELS), "tokens": query_tokens, }, @@ -230,23 +235,53 @@ async def map_unit_discovery( int(frequency) ) - path_stats = await _build_bm25_stats( - db, + index_result = await db.execute( + text( + cte + + """ + SELECT indexes.average_idf_path, indexes.average_idf_content, + indexes.unit_count + FROM document_map_unit_indexes AS indexes + JOIN ( + SELECT DISTINCT document_id, job_result_id FROM scoped_units + ) AS scoped_revisions + ON indexes.document_id = scoped_revisions.document_id + AND indexes.job_result_id = scoped_revisions.job_result_id + """ + ), + params, + ) + index_parts = [ + (float(path_idf or 0.0), float(content_idf or 0.0), int(unit_count or 0)) + for path_idf, content_idf, unit_count in index_result.all() + ] + average_idf_path = combine_average_idf( + [(path_idf, unit_count) for path_idf, _content_idf, unit_count in index_parts] + ) + average_idf_content = combine_average_idf( + [ + (content_idf, unit_count) + for _path_idf, content_idf, unit_count in index_parts + ] + ) + + path_stats = build_channel_bm25_stats( unit_rows=unit_rows, - map_unit_ids=map_unit_ids, + map_unit_id_field="map_unit_id", + length_field="path_token_count", channel="path", query_tokens=query_tokens, frequencies=frequencies, - length_field="path_token_count", + average_idf=average_idf_path, ) - content_stats = await _build_bm25_stats( - db, + content_stats = build_channel_bm25_stats( unit_rows=unit_rows, - map_unit_ids=map_unit_ids, + map_unit_id_field="map_unit_id", + length_field="content_token_count", channel="content", query_tokens=query_tokens, frequencies=frequencies, - length_field="content_token_count", + average_idf=average_idf_content, ) corpus = PersistedScoreCorpus( @@ -310,58 +345,6 @@ async def map_unit_discovery( return DiscoveryResult(status="error", error=str(exc), latency_ms=latency) -async def _build_bm25_stats( - session: AsyncSession, - *, - unit_rows: list[dict[str, Any]], - map_unit_ids: list[str], - channel: str, - query_tokens: list[str], - frequencies: dict[tuple[str, str], dict[str, int]], - length_field: str, -) -> PersistedBm25Stats: - lengths = [ - int(row[length_field]) for row in unit_rows if int(row[length_field]) > 0 - ] - document_count = len(lengths) - document_frequency = { - token: sum( - 1 - for row in unit_rows - if frequencies.get((row["map_unit_id"], channel), {}).get(token, 0) > 0 - ) - for token in query_tokens - } - needs_average_idf = any( - frequency > document_count / 2 for frequency in document_frequency.values() - ) - average_idf = 0.0 - if needs_average_idf and map_unit_ids and document_count: - result = await session.execute( - text( - "SELECT COALESCE(AVG(LN((:document_count - frequencies.document_frequency " - "+ 0.5) / (frequencies.document_frequency + 0.5))), 0.0) " - "FROM (SELECT token, COUNT(*) AS document_frequency " - "FROM document_map_unit_tokens " - "WHERE map_unit_id = ANY(:unit_ids) AND channel = :channel " - "GROUP BY token) AS frequencies" - ), - { - "document_count": document_count, - "unit_ids": map_unit_ids, - "channel": channel, - }, - ) - row = result.first() - average_idf = float(row[0]) if row else 0.0 - return PersistedBm25Stats( - document_count=document_count, - total_length=sum(lengths), - document_frequency=document_frequency, - average_idf=average_idf, - ) - - def _as_metadata_dict(value: object) -> dict[str, Any]: if isinstance(value, dict): return value diff --git a/packages/shared-python/shared/services/retrieval/serving_manifest.py b/packages/shared-python/shared/services/retrieval/serving_manifest.py index 5fe979e0..88811320 100644 --- a/packages/shared-python/shared/services/retrieval/serving_manifest.py +++ b/packages/shared-python/shared/services/retrieval/serving_manifest.py @@ -58,14 +58,6 @@ def build_revision_serving_payload( ) ) ) - map_units = list( - db.scalars( - select(DocumentMapUnit) - .where(DocumentMapUnit.document_id == scope.document_id) - .where(DocumentMapUnit.job_result_id == scope.job_result_id) - .order_by(DocumentMapUnit.sort_order, DocumentMapUnit.unit_id) - ) - ) section_path_by_id = { section.section_id: section.section_path for section in sections } @@ -121,18 +113,6 @@ def build_revision_serving_payload( } for chunk in chunks ], - "map_units": [ - { - "row_id": unit.id, - "unit_id": unit.unit_id, - "section_id": unit.section_id, - "unit_kind": unit.unit_kind, - "path_token_count": unit.path_token_count, - "content_token_count": unit.content_token_count, - "sort_order": unit.sort_order, - } - for unit in map_units - ], "root_asset_ids": sorted(root_asset_ids), "remounted_assets_by_section": remounted_assets, } @@ -143,7 +123,12 @@ def build_revision_statistics_payload( *, scope: DocumentPublicationScope, ) -> dict[str, Any]: - """Build compressed scoring contributions for one revision.""" + """Build compressed scoring contributions for one revision. + + Stores aggregate token frequencies for namespace statistics rebuild. + Per-unit frequencies stay in ``document_map_unit_tokens`` and are loaded + at query time by the query tokens only. + """ units = list( db.scalars( select(DocumentMapUnit) @@ -153,9 +138,8 @@ def build_revision_statistics_payload( ) unit_ids = [unit.id for unit in units] frequencies: dict[str, dict[str, int]] = {"path": {}, "content": {}} - unit_frequencies: dict[str, dict[str, dict[str, int]]] = {} if unit_ids: - for map_unit_id, channel, token, frequency in db.execute( + for _map_unit_id, channel, token, frequency in db.execute( select( DocumentMapUnitToken.map_unit_id, DocumentMapUnitToken.channel, @@ -164,15 +148,12 @@ def build_revision_statistics_payload( ).where(DocumentMapUnitToken.map_unit_id.in_(unit_ids)) ).all(): channel_key = str(channel) - if channel_key in frequencies: - token_key = str(token) - frequency_value = int(frequency) - frequencies[channel_key][token_key] = ( - frequencies[channel_key].get(token_key, 0) + frequency_value - ) - unit_frequencies.setdefault(str(map_unit_id), {}).setdefault( - channel_key, {} - )[token_key] = frequency_value + if channel_key not in frequencies: + continue + token_key = str(token) + frequencies[channel_key][token_key] = frequencies[channel_key].get( + token_key, 0 + ) + int(frequency) return { "document_id": scope.document_id, "job_result_id": scope.job_result_id, @@ -182,7 +163,6 @@ def build_revision_statistics_payload( int(unit.content_token_count or 0) for unit in units ), "token_frequencies": frequencies, - "unit_frequencies": unit_frequencies, } From 496501b77d66e5f173c038674ebd4158b092095f Mon Sep 17 00:00:00 2001 From: chengke <404835780@qq.com> Date: Mon, 31 Aug 2026 18:21:20 +0800 Subject: [PATCH 15/19] chore: drop unused serving_manifest import in nav_knowhere Co-authored-by: Cursor --- .../shared-python/shared/services/retrieval/nav/nav_knowhere.py | 1 - 1 file changed, 1 deletion(-) diff --git a/packages/shared-python/shared/services/retrieval/nav/nav_knowhere.py b/packages/shared-python/shared/services/retrieval/nav/nav_knowhere.py index 577b1b95..716ff1ec 100644 --- a/packages/shared-python/shared/services/retrieval/nav/nav_knowhere.py +++ b/packages/shared-python/shared/services/retrieval/nav/nav_knowhere.py @@ -45,7 +45,6 @@ PersistedScoreUnit, tokenize_query_for_ranker, ) -from shared.services.retrieval.serving_manifest import decode_serving_manifest _ASSET_TYPES = ("table", "image") # Knowhere sentinel path for the virtual document container (not a collectable leaf). From 5fee16666a8c30a65d43b3e96ccb0938f415b1a4 Mon Sep 17 00:00:00 2001 From: chengke <404835780@qq.com> Date: Mon, 31 Aug 2026 18:37:47 +0800 Subject: [PATCH 16/19] feat: add read-only --check for serving-index fallback readiness Co-authored-by: Cursor --- apps/api/scripts/backfill_map_unit_indexes.py | 214 +++++++++++++++++- 1 file changed, 213 insertions(+), 1 deletion(-) diff --git a/apps/api/scripts/backfill_map_unit_indexes.py b/apps/api/scripts/backfill_map_unit_indexes.py index be2b4dbb..6eed78ef 100644 --- a/apps/api/scripts/backfill_map_unit_indexes.py +++ b/apps/api/scripts/backfill_map_unit_indexes.py @@ -6,6 +6,10 @@ derived tables leave them empty intentionally. Run this command after deployment with ``--apply`` so each revision is rebuilt and committed independently; without ``--apply`` it is a read-only inventory. + +Use ``--check`` after backfill to verify whether query-time snapshot +fallbacks (manifest_merge / table_scan) would still fire, and whether +map-unit indexes are complete for scoring. """ # ruff: noqa: E402 @@ -15,6 +19,8 @@ import argparse import os import sys +from collections import defaultdict +from dataclasses import dataclass from pathlib import Path @@ -45,7 +51,12 @@ def _bootstrap_python_path() -> None: from sqlalchemy import select from shared.core.database_sync import get_sync_session_factory -from shared.models.database.document import Document +from shared.models.database.document import ( + Document, + DocumentMapUnitIndex, + RetrievalNamespaceMapSnapshot, + RetrievalServingRevisionManifest, +) from shared.services.retrieval.map_unit_index import replace_document_map_units from shared.services.retrieval.namespace_map_snapshot import ( patch_namespace_map_snapshot, @@ -56,6 +67,7 @@ def _bootstrap_python_path() -> None: lock_namespace_generation, ) from shared.services.retrieval.serving_manifest import ( + decode_serving_manifest, persist_revision_serving_state, rebuild_namespace_serving_statistics, ) @@ -70,6 +82,14 @@ def _build_parser() -> argparse.ArgumentParser: action="store_true", help="Build and commit each current revision index.", ) + parser.add_argument( + "--check", + action="store_true", + help=( + "Read-only: report whether snapshot fallbacks would still fire " + "and whether map-unit indexes are complete. Exit 1 if not ready." + ), + ) parser.add_argument( "--document-id", default="", help="Limit the backfill to one document." ) @@ -90,6 +110,192 @@ def _load_documents(document_id: str) -> list[Document]: return list(db.scalars(statement).all()) +@dataclass(frozen=True) +class NamespaceFallbackReport: + user_id: str + namespace: str + active_docs: int + snapshot_status: str + missing_from_snapshot: int + missing_map_index: int + missing_revision_manifest: int + suspicious_zero_idf: int + would_hit_snapshot_fallback: bool + scoring_incomplete: bool + + @property + def ready(self) -> bool: + return not self.would_hit_snapshot_fallback and not self.scoring_incomplete + + +def check_fallback_readiness(*, document_id: str = "") -> list[NamespaceFallbackReport]: + """Inspect active revisions for snapshot coverage and map-unit indexes.""" + documents = _load_documents(document_id) + by_scope: dict[tuple[str, str], list[Document]] = defaultdict(list) + for document in documents: + by_scope[(document.user_id, document.namespace)].append(document) + + session_factory = get_sync_session_factory() + reports: list[NamespaceFallbackReport] = [] + with session_factory() as db: + for (user_id, namespace), scoped_docs in sorted( + by_scope.items(), key=lambda item: (item[0][0], item[0][1]) + ): + revisions = [ + (doc.document_id, str(doc.current_job_result_id)) + for doc in scoped_docs + if doc.current_job_result_id + ] + if not revisions: + continue + + snapshot = db.execute( + select(RetrievalNamespaceMapSnapshot) + .where(RetrievalNamespaceMapSnapshot.user_id == user_id) + .where(RetrievalNamespaceMapSnapshot.namespace == namespace) + ).scalar_one_or_none() + + snapshot_documents: dict[str, object] | None = None + snapshot_status = "missing" + if snapshot is None: + snapshot_status = "missing" + else: + try: + payload = decode_serving_manifest( + bytes(snapshot.payload_zlib), + checksum=str(snapshot.checksum), + format_version=int(snapshot.format_version), + ) + decoded = payload.get("documents") + if isinstance(decoded, dict): + snapshot_documents = decoded + snapshot_status = "ok" + else: + snapshot_status = "corrupt" + except (TypeError, ValueError): + snapshot_status = "corrupt" + + missing_from_snapshot = 0 + if snapshot_documents is None: + missing_from_snapshot = len(revisions) + else: + for document_id_value, job_result_id in revisions: + entry = snapshot_documents.get(document_id_value) + if ( + not isinstance(entry, dict) + or str(entry.get("job_result_id") or "") != job_result_id + ): + missing_from_snapshot += 1 + if missing_from_snapshot and snapshot_status == "ok": + snapshot_status = "stale" + + index_rows = list( + db.execute( + select( + DocumentMapUnitIndex.document_id, + DocumentMapUnitIndex.job_result_id, + DocumentMapUnitIndex.unit_count, + DocumentMapUnitIndex.average_idf_path, + DocumentMapUnitIndex.average_idf_content, + ).where( + DocumentMapUnitIndex.document_id.in_( + [document_id_value for document_id_value, _ in revisions] + ) + ) + ).all() + ) + index_by_revision = { + (str(document_id_value), str(job_result_id)): ( + int(unit_count or 0), + float(average_idf_path or 0.0), + float(average_idf_content or 0.0), + ) + for document_id_value, job_result_id, unit_count, average_idf_path, average_idf_content in index_rows + } + missing_map_index = 0 + suspicious_zero_idf = 0 + for document_id_value, job_result_id in revisions: + stats = index_by_revision.get((document_id_value, job_result_id)) + if stats is None: + missing_map_index += 1 + continue + unit_count, average_idf_path, average_idf_content = stats + if ( + unit_count > 0 + and average_idf_path == 0.0 + and average_idf_content == 0.0 + ): + suspicious_zero_idf += 1 + + manifest_rows = list( + db.execute( + select( + RetrievalServingRevisionManifest.document_id, + RetrievalServingRevisionManifest.job_result_id, + ).where( + RetrievalServingRevisionManifest.document_id.in_( + [document_id_value for document_id_value, _ in revisions] + ) + ) + ).all() + ) + manifest_keys = { + (str(document_id_value), str(job_result_id)) + for document_id_value, job_result_id in manifest_rows + } + missing_revision_manifest = sum( + 1 + for document_id_value, job_result_id in revisions + if (document_id_value, job_result_id) not in manifest_keys + ) + + would_hit_snapshot_fallback = ( + snapshot_status != "ok" or missing_from_snapshot > 0 + ) + scoring_incomplete = missing_map_index > 0 or suspicious_zero_idf > 0 + reports.append( + NamespaceFallbackReport( + user_id=user_id, + namespace=namespace, + active_docs=len(revisions), + snapshot_status=snapshot_status, + missing_from_snapshot=missing_from_snapshot, + missing_map_index=missing_map_index, + missing_revision_manifest=missing_revision_manifest, + suspicious_zero_idf=suspicious_zero_idf, + would_hit_snapshot_fallback=would_hit_snapshot_fallback, + scoring_incomplete=scoring_incomplete, + ) + ) + return reports + + +def print_fallback_check(reports: list[NamespaceFallbackReport]) -> int: + """Print readiness report. Returns process exit code (0 ready, 1 not).""" + if not reports: + print("check: no active documents found") + return 0 + + failed = 0 + for report in reports: + status = "READY" if report.ready else "NOT_READY" + if not report.ready: + failed += 1 + print( + f"check status={status} user={report.user_id} namespace={report.namespace} " + f"active_docs={report.active_docs} snapshot={report.snapshot_status} " + f"missing_from_snapshot={report.missing_from_snapshot} " + f"missing_map_index={report.missing_map_index} " + f"missing_revision_manifest={report.missing_revision_manifest} " + f"suspicious_zero_idf={report.suspicious_zero_idf} " + f"would_hit_snapshot_fallback={report.would_hit_snapshot_fallback} " + f"scoring_incomplete={report.scoring_incomplete}" + ) + ready_count = len(reports) - failed + print(f"check namespaces_ready={ready_count}/{len(reports)}") + return 1 if failed else 0 + + def backfill_map_unit_indexes(*, apply: bool, document_id: str = "") -> int: documents = _load_documents(document_id) if not apply: @@ -157,6 +363,12 @@ def backfill_map_unit_indexes(*, apply: bool, document_id: str = "") -> int: def main() -> None: arguments = _build_parser().parse_args() + if arguments.check: + if arguments.apply: + raise SystemExit("use either --check or --apply, not both") + reports = check_fallback_readiness(document_id=str(arguments.document_id)) + raise SystemExit(print_fallback_check(reports)) + count = backfill_map_unit_indexes( apply=bool(arguments.apply), document_id=str(arguments.document_id) ) From d044c58e1075178b1e46eae3677a25a0996482ba Mon Sep 17 00:00:00 2001 From: suguanYang Date: Mon, 31 Aug 2026 20:24:10 +0800 Subject: [PATCH 17/19] fix: resolve PR 358 codeql comments --- .../4a5b6c7d8e9f_add_retrieval_serving_generations.py | 9 +++++++++ .../versions/5b6c7d8e9f0a_add_term_trigram_indexes.py | 9 +++++++++ ...6c7d8e9f0a1b_add_retrieval_namespace_map_snapshots.py | 9 +++++++++ .../test_retrieval_snapshot_large_corpus_contract.py | 5 ++++- 4 files changed, 31 insertions(+), 1 deletion(-) diff --git a/apps/api/alembic/versions/4a5b6c7d8e9f_add_retrieval_serving_generations.py b/apps/api/alembic/versions/4a5b6c7d8e9f_add_retrieval_serving_generations.py index 426ff6e0..5972f228 100644 --- a/apps/api/alembic/versions/4a5b6c7d8e9f_add_retrieval_serving_generations.py +++ b/apps/api/alembic/versions/4a5b6c7d8e9f_add_retrieval_serving_generations.py @@ -13,6 +13,15 @@ branch_labels: Sequence[str] | None = None depends_on: Sequence[str] | None = None +__all__: list[str] = [ + "revision", + "down_revision", + "branch_labels", + "depends_on", + "upgrade", + "downgrade", +] + def upgrade() -> None: if not sa.inspect(op.get_bind()).has_table("retrieval_namespace_generations"): diff --git a/apps/api/alembic/versions/5b6c7d8e9f0a_add_term_trigram_indexes.py b/apps/api/alembic/versions/5b6c7d8e9f0a_add_term_trigram_indexes.py index 71720ff0..50126978 100644 --- a/apps/api/alembic/versions/5b6c7d8e9f0a_add_term_trigram_indexes.py +++ b/apps/api/alembic/versions/5b6c7d8e9f0a_add_term_trigram_indexes.py @@ -12,6 +12,15 @@ branch_labels: Sequence[str] | None = None depends_on: Sequence[str] | None = None +__all__: list[str] = [ + "revision", + "down_revision", + "branch_labels", + "depends_on", + "upgrade", + "downgrade", +] + _MAP_UNIT_INDEX = "idx_document_map_units_term_trgm" _CHUNK_INDEX = "idx_document_chunks_term_trgm" diff --git a/apps/api/alembic/versions/6c7d8e9f0a1b_add_retrieval_namespace_map_snapshots.py b/apps/api/alembic/versions/6c7d8e9f0a1b_add_retrieval_namespace_map_snapshots.py index 1f7b7c01..b561548f 100644 --- a/apps/api/alembic/versions/6c7d8e9f0a1b_add_retrieval_namespace_map_snapshots.py +++ b/apps/api/alembic/versions/6c7d8e9f0a1b_add_retrieval_namespace_map_snapshots.py @@ -13,6 +13,15 @@ branch_labels: Sequence[str] | None = None depends_on: Sequence[str] | None = None +__all__: list[str] = [ + "revision", + "down_revision", + "branch_labels", + "depends_on", + "upgrade", + "downgrade", +] + def upgrade() -> None: if not sa.inspect(op.get_bind()).has_table("retrieval_namespace_map_snapshots"): diff --git a/apps/api/tests/contract/test_retrieval_snapshot_large_corpus_contract.py b/apps/api/tests/contract/test_retrieval_snapshot_large_corpus_contract.py index 0b45a703..f200ff84 100644 --- a/apps/api/tests/contract/test_retrieval_snapshot_large_corpus_contract.py +++ b/apps/api/tests/contract/test_retrieval_snapshot_large_corpus_contract.py @@ -16,7 +16,7 @@ _REVISION_GROUP_SIZE, load_nav_snapshot, ) -import shared.services.retrieval.nav_snapshot as nav_snapshot_module +from shared.services.retrieval import nav_snapshot as nav_snapshot_module from sqlalchemy import Executable, Result, select from sqlalchemy.engine import Row from sqlalchemy.sql.selectable import Select @@ -64,6 +64,9 @@ async def execute(self, statement: Executable) -> Result[tuple[object, ...]]: result = await self._session.execute(statement) return cast(Result[tuple[object, ...]], result) + async def rollback(self) -> None: + await self._session.rollback() + async def _seed_large_retrieval_corpus(namespace: str) -> None: await ContractDatabase.execute( From 0e9ed9aea0af161165efae39d63c7d193d052fd3 Mon Sep 17 00:00:00 2001 From: suguanYang Date: Mon, 31 Aug 2026 20:44:05 +0800 Subject: [PATCH 18/19] fix: tighten retrieval fallback handling --- .../services/retrieval/execution/routes.py | 1 - .../shared/services/retrieval/nav_snapshot.py | 16 ++++++++++++++-- 2 files changed, 14 insertions(+), 3 deletions(-) diff --git a/packages/shared-python/shared/services/retrieval/execution/routes.py b/packages/shared-python/shared/services/retrieval/execution/routes.py index 673e2b6e..29973c6b 100644 --- a/packages/shared-python/shared/services/retrieval/execution/routes.py +++ b/packages/shared-python/shared/services/retrieval/execution/routes.py @@ -56,7 +56,6 @@ async def run_retrieval_route( async def _try_run_small_corpus_route( context: RetrievalRouteContext, ) -> RetrievalRouteOutcome | None: - total_chunk_count: int | None = None total_chunk_count = await count_scoped_chunks( context.db, user_id=context.user_id, diff --git a/packages/shared-python/shared/services/retrieval/nav_snapshot.py b/packages/shared-python/shared/services/retrieval/nav_snapshot.py index 95d24078..ea502c54 100644 --- a/packages/shared-python/shared/services/retrieval/nav_snapshot.py +++ b/packages/shared-python/shared/services/retrieval/nav_snapshot.py @@ -26,6 +26,7 @@ tuple_, ) from sqlalchemy.engine import Result +from sqlalchemy.exc import SQLAlchemyError from shared.models.database.document import ( Document, @@ -393,8 +394,14 @@ async def _resolve_namespace_snapshot_entries( ) try: row = (await db.execute(statement)).first() - except Exception: + except SQLAlchemyError as exc: await db.rollback() + _logger.warning( + "retrieval snapshot namespace lookup failed user_id=%s namespace=%s error=%s", + user_id, + namespace, + exc, + ) return None if row is None: return None @@ -470,8 +477,13 @@ async def _resolve_manifest_entries( ) try: rows = (await db.execute(statement)).all() - except Exception: + except SQLAlchemyError as exc: await db.rollback() + _logger.warning( + "retrieval manifest lookup failed revisions=%s error=%s", + revision_group, + exc, + ) return None if len(rows) != len(revision_group): return None From eeff527de250ea5b2fd3e751e76d14b942352f50 Mon Sep 17 00:00:00 2001 From: suguanYang Date: Mon, 31 Aug 2026 21:57:31 +0800 Subject: [PATCH 19/19] fix: make CodeQL migration exports explicit --- .../versions/4a5b6c7d8e9f_add_retrieval_serving_generations.py | 2 +- .../alembic/versions/5b6c7d8e9f0a_add_term_trigram_indexes.py | 2 +- .../6c7d8e9f0a1b_add_retrieval_namespace_map_snapshots.py | 2 +- 3 files changed, 3 insertions(+), 3 deletions(-) diff --git a/apps/api/alembic/versions/4a5b6c7d8e9f_add_retrieval_serving_generations.py b/apps/api/alembic/versions/4a5b6c7d8e9f_add_retrieval_serving_generations.py index 5972f228..f4fa59f1 100644 --- a/apps/api/alembic/versions/4a5b6c7d8e9f_add_retrieval_serving_generations.py +++ b/apps/api/alembic/versions/4a5b6c7d8e9f_add_retrieval_serving_generations.py @@ -13,7 +13,7 @@ branch_labels: Sequence[str] | None = None depends_on: Sequence[str] | None = None -__all__: list[str] = [ +__all__ = [ "revision", "down_revision", "branch_labels", diff --git a/apps/api/alembic/versions/5b6c7d8e9f0a_add_term_trigram_indexes.py b/apps/api/alembic/versions/5b6c7d8e9f0a_add_term_trigram_indexes.py index 50126978..c49e770d 100644 --- a/apps/api/alembic/versions/5b6c7d8e9f0a_add_term_trigram_indexes.py +++ b/apps/api/alembic/versions/5b6c7d8e9f0a_add_term_trigram_indexes.py @@ -12,7 +12,7 @@ branch_labels: Sequence[str] | None = None depends_on: Sequence[str] | None = None -__all__: list[str] = [ +__all__ = [ "revision", "down_revision", "branch_labels", diff --git a/apps/api/alembic/versions/6c7d8e9f0a1b_add_retrieval_namespace_map_snapshots.py b/apps/api/alembic/versions/6c7d8e9f0a1b_add_retrieval_namespace_map_snapshots.py index b561548f..16380933 100644 --- a/apps/api/alembic/versions/6c7d8e9f0a1b_add_retrieval_namespace_map_snapshots.py +++ b/apps/api/alembic/versions/6c7d8e9f0a1b_add_retrieval_namespace_map_snapshots.py @@ -13,7 +13,7 @@ branch_labels: Sequence[str] | None = None depends_on: Sequence[str] | None = None -__all__: list[str] = [ +__all__ = [ "revision", "down_revision", "branch_labels",