diff --git a/CHANGELOG.md b/CHANGELOG.md index cc000c6df..6de2d3fc1 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -23,6 +23,19 @@ 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. + - **#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/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..9c0ec09b9 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,140 @@ 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_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."""