From e9c4bff93ce0f52f5dd84d4040b2a2faf5f1e0ce Mon Sep 17 00:00:00 2001 From: phernandez Date: Tue, 15 Sep 2026 17:53:56 -0500 Subject: [PATCH] perf(core): read only the projects in scope from the vector indexes Both built-in vector adapters took the k nearest vectors across the whole database and applied the scope afterwards. In a shared database a small project could get an under-filled or empty page for a query its own notes answered, because a larger neighbour's vectors filled the window first. On a 20k-chunk corpus where one project holds 2% of the vectors, a limit-100 query for that project returned 1 row from sqlite-vec. sqlite-vec: `project_id` becomes a vec0 partition key, so the KNN ranks each partition in scope and the outer ORDER BY merges them (100 of 100 rows, 0.2 ms against 3.1 ms). Existing local storage is carried into the partitioned table through a temporary table inside one transaction; vectors and ready manifests survive, nothing is re-embedded. Dimension and source_hash mismatches still rebuild as before. pgvector: the statement's tie-break sort keys had kept the planner off the HNSW index entirely, so every vector query was an exact scan of the table. The window is now taken by distance alone inside a materialized CTE and re-sorted with the tie-breaks outside it. `hnsw.ef_search` is set per transaction to the window the query must fill (the default 40 capped every larger candidate pool), and `hnsw.iterative_scan = relaxed_order` lets the scan continue until the scope and manifest filters have admitted enough rows. That scan needs pgvector 0.8; an older extension now raises SemanticDependenciesMissingError at first use rather than returning short windows. Neon tenants run 0.8.0. Refs #1558 Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_019YW9ysxugGGBCNEGzsxtFV Signed-off-by: phernandez --- CHANGELOG.md | 13 ++ src/basic_memory/models/search.py | 9 +- src/basic_memory/repository/pgvector_index.py | 64 +++++++- .../repository/sqlite_vec_index.py | 60 +++++++- tests/repository/test_pgvector_index.py | 86 ++++++++++- .../test_postgres_search_repository_unit.py | 3 + .../test_sqlite_vector_search_repository.py | 142 +++++++++++++++++- 7 files changed, 361 insertions(+), 16 deletions(-) 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."""