diff --git a/CHANGELOG.md b/CHANGELOG.md index 6de2d3fc1..f4811c552 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -36,6 +36,15 @@ needs pgvector 0.8 or later; an older extension is reported as a dependency error instead of quietly returning short windows. +- **#1558**: A vector or hybrid search with structured filters (note types, dates, + categories, metadata, path prefixes, valid time) now fills its candidate window. + The vector index ranks by similarity alone, so a window taken straight from it and + filtered afterwards could hold few admitted rows while more sat just past it, and + the page came back short although matches existed. The reader re-reads the window + with a bounded geometric overfetch until it holds enough admitted rows, the ranking + is exhausted, or its tail falls below the similarity threshold. Unfiltered searches + read their window once, as before. + - **#1512**: Word, PowerPoint, and CSV files get the same sidecar Markdown note a PDF gets. `bm import document ` indexes the project, extracts the file, and writes `..md` next to it plus a run note under diff --git a/src/basic_memory/repository/search_reader.py b/src/basic_memory/repository/search_reader.py index 8d5d80b28..ed2a31fa8 100644 --- a/src/basic_memory/repository/search_reader.py +++ b/src/basic_memory/repository/search_reader.py @@ -182,6 +182,27 @@ class HydratedChunk: similarity: float +@dataclass(frozen=True, slots=True) +class CandidateWindow: + """The search rows one vector candidate window resolved to, in adapter order. + + ``similarity_by_key`` holds the best chunk similarity of every row above the + threshold; ``rows`` holds those of them that exist and that the query's filters + admit, so a key present in the first and absent from the second was rejected. + """ + + similarity_by_key: dict[SearchIndexKey, float] + chunks_by_key: dict[SearchIndexKey, list[tuple[float, str]]] + rows: dict[SearchIndexKey, SearchIndexRow] + chunk_count: int + vector_query_ms: float = 0.0 + hydrate_ms: float = 0.0 + + @property + def admitted(self) -> int: + return sum(1 for key in self.similarity_by_key if key in self.rows) + + # --- Vector and hybrid retrieval --- @@ -626,33 +647,26 @@ async def vector_only( with logfire.span("search.embed_query", query_chars=len(query_text)): query_embedding = await self.vector.embedding_provider.embed_query(query_text) embed_ms = (time.perf_counter() - embed_start) * 1000 - vector_query_start = time.perf_counter() - - # Constraint: vector adapters may open their own session, while the SQLite - # test/runtime pool can contain only one connection. A plain AsyncSession - # defers checkout until hydration runs after adapter search has released it. - async with self.session_maker() as session: - vector_rows = await self._run_vector_query( - session, - query_embedding, - candidate_limit, - trace=trace, - ) - vector_query_ms = (time.perf_counter() - vector_query_start) * 1000 - vector_row_count = len(vector_rows) - hydrate_ms = 0.0 + # Per-query min_similarity overrides the configured default. + effective_min_similarity = ( + query.min_similarity if query.min_similarity is not None else self.vector.min_similarity + ) + + window = await self._candidate_window( + query, + query_embedding, + candidate_limit, + min_similarity=effective_min_similarity, + trace=trace, + ) if trace is not None: trace.vector = build_vector_stage( previous=trace.vector, - effective_min_similarity=( - query.min_similarity - if query.min_similarity is not None - else self.vector.min_similarity - ), + effective_min_similarity=effective_min_similarity, min_similarity_source=("query" if query.min_similarity is not None else "config"), embed_ms=embed_ms, - vector_query_ms=vector_query_ms, + vector_query_ms=window.vector_query_ms, ) def _log_vector_summary() -> None: @@ -671,92 +685,20 @@ def _log_vector_summary() -> None: retrieval_mode="vector", query_length=len(query_text), candidate_limit=candidate_limit, - vector_row_count=vector_row_count, + vector_row_count=window.chunk_count, embed_ms=embed_ms, - vector_query_ms=vector_query_ms, - hydrate_ms=hydrate_ms, + vector_query_ms=window.vector_query_ms, + hydrate_ms=window.hydrate_ms, total_ms=total_ms, ) - if not vector_rows: - _log_vector_summary() - return [] - - hydrate_start = time.perf_counter() - # Build per-search_index_row similarity scores from chunk-level results. - # Each chunk_key encodes the search_index row type and id; keep both as the - # key because different row types can share the same numeric id (#982). - # Track the best similarity per row (for ranking) and all chunks (for context). - similarity_by_si_key: dict[SearchIndexKey, float] = {} - chunks_by_si_key: dict[SearchIndexKey, list[tuple[float, str]]] = {} - for chunk in vector_rows: - try: - si_key = parse_chunk_key(chunk.chunk_key) - except (ValueError, IndexError): - # A chunk without a parseable key names no search row to rank. - continue - current = similarity_by_si_key.get(si_key) - if current is None or chunk.similarity > current: - similarity_by_si_key[si_key] = chunk.similarity - chunks_by_si_key.setdefault(si_key, []).append((chunk.similarity, chunk.chunk_text)) - - if not similarity_by_si_key: - hydrate_ms = (time.perf_counter() - hydrate_start) * 1000 + if not window.similarity_by_key: _log_vector_summary() return [] - # Filter out results below the minimum similarity threshold. - # Per-query min_similarity overrides the configured default. - effective_min_similarity = ( - query.min_similarity if query.min_similarity is not None else self.vector.min_similarity - ) - if effective_min_similarity > 0.0: - if trace is not None: - threshold_rejections = tuple( - BelowThreshold(key=key, similarity=value, threshold=effective_min_similarity) - for key, value in similarity_by_si_key.items() - if value < effective_min_similarity - ) - trace.vector = build_vector_stage( - previous=trace.vector, - threshold_rejections=threshold_rejections, - ) - similarity_by_si_key = { - k: v for k, v in similarity_by_si_key.items() if v >= effective_min_similarity - } - if not similarity_by_si_key: - hydrate_ms = (time.perf_counter() - hydrate_start) * 1000 - _log_vector_summary() - return [] - - # Fetch the actual search_index rows. Colliding (type, id) keys share one - # bare id, so deduplicate while preserving first-seen order. - si_ids = list(dict.fromkeys(si_id for _, si_id in similarity_by_si_key)) - search_index_rows = await self._fetch_search_index_rows_by_ids(si_ids) - if trace is not None: - trace.vector = build_vector_stage( - previous=trace.vector, - missing_search_rows=tuple( - MissingSearchRow(key=key) - for key in similarity_by_si_key - if key not in search_index_rows - ), - ) - - if query.has_filters: - allowed_keys = await self._filter_candidate_keys(list(search_index_rows), query) - if trace is not None: - trace.vector = build_vector_stage( - previous=trace.vector, - filter_rejections=tuple( - FilteredOut(key=key) for key in search_index_rows if key not in allowed_keys - ), - ) - search_index_rows = {k: v for k, v in search_index_rows.items() if k in allowed_keys} - ranked_rows: list[SearchIndexRow] = [] - for si_key, similarity in similarity_by_si_key.items(): - row = search_index_rows.get(si_key) + for si_key, similarity in window.similarity_by_key.items(): + row = window.rows.get(si_key) if row is None: continue @@ -766,7 +708,7 @@ def _log_vector_summary() -> None: if content_snippet and len(content_snippet) <= SMALL_NOTE_CONTENT_LIMIT: matched_chunk_text = content_snippet else: - si_chunks = chunks_by_si_key.get(si_key, []) + si_chunks = window.chunks_by_key.get(si_key, []) si_chunks.sort(key=lambda c: c[0], reverse=True) top_texts = [chunk_text for _, chunk_text in si_chunks[:TOP_CHUNKS_PER_RESULT]] matched_chunk_text = "\n---\n".join(top_texts) if top_texts else None @@ -780,7 +722,6 @@ def _log_vector_summary() -> None: ) ranked_rows.sort(key=lambda item: item.score or 0.0, reverse=True) - hydrate_ms = (time.perf_counter() - hydrate_start) * 1000 # Rerank over the wide candidate pool, then slice to the page. Suppressed when # hybrid calls this internally (apply_rerank=False): hybrid reranks its own # fused result, and _rerank_and_paginate is a plain slice without a reranker. @@ -817,6 +758,128 @@ def _log_vector_summary() -> None: _log_vector_summary() return output + async def _candidate_window( + self, + query: PreparedSearchQuery, + query_embedding: list[float], + candidate_limit: int, + *, + min_similarity: float, + trace: SearchTraceCollector | None = None, + ) -> CandidateWindow: + """Resolve the nearest chunks to admitted search rows, widening past rejections. + + Trigger: the query carries structured filters the adapter cannot evaluate. + Why: the adapter ranks by similarity alone, so a window taken straight from + its ranking can hold few admitted rows while more sit just past it, and a + page built from that window comes up short although matches exist. + Outcome: the window is re-read with a bounded geometric overfetch until it + holds ``candidate_limit`` admitted rows, the ranking is exhausted, or its + tail has fallen below the similarity threshold, past which nothing further + can qualify. A query without filters resolves its window once. + """ + scan_limit = candidate_limit + scanned = -1 + vector_query_ms = 0.0 + hydrate_ms = 0.0 + while True: + vector_query_start = time.perf_counter() + # Constraint: vector adapters may open their own session, while the SQLite + # test/runtime pool can contain only one connection. A plain AsyncSession + # defers checkout until hydration runs after adapter search has released it. + async with self.session_maker() as session: + chunks = await self._run_vector_query( + session, query_embedding, scan_limit, trace=trace + ) + vector_query_ms += (time.perf_counter() - vector_query_start) * 1000 + hydrate_start = time.perf_counter() + window = await self._resolve_rows(chunks, query, min_similarity, trace=trace) + hydrate_ms += (time.perf_counter() - hydrate_start) * 1000 + + exhausted = ( + len(chunks) < scan_limit + or len(chunks) <= scanned + or scan_limit >= VECTOR_FILTER_SCAN_LIMIT + ) + tail_below_threshold = bool(chunks) and chunks[-1].similarity < min_similarity + if ( + not query.has_filters + or window.admitted >= candidate_limit + or exhausted + or tail_below_threshold + ): + return replace(window, vector_query_ms=vector_query_ms, hydrate_ms=hydrate_ms) + scanned = len(chunks) + scan_limit = min(scan_limit * 2, VECTOR_FILTER_SCAN_LIMIT) + + async def _resolve_rows( + self, + chunks: list[HydratedChunk], + query: PreparedSearchQuery, + min_similarity: float, + *, + trace: SearchTraceCollector | None = None, + ) -> CandidateWindow: + """Turn ranked chunks into the search rows above threshold that the filters admit.""" + # Build per-search_index_row similarity scores from chunk-level results. + # Each chunk_key encodes the search_index row type and id; keep both as the + # key because different row types can share the same numeric id (#982). + # Track the best similarity per row (for ranking) and all chunks (for context). + similarity_by_key: dict[SearchIndexKey, float] = {} + chunks_by_key: dict[SearchIndexKey, list[tuple[float, str]]] = {} + for chunk in chunks: + try: + si_key = parse_chunk_key(chunk.chunk_key) + except (ValueError, IndexError): + # A chunk without a parseable key names no search row to rank. + continue + current = similarity_by_key.get(si_key) + if current is None or chunk.similarity > current: + similarity_by_key[si_key] = chunk.similarity + chunks_by_key.setdefault(si_key, []).append((chunk.similarity, chunk.chunk_text)) + + # Filter out results below the minimum similarity threshold. + if min_similarity > 0.0: + if trace is not None: + threshold_rejections = tuple( + BelowThreshold(key=key, similarity=value, threshold=min_similarity) + for key, value in similarity_by_key.items() + if value < min_similarity + ) + trace.vector = build_vector_stage( + previous=trace.vector, + threshold_rejections=threshold_rejections, + ) + similarity_by_key = {k: v for k, v in similarity_by_key.items() if v >= min_similarity} + if not similarity_by_key: + return CandidateWindow(similarity_by_key, chunks_by_key, {}, len(chunks)) + + # Fetch the actual search_index rows. Colliding (type, id) keys share one + # bare id, so deduplicate while preserving first-seen order. + si_ids = list(dict.fromkeys(si_id for _, si_id in similarity_by_key)) + search_index_rows = await self._fetch_search_index_rows_by_ids(si_ids) + if trace is not None: + trace.vector = build_vector_stage( + previous=trace.vector, + missing_search_rows=tuple( + MissingSearchRow(key=key) + for key in similarity_by_key + if key not in search_index_rows + ), + ) + + if query.has_filters: + allowed_keys = await self._filter_candidate_keys(list(search_index_rows), query) + if trace is not None: + trace.vector = build_vector_stage( + previous=trace.vector, + filter_rejections=tuple( + FilteredOut(key=key) for key in search_index_rows if key not in allowed_keys + ), + ) + search_index_rows = {k: v for k, v in search_index_rows.items() if k in allowed_keys} + return CandidateWindow(similarity_by_key, chunks_by_key, search_index_rows, len(chunks)) + # --- Hybrid score-based fusion --- async def hybrid( diff --git a/tests/repository/test_sqlite_vector_search_repository.py b/tests/repository/test_sqlite_vector_search_repository.py index 9c0ec09b9..865d3eaed 100644 --- a/tests/repository/test_sqlite_vector_search_repository.py +++ b/tests/repository/test_sqlite_vector_search_repository.py @@ -599,6 +599,69 @@ async def test_sqlite_vec_scope_is_a_partition_not_a_filter_on_the_nearest(searc assert [match.key.entity_id for match in nearest_two] == [931, 932] +@pytest.mark.asyncio +async def test_filtered_vector_search_fills_its_window_past_nearer_rejected_rows( + search_repository, +): + """Rows the filter rejects can sit in front of the ones it admits without hiding them. + + Twelve archive rows are nearer the query than three notes rows. With a candidate + window of ten chunks, a filter on the notes prefix used to see ten rejected + candidates and answer with nothing, although three notes matched. + """ + if not isinstance(search_repository, SQLiteSearchRepository): + pytest.skip("sqlite-vec search behavior is local SQLite-only.") + + _enable_semantic(search_repository) + await search_repository.init_search_index() + index = cast(SQLiteVecIndex, search_repository._semantic_vector_index) + project = search_repository.project_id + # limit 1 with vector_k 4 sizes the first window at ten chunks. + search_repository._semantic_vector_k = 4 + search_repository._semantic_min_similarity = 0.0 + + nearer = list(range(1101, 1113)) + farther = [1121, 1122, 1123] + for row_id in nearer: + await search_repository.index_item( + _entity_row( + project_id=project, + row_id=row_id, + entity_id=row_id, + title=f"Archive {row_id}", + permalink=f"archive/entry-{row_id}", + content_stems="auth token archive", + ) + ) + for row_id in farther: + await search_repository.index_item( + _entity_row( + project_id=project, + row_id=row_id, + entity_id=row_id, + title=f"Note {row_id}", + permalink=f"notes/entry-{row_id}", + content_stems="schema note", + ) + ) + # The query embeds to [1,0,0,0]; archive rows sit nearer it than notes rows. + await _seed_ready_vectors( + search_repository, + index, + [(row_id, project, f"[1,{(row_id - 1100) / 100:.2f},0,0]") for row_id in nearer] + + [(row_id, project, f"[0.5,1,{(row_id - 1120) / 100:.2f},0]") for row_id in farther], + ) + + results = await search_repository.search( + search_text="auth", + file_path_prefix="notes", + retrieval_mode=SearchRetrievalMode.VECTOR, + limit=1, + ) + + assert [row.id for row in results] == [1121] + + @pytest.mark.asyncio async def test_sqlite_vec_partitions_legacy_storage_without_re_embedding(search_repository): """Storage from before the partition key is carried over, vectors and readiness intact.""" diff --git a/tests/repository/test_vector_filter_window.py b/tests/repository/test_vector_filter_window.py new file mode 100644 index 000000000..5219b27a5 --- /dev/null +++ b/tests/repository/test_vector_filter_window.py @@ -0,0 +1,216 @@ +"""A filtered vector search fills its candidate window instead of stopping short. + +The adapter ranks by similarity alone and knows nothing of structured filters. A window +taken straight from that ranking and filtered afterwards can hold few admitted rows while +more sit just past it, so a page came back short although matches existed. The reader now +re-reads the window with a bounded geometric overfetch until it holds enough admitted rows, +the ranking is exhausted, or its tail has fallen below the similarity threshold. +""" + +from collections.abc import Sequence +from typing import Any, cast, override +from unittest.mock import AsyncMock, patch + +import pytest +from sqlalchemy.ext.asyncio import AsyncSession + +from basic_memory.repository.search_index_row import SearchIndexKey, SearchIndexRow +from basic_memory.repository.search_query import PreparedSearchQuery +from basic_memory.repository.search_reader import ( + VECTOR_FILTER_SCAN_LIMIT, + HydratedChunk, + SemanticSearch, +) +from basic_memory.repository.search_scope import ProjectScope +from basic_memory.repository.search_trace import SearchTraceCollector +from basic_memory.schemas.search import SearchRetrievalMode +from tests.repository.test_hybrid_fusion import FakeFts, FakeRow, fake_vector_retrieval +from tests.repository.test_vector_threshold import fake_session_maker + +FILTERED = PreparedSearchQuery( + search_text="test", retrieval_mode=SearchRetrievalMode.VECTOR, file_path_prefix="notes" +) +UNFILTERED = PreparedSearchQuery(search_text="test", retrieval_mode=SearchRetrievalMode.VECTOR) + + +class AdmittingFts(FakeFts): + """Answers a filter pass with exactly the candidates in ``admitted``.""" + + def __init__(self, admitted: set[int]) -> None: + super().__init__() + self.admitted = admitted + + @override + async def search( + self, + scope: ProjectScope, + query: PreparedSearchQuery, + *, + limit: int, + offset: int, + allow_relaxed: bool = False, + session: AsyncSession | None = None, + candidate_keys: Sequence[SearchIndexKey] | None = None, + trace: SearchTraceCollector | None = None, + ) -> list[SearchIndexRow]: + self.calls.append({"candidate_keys": list(candidate_keys or [])}) + return cast( + list[SearchIndexRow], + [ + FakeRow(id=row_id) + for _type, row_id in candidate_keys or [] + if row_id in self.admitted + ], + ) + + +def _ranking(count: int, *, top: float = 0.99, step: float = 0.01) -> list[HydratedChunk]: + """Rows 1..count, nearest first, similarity falling by ``step`` per rank.""" + return [ + HydratedChunk( + entity_id=row_id, + chunk_key=f"entity:{row_id}:0", + chunk_text=f"chunk {row_id}", + similarity=top - (row_id - 1) * step, + ) + for row_id in range(1, count + 1) + ] + + +def _adapter(ranking: list[HydratedChunk]) -> AsyncMock: + """A neighbour stage that hands back the top ``candidate_limit`` of a fixed ranking.""" + + async def run(session: Any, embedding: Any, candidate_limit: int, *, trace: Any = None): + return ranking[:candidate_limit] + + return AsyncMock(side_effect=run) + + +def _rows() -> AsyncMock: + async def fetch(row_ids: list[int]) -> dict[SearchIndexKey, Any]: + return { + ("entity", row_id): FakeRow(id=row_id, file_path=f"row-{row_id}.md") + for row_id in row_ids + } + + return AsyncMock(side_effect=fetch) + + +def _semantic(fts: FakeFts, *, vector_k: int = 4, min_similarity: float = 0.0) -> SemanticSearch: + return SemanticSearch( + fake_session_maker(), + ProjectScope.single(1), + fts, + fake_vector_retrieval(vector_k=vector_k, min_similarity=min_similarity), + ) + + +async def _search( + semantic: SemanticSearch, + adapter: AsyncMock, + *, + query: PreparedSearchQuery, + limit: int = 2, +) -> list[int]: + with ( + patch.object(semantic, "_run_vector_query", adapter), + patch.object(semantic, "_fetch_search_index_rows_by_ids", _rows()), + ): + rows = await semantic.vector_only(query, limit=limit, offset=0) + return [row.id for row in rows if row.id is not None] + + +def _windows(adapter: AsyncMock) -> list[int]: + return [call.args[2] for call in adapter.await_args_list] + + +@pytest.mark.asyncio +async def test_the_window_widens_until_the_filter_has_admitted_enough_rows(): + """Twenty rejected neighbours in front of five admitted ones still yield a full page.""" + fts = AdmittingFts(admitted={21, 22, 23, 24, 25}) + adapter = _adapter(_ranking(25)) + + # limit 2 with vector_k 4 sizes the window at 20 chunks, all of them rejected. + found = await _search(_semantic(fts), adapter, query=FILTERED) + + assert found == [21, 22] + assert _windows(adapter) == [20, 40] + + +@pytest.mark.asyncio +async def test_a_query_without_filters_reads_its_window_once(): + fts = AdmittingFts(admitted=set()) + adapter = _adapter(_ranking(25)) + + found = await _search(_semantic(fts), adapter, query=UNFILTERED) + + assert found == [1, 2] + assert _windows(adapter) == [20] + assert fts.calls == [] + + +@pytest.mark.asyncio +async def test_an_exhausted_ranking_ends_the_search_with_what_it_admitted(): + """When the adapter has nothing past the window, a short answer is the true answer.""" + fts = AdmittingFts(admitted={5}) + adapter = _adapter(_ranking(12)) + + found = await _search(_semantic(fts), adapter, query=FILTERED) + + assert found == [5] + # The first read came back short of its own window: nothing further exists. + assert _windows(adapter) == [20] + + +@pytest.mark.asyncio +async def test_a_ranking_that_stops_growing_is_exhausted(): + """An adapter capped below the widened window cannot be asked forever.""" + fts = AdmittingFts(admitted=set()) + adapter = _adapter(_ranking(20)) + + found = await _search(_semantic(fts), adapter, query=FILTERED) + + assert found == [] + assert _windows(adapter) == [20, 40] + + +@pytest.mark.asyncio +async def test_a_tail_below_the_threshold_ends_the_widening(): + """Nothing past a sub-threshold tail can qualify, whatever the filter would admit.""" + fts = AdmittingFts(admitted={30}) + # Similarities fall from 0.99 by 0.03 per rank: rank 20 sits at 0.42, under 0.5. + adapter = _adapter(_ranking(40, step=0.03)) + + found = await _search(_semantic(fts, min_similarity=0.5), adapter, query=FILTERED) + + assert found == [] + assert _windows(adapter) == [20] + + +@pytest.mark.asyncio +async def test_the_widening_is_bounded_by_the_scan_limit(): + fts = AdmittingFts(admitted=set()) + adapter = _adapter(_ranking(VECTOR_FILTER_SCAN_LIMIT + 10, step=0.0)) + + found = await _search(_semantic(fts), adapter, query=FILTERED) + + assert found == [] + windows = _windows(adapter) + assert windows[0] == 20 + assert windows[-1] == VECTOR_FILTER_SCAN_LIMIT + assert all( + later == min(earlier * 2, VECTOR_FILTER_SCAN_LIMIT) + for earlier, later in zip(windows, windows[1:]) + ) + + +@pytest.mark.asyncio +async def test_each_round_asks_the_filter_only_about_rows_that_exist(): + """The filter pass is bounded by the candidates, never a page of the whole match set.""" + fts = AdmittingFts(admitted={21, 22}) + adapter = _adapter(_ranking(25)) + + await _search(_semantic(fts), adapter, query=FILTERED) + + asked = [sorted(row_id for _type, row_id in call["candidate_keys"]) for call in fts.calls] + assert asked == [list(range(1, 21)), list(range(1, 26))]