diff --git a/CHANGELOG.md b/CHANGELOG.md index cc000c6df..f4811c552 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -23,6 +23,28 @@ database is skipped with a warning and an inexact total, and a retryable outage or a server too old to attribute its results fails the whole page. +- **#1558**: Vector retrieval reads only the projects in scope and fills the window it + asks for. The sqlite-vec table gains a `project_id` partition key, so a scoped + nearest-neighbour query ranks each project's own vectors instead of taking the k + nearest across the whole database and discarding the out-of-scope ones, which + could leave a small project with an empty page for a query its notes answered. + Existing local storage is carried into the partitioned table without re-embedding. + On Postgres the nearest-neighbour statement now runs on the HNSW index (its + tie-break sort keys had kept the planner on an exact scan of every vector), with + `hnsw.ef_search` sized to the candidate window and an iterative scan that keeps + going until the scope and manifest filters have admitted enough rows. That scan + 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/models/search.py b/src/basic_memory/models/search.py index 70badd098..d192c5250 100644 --- a/src/basic_memory/models/search.py +++ b/src/basic_memory/models/search.py @@ -223,11 +223,18 @@ def create_sqlite_search_vector_embeddings(dimensions: int) -> DDL: - """Build sqlite-vec virtual table DDL for the configured embedding dimension.""" + """Build sqlite-vec virtual table DDL for the configured embedding dimension. + + ``project_id`` is a vec0 partition key: a scoped nearest-neighbour query reads + only the partitions in scope, instead of ranking every project's vectors in the + database and discarding the out-of-scope ones afterwards, which left a small + project with an under-filled window whenever a larger neighbour sat closer. + """ return DDL( f""" CREATE VIRTUAL TABLE IF NOT EXISTS search_vector_embeddings USING vec0( + project_id integer partition key, embedding float[{dimensions}], +source_hash text ) diff --git a/src/basic_memory/repository/pgvector_index.py b/src/basic_memory/repository/pgvector_index.py index 1db0be8c5..41ad25e31 100644 --- a/src/basic_memory/repository/pgvector_index.py +++ b/src/basic_memory/repository/pgvector_index.py @@ -3,6 +3,7 @@ from __future__ import annotations import asyncio +import re from collections.abc import Sequence from loguru import logger @@ -23,6 +24,32 @@ ) +# pgvector's HNSW scan hands back at most ``hnsw.ef_search`` rows, 40 by default, +# whatever the LIMIT asks for; the server caps the setting at 1000. +HNSW_EF_SEARCH_DEFAULT = 40 +HNSW_EF_SEARCH_MAX = 1000 + +_PGVECTOR_VERSION = re.compile(r"^(\d+)\.(\d+)") + + +def pgvector_supports_iterative_scan(extversion: str) -> bool: + """Whether this pgvector keeps scanning an HNSW index until the LIMIT is filled. + + Iterative index scans arrived in pgvector 0.8.0. Before that, a scan stops at + ``hnsw.ef_search`` candidates, and a filter applied afterwards (the manifest + join, a scope narrower than the table) leaves the window under-filled, so the + adapter cannot promise the nearest ``limit`` rows in scope. A version string + the pattern cannot read counts as older. + """ + match = _PGVECTOR_VERSION.match(extversion) + return match is not None and (int(match.group(1)), int(match.group(2))) >= (0, 8) + + +def hnsw_ef_search_for(limit: int) -> int: + """The candidate-list size that lets one HNSW scan return ``limit`` rows.""" + return min(max(limit, HNSW_EF_SEARCH_DEFAULT), HNSW_EF_SEARCH_MAX) + + class PgVectorIndex: """Persist and query semantic vectors in PostgreSQL with pgvector.""" @@ -56,6 +83,19 @@ async def initialize(self) -> None: raise SemanticDependenciesMissingError( "pgvector extension is unavailable for this Postgres database." ) from exc + version = await session.execute( + text("SELECT extversion FROM pg_extension WHERE extname = 'vector'") + ) + extversion = str(version.scalar_one()) + # Trigger: the installed pgvector predates iterative index scans. + # Why: without them a scoped query silently returns fewer rows than + # the window asked for; a deployment gap should read as one. + # Outcome: a typed dependency error the API reports as a bad request. + if not pgvector_supports_iterative_scan(extversion): + raise SemanticDependenciesMissingError( + f"pgvector {extversion} predates iterative index scans; semantic " + "search needs pgvector 0.8 or later (ALTER EXTENSION vector UPDATE)." + ) existing_dimensions = await self._existing_dimensions(session) storage_missing = existing_dimensions is None @@ -322,10 +362,25 @@ async def search( embeddings_in_scope = projects.predicate("e.project_id", params) chunks_in_scope = projects.predicate("c.project_id", params) async with db.scoped_session(self._session_maker) as session: + # Both settings are transaction-local, so they last exactly as long as + # this scoped session. The scan is sized to the window it must fill + # and continues past that until the scope and manifest filters have + # admitted enough rows, instead of stopping at the first ef_search + # candidates and returning whatever of them survived. + await session.execute( + text("SELECT set_config('hnsw.ef_search', :ef_search, true)"), + {"ef_search": str(hnsw_ef_search_for(limit))}, + ) + await session.execute( + text("SELECT set_config('hnsw.iterative_scan', 'relaxed_order', true)") + ) + # A relaxed iterative scan may hand rows back slightly out of distance + # order, so the window is taken by distance alone and sorted once more. result = await session.execute( text( + "WITH nearest AS MATERIALIZED (" "SELECT c.entity_id, c.chunk_key, " - "1 - (e.embedding <=> CAST(:query AS vector)) AS similarity " + "e.embedding <=> CAST(:query AS vector) AS distance " "FROM search_vector_embeddings e " "JOIN search_vector_chunks c ON c.id = e.chunk_id " f"WHERE {embeddings_in_scope} " @@ -335,9 +390,12 @@ async def search( "AND c.embedding_status = 'ready' " "AND c.embedding_model = :embedding_identity " "AND e.source_hash = c.source_hash " - "ORDER BY e.embedding <=> CAST(:query AS vector), " - "c.entity_id ASC, c.chunk_key ASC " + "ORDER BY e.embedding <=> CAST(:query AS vector) " "LIMIT :limit" + ") " + "SELECT entity_id, chunk_key, 1 - distance AS similarity " + "FROM nearest " + "ORDER BY distance ASC, entity_id ASC, chunk_key ASC" ), params, ) 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/src/basic_memory/repository/sqlite_vec_index.py b/src/basic_memory/repository/sqlite_vec_index.py index cfa45ddf7..9f81baf06 100644 --- a/src/basic_memory/repository/sqlite_vec_index.py +++ b/src/basic_memory/repository/sqlite_vec_index.py @@ -115,6 +115,7 @@ async def initialize(self) -> None: expected_dimensions = f"float[{self.scope.dimensions}]" dimensions_changed = bool(vector_sql and expected_dimensions not in vector_sql) source_hash_missing = bool(vector_sql and "+source_hash text" not in vector_sql) + partitions_missing = bool(vector_sql and "partition key" not in vector_sql) if dimensions_changed or source_hash_missing: logger.warning( "SQLite vector storage schema mismatch " @@ -124,6 +125,19 @@ async def initialize(self) -> None: source_hash_missing=source_hash_missing, ) await session.execute(text("DROP TABLE IF EXISTS search_vector_embeddings")) + elif partitions_missing: + # Trigger: storage predates the project_id partition key, and its + # vectors are otherwise current. + # Why: vec0 cannot add a column in place, and re-embedding a whole + # vault only to change how rows are partitioned would cost every + # local user a full embedding pass for nothing new. + # Outcome: the vectors are carried into partitioned storage, each + # keyed by the project its manifest row names; manifests stay ready. + logger.info( + "SQLite vector storage predates project partitions; " + "carrying vectors into partitioned storage" + ) + await self._partition_existing_storage(session) await session.execute(create_sqlite_search_vector_embeddings(self.scope.dimensions)) # Missing or dimension-rebuilt vec storage has no vectors, so ready @@ -138,6 +152,35 @@ async def initialize(self) -> None: await session.commit() self._initialized = True + async def _partition_existing_storage(self, session: AsyncSession) -> None: + """Rebuild vec storage with the partition key, keeping every current vector. + + SQLite DDL is transactional, so the copy out, drop, recreate, and copy back + either all land or none do. A vector whose manifest row is gone has no + project to file under and is left behind, which is what the orphan sweep + would have done to it anyway. + """ + await session.execute( + text( + "CREATE TEMP TABLE search_vector_embeddings_carry AS " + "SELECT e.rowid AS id, c.project_id AS project_id, " + "e.embedding AS embedding, e.source_hash AS source_hash " + "FROM search_vector_embeddings e " + "JOIN search_vector_chunks c ON c.id = e.rowid" + ) + ) + await session.execute(text("DROP TABLE search_vector_embeddings")) + await session.execute(create_sqlite_search_vector_embeddings(self.scope.dimensions)) + await session.execute( + text( + "INSERT INTO search_vector_embeddings " + "(rowid, project_id, embedding, source_hash) " + "SELECT id, project_id, embedding, source_hash " + "FROM search_vector_embeddings_carry" + ) + ) + await session.execute(text("DROP TABLE search_vector_embeddings_carry")) + async def upsert(self, project_id: int, records: Sequence[VectorRecord]) -> None: if not records: return @@ -192,12 +235,14 @@ async def upsert(self, project_id: int, records: Sequence[VectorRecord]) -> None ) await session.execute( text( - "INSERT INTO search_vector_embeddings (rowid, embedding, source_hash) " - "VALUES (:rowid, :embedding, :source_hash)" + "INSERT INTO search_vector_embeddings " + "(rowid, project_id, embedding, source_hash) " + "VALUES (:rowid, :project_id, :embedding, :source_hash)" ), [ { "rowid": rowids_by_key[record.key], + "project_id": project_id, "embedding": json.dumps(record.values), "source_hash": record.source_hash, } @@ -323,21 +368,24 @@ async def search( "embedding_identity": self.scope.embedding_identity, "limit": limit, } - chunks_in_scope = projects.predicate("c.project_id", params) + # vec0 ranks the k nearest within each partition in scope, so a small + # project is never crowded out of its own window by a larger neighbour that + # shares the database; the outer ORDER BY merges the partitions. + partitions_in_scope = projects.predicate("project_id", params) async with db.scoped_session(self._session_maker) as session: await self._ensure_loaded(session) result = await session.execute( text( "WITH vector_matches AS MATERIALIZED (" " SELECT rowid, distance, source_hash FROM search_vector_embeddings " - " WHERE embedding MATCH :query AND k = :vector_k" + f" WHERE {partitions_in_scope} " + " AND embedding MATCH :query AND k = :vector_k" ") " "SELECT c.entity_id, c.chunk_key, vector_matches.distance " "FROM vector_matches " "JOIN search_vector_chunks c ON c.id = vector_matches.rowid " "AND c.source_hash = vector_matches.source_hash " - f"WHERE {chunks_in_scope} " - "AND c.vector_index = 'sqlite-vec' " + "WHERE c.vector_index = 'sqlite-vec' " "AND c.embedding_status = 'ready' " "AND c.embedding_model = :embedding_identity " "ORDER BY vector_matches.distance ASC, " diff --git a/tests/repository/test_pgvector_index.py b/tests/repository/test_pgvector_index.py index c46abca42..ea9b5908e 100644 --- a/tests/repository/test_pgvector_index.py +++ b/tests/repository/test_pgvector_index.py @@ -9,7 +9,11 @@ import pytest from basic_memory.repository import pgvector_index as pgvector_index_module -from basic_memory.repository.pgvector_index import PgVectorIndex +from basic_memory.repository.pgvector_index import ( + PgVectorIndex, + hnsw_ef_search_for, + pgvector_supports_iterative_scan, +) from basic_memory.repository.search_scope import ProjectScope from basic_memory.repository.semantic_errors import SemanticDependenciesMissingError from basic_memory.repository.semantic_vector_index import ( @@ -40,6 +44,10 @@ def fetchone(self) -> object | None: def scalar_one_or_none(self) -> object | None: return self._scalar + def scalar_one(self) -> object: + assert self._scalar is not None + return self._scalar + def mappings(self) -> FakeResult: return self @@ -59,6 +67,7 @@ def __init__( chunk_rows: list[dict[str, object]] | None = None, search_rows: list[dict[str, object]] | None = None, fail_extension: bool = False, + pgvector_version: str = "0.8.0", ) -> None: self.table_exists = table_exists self.dimensions = dimensions @@ -66,6 +75,7 @@ def __init__( self.chunk_rows = chunk_rows or [] self.search_rows = search_rows or [] self.fail_extension = fail_extension + self.pgvector_version = pgvector_version self.calls: list[tuple[str, dict[str, object] | None]] = [] self.commit_count = 0 @@ -78,6 +88,8 @@ async def execute( self.calls.append((sql, params)) if "CREATE EXTENSION" in sql and self.fail_extension: raise RuntimeError("extension unavailable") + if "SELECT extversion" in sql: + return FakeResult(scalar=self.pgvector_version) if "information_schema.tables" in sql: return FakeResult(fetchone=(1,) if self.table_exists else None) if "attname = 'source_hash'" in sql: @@ -330,7 +342,7 @@ async def test_search_returns_normalized_stable_matches(monkeypatch) -> None: (15, 0.0), ] search_call = next(call for call in session.calls if "AS similarity" in call[0]) - assert "c.entity_id ASC, c.chunk_key ASC" in search_call[0] + assert "ORDER BY distance ASC, entity_id ASC, chunk_key ASC" in search_call[0] assert search_call[1] == { "query": "[1,0,0,0]", "scope_0": 7, @@ -356,3 +368,73 @@ async def test_search_binds_every_project_in_scope(monkeypatch) -> None: assert "AND c.project_id IN (:scope_0, :scope_1)" in sql assert params["scope_0"] == 7 assert params["scope_1"] == 9 + + +def _settings(session: FakeSession) -> list[tuple[str, dict[str, object] | None]]: + return [call for call in session.calls if "set_config" in call[0]] + + +@pytest.mark.parametrize( + ("extversion", "expected"), + [("0.7.4", False), ("0.8.0", True), ("0.8.1", True), ("1.0.0", True), ("garbage", False)], +) +def test_iterative_scan_arrived_in_pgvector_0_8(extversion: str, expected: bool) -> None: + assert pgvector_supports_iterative_scan(extversion) is expected + + +def test_ef_search_is_sized_to_the_window_within_the_server_bounds() -> None: + assert hnsw_ef_search_for(5) == 40 + assert hnsw_ef_search_for(40) == 40 + assert hnsw_ef_search_for(250) == 250 + assert hnsw_ef_search_for(5000) == 1000 + + +@pytest.mark.asyncio +async def test_initialize_requires_a_pgvector_that_can_keep_scanning(monkeypatch) -> None: + """An extension too old to fill a filtered window is a deployment error, not a quiet gap.""" + older = FakeSession(pgvector_version="0.7.4") + _install_session(monkeypatch, older) + index = PgVectorIndex(MagicMock(), _scope()) + + with pytest.raises(SemanticDependenciesMissingError, match="pgvector 0.7.4 predates"): + await index.initialize() + + assert not any("CREATE TABLE" in sql for sql in _sql_calls(older)) + assert index._initialized is False + + +@pytest.mark.asyncio +async def test_search_sizes_the_scan_to_the_window_it_must_fill(monkeypatch) -> None: + """A 250-row candidate window asks HNSW for 250 candidates, not the default 40.""" + session = FakeSession(search_rows=[]) + _install_session(monkeypatch, session) + index = PgVectorIndex(MagicMock(), _scope()) + index._initialized = True + + await index.search([1.0, 0.0, 0.0, 0.0], limit=250, projects=PROJECTS) + + ef_search = next(call for call in _settings(session) if "hnsw.ef_search" in call[0]) + assert "set_config('hnsw.ef_search', :ef_search, true)" in ef_search[0] + assert ef_search[1] == {"ef_search": "250"} + search_sql = next(sql for sql, _params in session.calls if "AS similarity" in sql) + # The window is taken by distance inside the CTE, then re-sorted with tie-breaks. + assert "ORDER BY e.embedding <=> CAST(:query AS vector) LIMIT :limit" in search_sql + + +@pytest.mark.asyncio +async def test_search_keeps_scanning_until_the_window_fills(monkeypatch) -> None: + session = FakeSession(search_rows=[]) + _install_session(monkeypatch, session) + index = PgVectorIndex(MagicMock(), _scope()) + index._initialized = True + + await index.search([1.0, 0.0, 0.0, 0.0], limit=5, projects=PROJECTS) + + settings = [sql for sql, _params in _settings(session)] + assert any("'hnsw.ef_search'" in sql for sql in settings) + assert any("'hnsw.iterative_scan', 'relaxed_order', true" in sql for sql in settings) + # Settings precede the scan they configure, inside the same session. + ordered = [sql for sql, _params in session.calls] + assert max(ordered.index(sql) for sql in settings) < next( + position for position, sql in enumerate(ordered) if "AS similarity" in sql + ) diff --git a/tests/repository/test_postgres_search_repository_unit.py b/tests/repository/test_postgres_search_repository_unit.py index 30a468e28..df18a17d8 100644 --- a/tests/repository/test_postgres_search_repository_unit.py +++ b/tests/repository/test_postgres_search_repository_unit.py @@ -225,10 +225,13 @@ async def fake_scoped_session(session_maker): ) missing_table = MagicMock() missing_table.fetchone.return_value = None + pgvector_version = MagicMock() + pgvector_version.scalar_one.return_value = "0.8.0" session.execute.side_effect = [ MagicMock(), MagicMock(), MagicMock(), + pgvector_version, missing_table, MagicMock(), MagicMock(), diff --git a/tests/repository/test_sqlite_vector_search_repository.py b/tests/repository/test_sqlite_vector_search_repository.py index c53291b26..865d3eaed 100644 --- a/tests/repository/test_sqlite_vector_search_repository.py +++ b/tests/repository/test_sqlite_vector_search_repository.py @@ -492,12 +492,12 @@ async def test_sqlite_vec_search_reads_every_project_in_scope(search_repository) ) await session.execute( text( - "INSERT INTO search_vector_embeddings (rowid, embedding, source_hash) " - "VALUES (:rowid, :embedding, 'hash')" + "INSERT INTO search_vector_embeddings (rowid, project_id, embedding, source_hash) " + "VALUES (:rowid, :project_id, :embedding, 'hash')" ), [ - {"rowid": 911, "embedding": "[1,0,0,0]"}, - {"rowid": 912, "embedding": "[0,1,0,0]"}, + {"rowid": 911, "project_id": own_project, "embedding": "[1,0,0,0]"}, + {"rowid": 912, "project_id": other_project, "embedding": "[0,1,0,0]"}, ], ) await session.commit() @@ -514,6 +514,203 @@ async def test_sqlite_vec_search_reads_every_project_in_scope(search_repository) assert nothing == [] +async def _seed_ready_vectors( + search_repository: SQLiteSearchRepository, + index: SQLiteVecIndex, + rows: list[tuple[int, int, str]], + *, + partitioned: bool = True, +) -> None: + """Insert ready manifest rows and their vectors: ``(rowid, project_id, embedding)``.""" + embedding_identity = search_repository._embedding_model_key() + async with db.scoped_session(search_repository.session_maker) as session: + await index._ensure_loaded(session) + await session.execute( + text( + "INSERT INTO search_vector_chunks (" + "id, entity_id, project_id, chunk_key, chunk_text, source_hash, " + "entity_fingerprint, embedding_model, vector_index, embedding_status" + ") VALUES (" + ":id, :id, :project_id, :chunk_key, 'text', 'hash', " + "'fingerprint', :embedding_model, 'sqlite-vec', 'ready')" + ), + [ + { + "id": rowid, + "project_id": project_id, + "chunk_key": f"entity:{rowid}:0", + "embedding_model": embedding_identity, + } + for rowid, project_id, _embedding in rows + ], + ) + if partitioned: + await session.execute( + text( + "INSERT INTO search_vector_embeddings " + "(rowid, project_id, embedding, source_hash) " + "VALUES (:rowid, :project_id, :embedding, 'hash')" + ), + [ + {"rowid": rowid, "project_id": project_id, "embedding": embedding} + for rowid, project_id, embedding in rows + ], + ) + else: + await session.execute( + text( + "INSERT INTO search_vector_embeddings (rowid, embedding, source_hash) " + "VALUES (:rowid, :embedding, 'hash')" + ), + [{"rowid": rowid, "embedding": embedding} for rowid, _project, embedding in rows], + ) + await session.commit() + + +@pytest.mark.asyncio +async def test_sqlite_vec_scope_is_a_partition_not_a_filter_on_the_nearest(search_repository): + """A small project fills its window even when a neighbour's vectors sit closer. + + The k nearest across the whole database used to be taken first and the scope + applied afterwards, so a project holding a few vectors among a large + neighbour's could get an empty page for a query its own notes answered. + """ + 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) + small = search_repository.project_id + large = small + 1 + + # The large project's vectors are all nearer the query than the small one's. + await _seed_ready_vectors( + search_repository, + index, + [(921, large, "[1,0,0,0]"), (922, large, "[0.9,0.1,0,0]"), (923, large, "[0.8,0.2,0,0]")] + + [(931, small, "[0,1,0,0]"), (932, small, "[0,0,1,0]")], + ) + + nearest_two = await index.search( + [1.0, 0.0, 0.0, 0.0], limit=2, projects=ProjectScope.single(small) + ) + + 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.""" + if not isinstance(search_repository, SQLiteSearchRepository): + pytest.skip("sqlite-vec storage upgrade 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 + dimensions = search_repository._vector_dimensions + + async with db.scoped_session(search_repository.session_maker) as session: + await index._ensure_loaded(session) + await session.execute(text("DROP TABLE search_vector_embeddings")) + await session.execute( + text( + "CREATE VIRTUAL TABLE search_vector_embeddings USING vec0(" + f"embedding float[{dimensions}], +source_hash text)" + ) + ) + await session.commit() + await _seed_ready_vectors( + search_repository, + index, + [(941, project, "[1,0,0,0]"), (942, project, "[0,1,0,0]")], + partitioned=False, + ) + + index.invalidate_initialization() + await index.initialize() + + async with db.scoped_session(search_repository.session_maker) as session: + table_sql = await session.scalar( + text("SELECT sql FROM sqlite_master WHERE name = 'search_vector_embeddings'") + ) + statuses = await session.execute( + text("SELECT embedding_status FROM search_vector_chunks WHERE id IN (941, 942)") + ) + carried = await session.execute( + text("SELECT rowid, project_id FROM search_vector_embeddings ORDER BY rowid") + ) + assert table_sql is not None and "project_id integer partition key" in table_sql + assert statuses.scalars().all() == ["ready", "ready"] + assert carried.all() == [(941, project), (942, project)] + found = await index.search([1.0, 0.0, 0.0, 0.0], limit=5, projects=ProjectScope.single(project)) + assert [match.key.entity_id for match in found] == [941, 942] + + @pytest.mark.asyncio async def test_sqlite_vec_delete_requires_pending_source_generation(search_repository): """A stale delete cannot remove a same-source vector that is already ready.""" 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))]