From 0470b7d9db269c53bbbd8e21144192673da5af09 Mon Sep 17 00:00:00 2001 From: phernandez Date: Wed, 22 Jul 2026 00:30:59 -0500 Subject: [PATCH 01/24] feat(core): add pluggable semantic vector indexes Signed-off-by: phernandez --- docs/semantic-search.md | 86 ++- ...0l1m2n3_add_vector_index_manifest_state.py | 78 +++ src/basic_memory/config_models.py | 9 + src/basic_memory/models/search.py | 4 + src/basic_memory/repository/pgvector_index.py | 299 +++++++++++ .../repository/postgres_search_repository.py | 258 ++------- .../repository/search_repository.py | 26 + .../repository/search_repository_base.py | 360 ++++++++++++- .../repository/semantic_errors.py | 4 + .../repository/semantic_vector_index.py | 115 ++++ .../semantic_vector_index_factory.py | 136 +++++ .../repository/semantic_vector_sync.py | 55 +- .../repository/sqlite_search_repository.py | 206 ++------ .../repository/sqlite_vec_index.py | 286 ++++++++++ src/basic_memory/services/project_service.py | 173 +++---- src/basic_memory/services/search_service.py | 51 +- test-int/test_embedding_status_vec0.py | 60 ++- tests/repository/test_pgvector_index.py | 253 +++++++++ .../test_postgres_search_repository_unit.py | 53 +- .../repository/test_semantic_vector_index.py | 296 +++++++++++ tests/repository/test_semantic_vector_sync.py | 6 +- .../test_sqlite_vector_search_repository.py | 490 +++++++++++++++++- tests/services/test_project_removal_bug.py | 10 +- .../test_project_service_embedding_status.py | 215 ++++---- tests/services/test_semantic_search.py | 21 +- tests/test_semantic_vector_index_migration.py | 100 ++++ 26 files changed, 2870 insertions(+), 780 deletions(-) create mode 100644 src/basic_memory/alembic/versions/o8j9k0l1m2n3_add_vector_index_manifest_state.py create mode 100644 src/basic_memory/repository/pgvector_index.py create mode 100644 src/basic_memory/repository/semantic_vector_index.py create mode 100644 src/basic_memory/repository/semantic_vector_index_factory.py create mode 100644 src/basic_memory/repository/sqlite_vec_index.py create mode 100644 tests/repository/test_pgvector_index.py create mode 100644 tests/repository/test_semantic_vector_index.py create mode 100644 tests/test_semantic_vector_index_migration.py diff --git a/docs/semantic-search.md b/docs/semantic-search.md index fb54f4d37..aea71af00 100644 --- a/docs/semantic-search.md +++ b/docs/semantic-search.md @@ -99,6 +99,7 @@ All settings are fields on `BasicMemoryConfig` and can be set via environment va | Config Field | Env Var | Default | Description | |---|---|---|---| | `semantic_search_enabled` | `BASIC_MEMORY_SEMANTIC_SEARCH_ENABLED` | Auto (`true` when semantic deps are available) | Enable semantic search. Required before vector/hybrid modes work. | +| `semantic_vector_index` | `BASIC_MEMORY_SEMANTIC_VECTOR_INDEX` | `"pgvector"` | Postgres vector storage adapter. `"pgvector"` is built in; other names resolve through the `basic_memory.semantic_vector_indexes` Python entry-point group. SQLite always uses its built-in `sqlite-vec` adapter. | | `semantic_embedding_provider` | `BASIC_MEMORY_SEMANTIC_EMBEDDING_PROVIDER` | `"fastembed"` | Embedding provider: `"fastembed"` (local), `"openai"` (API), or `"litellm"` (multi-provider API, **experimental** — advanced users only). | | `semantic_embedding_model` | `BASIC_MEMORY_SEMANTIC_EMBEDDING_MODEL` | `"bge-small-en-v1.5"` | Model identifier. Auto-adjusted per provider if left at default. | | `semantic_embedding_api_base` | `BASIC_MEMORY_SEMANTIC_EMBEDDING_API_BASE` | Unset | Optional custom endpoint for the LiteLLM provider, including local or self-hosted OpenAI-compatible servers. | @@ -349,6 +350,7 @@ bm reindex -p my-project - **Dimension change**: After changing `semantic_embedding_dimensions` - **LiteLLM role change**: After changing `semantic_embedding_document_input_type` or `semantic_embedding_query_input_type` - **Literal prefix change**: After changing `semantic_embedding_document_prefix` or `semantic_embedding_query_prefix` +- **Vector index change**: After changing `semantic_vector_index` (the normal incremental sync also detects the change, while an explicit reindex completes the migration immediately) The reindex command shows progress with embedded/skipped/error counts: @@ -405,10 +407,92 @@ The sqlite-vec extension is loaded per-connection. Vector tables are created laz ### Postgres (cloud) -- **Vector storage**: [pgvector](https://github.com/pgvector/pgvector) with HNSW indexing +- **Default vector storage**: [pgvector](https://github.com/pgvector/pgvector) with HNSW indexing - **Local Docker**: use `docker-compose-postgres.yml` (`pgvector/pgvector:pg17`). Plain `postgres:17` lacks the extension; run `CREATE EXTENSION IF NOT EXISTS vector;` on any external instance before first migration. - **Chunk metadata table**: Created via Alembic migration (`search_vector_chunks` with `BIGSERIAL` primary key) - **Embedding table**: `search_vector_embeddings` created at runtime (dimension-dependent, same pattern as SQLite) - **Index**: HNSW index on the embedding column for fast approximate nearest-neighbour queries The Alembic migration creates the dimension-independent chunks table. The embeddings table and HNSW index are deferred to runtime because they depend on the configured vector dimensions. + +## Pluggable Vector Indexes + +Postgres deployments can replace pgvector storage and nearest-neighbour lookup without +replacing Basic Memory's SQL repositories or embedding providers: + +```bash +export BASIC_MEMORY_SEMANTIC_VECTOR_INDEX=milvus +``` + +The named extension must be installed in the same Python environment as Basic Memory. A +configured extension that is missing, duplicated, invalid, or returns an incompatible adapter +fails explicitly at startup. Basic Memory does not silently fall back to pgvector, because doing +so would split vectors across stores while appearing healthy. + +SQLite remains automatic in this version: local SQLite databases always select `sqlite-vec`, even +if `semantic_vector_index` is set. The selector controls Postgres-backed runtimes only. + +### Extension Package Contract + +A separately distributed package registers one factory under the +`basic_memory.semantic_vector_indexes` entry-point group: + +```toml +[project.entry-points."basic_memory.semantic_vector_indexes"] +milvus = "basic_memory_milvus:create_index" +``` + +The factory receives an explicit scope and the validated Basic Memory configuration: + +```python +from basic_memory.config import BasicMemoryConfig +from basic_memory.repository.semantic_vector_index import ( + SemanticVectorIndex, + VectorIndexScope, +) + + +def create_index( + *, + scope: VectorIndexScope, + app_config: BasicMemoryConfig, +) -> SemanticVectorIndex: + ... +``` + +`VectorIndexScope` contains a stable, credential-free database namespace, project ID, embedding +identity, and vector dimensions. Extensions must isolate storage by the complete scope. They own +their client lifecycle, credentials, collection/index creation, vector persistence, and +nearest-neighbour implementation. + +The returned `SemanticVectorIndex` has five asynchronous operations: + +- `initialize()` validates or creates backend storage. +- `upsert(records)` idempotently writes vectors by `(entity_id, chunk_key)`. +- `delete(keys)` removes stable keys; missing keys are successful no-ops. +- `delete_entity(entity_id)` removes all vectors for one entity in the scope. +- `search(query, limit)` returns stable keys with normalized cosine similarity in `[0, 1]`. + +The adapter never receives a SQLAlchemy session and never calls the embedding provider. Basic +Memory owns chunking and embedding, while the extension owns vector storage and lookup. + +Adapters may also implement the separate `SemanticVectorIndexReconciler` capability. After a +vector reindex, Basic Memory passes it the complete set of current ready keys so the adapter can +delete scoped external orphans. Keeping reconciliation separate preserves the narrow required +storage protocol while allowing external stores to reclaim records left by interrupted deletes or +ready-state commits. + +### SQL Manifest and Failure Recovery + +`search_vector_chunks` remains the authoritative manifest even when vectors live in an external +store. Each row records the selected `vector_index`, embedding identity, stable chunk key, and an +`embedding_status` of `pending` or `ready`. + +Writes and deletes commit `pending` before calling the adapter. A successful adapter operation then +makes the manifest row ready or removes it. If the external operation fails, the pending row is not +searchable and the next sync safely retries the idempotent operation. Adapter search results are +hydrated only through current, ready manifest rows, so stale or orphaned external matches fail +closed. + +Switching `semantic_vector_index` invalidates existing manifest rows for incremental re-embedding. +Run `bm reindex --embeddings` after a switch to migrate all eligible content immediately. diff --git a/src/basic_memory/alembic/versions/o8j9k0l1m2n3_add_vector_index_manifest_state.py b/src/basic_memory/alembic/versions/o8j9k0l1m2n3_add_vector_index_manifest_state.py new file mode 100644 index 000000000..f9a23a2d8 --- /dev/null +++ b/src/basic_memory/alembic/versions/o8j9k0l1m2n3_add_vector_index_manifest_state.py @@ -0,0 +1,78 @@ +"""Add vector index identity and readiness to the semantic manifest. + +Revision ID: o8j9k0l1m2n3 +Revises: n7i8j9k0l1m2 +Create Date: 2026-07-21 00:00:00.000000 + +""" + +from typing import Sequence, Union + +from alembic import op +from sqlalchemy import inspect + + +revision: str = "o8j9k0l1m2n3" +down_revision: Union[str, None] = "n7i8j9k0l1m2" +branch_labels: Union[str, Sequence[str], None] = None +depends_on: Union[str, Sequence[str], None] = None + + +def upgrade() -> None: + """Make PostgreSQL chunk rows an authoritative vector-write manifest. + + SQLite creates this derived table lazily at runtime. Its schema check + rebuilds older vector tables on first semantic initialization, so only the + migration-owned PostgreSQL table needs an in-place upgrade here. + """ + connection = op.get_bind() + if connection.dialect.name != "postgresql": + return + if "search_vector_chunks" not in inspect(connection).get_table_names(): + return + + op.execute("ALTER TABLE search_vector_chunks ADD COLUMN IF NOT EXISTS vector_index TEXT") + op.execute("ALTER TABLE search_vector_chunks ADD COLUMN IF NOT EXISTS embedding_status TEXT") + op.execute("UPDATE search_vector_chunks SET vector_index = 'pgvector'") + + tables = set(inspect(connection).get_table_names()) + if "search_vector_embeddings" in tables: + op.execute( + """ + UPDATE search_vector_chunks AS chunks + SET embedding_status = CASE + WHEN EXISTS ( + SELECT 1 FROM search_vector_embeddings AS embeddings + WHERE embeddings.chunk_id = chunks.id + ) THEN 'ready' + ELSE 'pending' + END + """ + ) + else: + op.execute("UPDATE search_vector_chunks SET embedding_status = 'pending'") + + op.execute("ALTER TABLE search_vector_chunks ALTER COLUMN vector_index SET NOT NULL") + op.execute("ALTER TABLE search_vector_chunks ALTER COLUMN embedding_status SET NOT NULL") + op.create_check_constraint( + "ck_search_vector_chunks_embedding_status", + "search_vector_chunks", + "embedding_status IN ('pending', 'ready')", + ) + + +def downgrade() -> None: + """Remove vector index manifest state from PostgreSQL.""" + connection = op.get_bind() + if connection.dialect.name != "postgresql": + return + if "search_vector_chunks" not in inspect(connection).get_table_names(): + return + + op.drop_constraint( + "ck_search_vector_chunks_embedding_status", + "search_vector_chunks", + type_="check", + ) + op.execute("ALTER TABLE search_vector_chunks DROP COLUMN IF EXISTS embedding_status") + op.execute("ALTER TABLE search_vector_chunks DROP COLUMN IF EXISTS vector_index") diff --git a/src/basic_memory/config_models.py b/src/basic_memory/config_models.py index 188039cd2..b5ad9c247 100644 --- a/src/basic_memory/config_models.py +++ b/src/basic_memory/config_models.py @@ -238,6 +238,15 @@ def __init__(self, **data: Any) -> None: ... default_factory=_default_semantic_search_enabled, description="Enable semantic search (vector/hybrid retrieval). Works on both SQLite and Postgres backends. Requires semantic dependencies (included by default).", ) + semantic_vector_index: str = Field( + default="pgvector", + description=( + "Semantic vector index backend for Postgres deployments. 'pgvector' is built in; " + "other names resolve through the basic_memory.semantic_vector_indexes entry-point " + "group. SQLite continues to use sqlite-vec." + ), + min_length=1, + ) semantic_embedding_provider: str = Field( default="fastembed", description="Embedding provider for local semantic indexing/search.", diff --git a/src/basic_memory/models/search.py b/src/basic_memory/models/search.py index 75b2c67d5..ee50c7a36 100644 --- a/src/basic_memory/models/search.py +++ b/src/basic_memory/models/search.py @@ -106,6 +106,8 @@ source_hash TEXT NOT NULL, entity_fingerprint TEXT NOT NULL, embedding_model TEXT NOT NULL, + vector_index TEXT NOT NULL, + embedding_status TEXT NOT NULL CHECK (embedding_status IN ('pending', 'ready')), updated_at TIMESTAMPTZ NOT NULL DEFAULT NOW(), UNIQUE (project_id, entity_id, chunk_key) ) @@ -128,6 +130,8 @@ source_hash TEXT NOT NULL, entity_fingerprint TEXT NOT NULL, embedding_model TEXT NOT NULL, + vector_index TEXT NOT NULL, + embedding_status TEXT NOT NULL CHECK (embedding_status IN ('pending', 'ready')), updated_at TEXT NOT NULL DEFAULT CURRENT_TIMESTAMP ) """) diff --git a/src/basic_memory/repository/pgvector_index.py b/src/basic_memory/repository/pgvector_index.py new file mode 100644 index 000000000..7225d8087 --- /dev/null +++ b/src/basic_memory/repository/pgvector_index.py @@ -0,0 +1,299 @@ +"""Built-in pgvector implementation of the semantic vector index contract.""" + +from __future__ import annotations + +import asyncio +from collections.abc import Sequence + +from loguru import logger +from sqlalchemy import text +from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker + +from basic_memory import db +from basic_memory.repository.semantic_errors import SemanticDependenciesMissingError +from basic_memory.repository.semantic_vector_index import ( + SemanticVectorIndex, + VectorIndexScope, + VectorKey, + VectorMatch, + VectorRecord, + validate_query_dimensions, + validate_vector_dimensions, +) + + +class PgVectorIndex(SemanticVectorIndex): + """Persist and query semantic vectors in PostgreSQL with pgvector.""" + + def __init__( + self, + session_maker: async_sessionmaker[AsyncSession], + scope: VectorIndexScope, + ) -> None: + self._session_maker = session_maker + self.scope = scope + self._initialized = False + self._initialize_lock = asyncio.Lock() + + @staticmethod + def _format_vector(vector: Sequence[float]) -> str: + values = ",".join(f"{float(value):.12g}" for value in vector) + return f"[{values}]" + + async def initialize(self) -> None: + if self._initialized: + return + + async with self._initialize_lock: + if self._initialized: + return + + async with db.scoped_session(self._session_maker) as session: + try: + await session.execute(text("CREATE EXTENSION IF NOT EXISTS vector")) + except Exception as exc: + raise SemanticDependenciesMissingError( + "pgvector extension is unavailable for this Postgres database." + ) from exc + + existing_dimensions = await self._existing_dimensions(session) + dimensions_changed = ( + existing_dimensions is not None and existing_dimensions != self.scope.dimensions + ) + if dimensions_changed: + logger.warning( + "Embedding dimension mismatch: table has {existing}, " + "provider expects {expected}. Recreating vector storage.", + existing=existing_dimensions, + expected=self.scope.dimensions, + ) + await session.execute(text("DROP TABLE IF EXISTS search_vector_embeddings")) + + await session.execute( + text(f""" + CREATE TABLE IF NOT EXISTS search_vector_embeddings ( + chunk_id BIGINT PRIMARY KEY + REFERENCES search_vector_chunks(id) ON DELETE CASCADE, + project_id INTEGER NOT NULL, + embedding vector({self.scope.dimensions}) NOT NULL, + embedding_dims INTEGER NOT NULL, + updated_at TIMESTAMPTZ NOT NULL DEFAULT NOW() + ) + """) + ) + await session.execute( + text( + "CREATE INDEX IF NOT EXISTS " + "idx_search_vector_embeddings_project_dims " + "ON search_vector_embeddings (project_id, embedding_dims)" + ) + ) + await session.execute( + text( + "CREATE INDEX IF NOT EXISTS idx_search_vector_embeddings_hnsw " + "ON search_vector_embeddings " + "USING hnsw (embedding vector_cosine_ops) " + "WITH (m = 16, ef_construction = 64)" + ) + ) + + # Trigger: pgvector's fixed-width column was recreated for a new model. + # Why: SQL manifest rows can otherwise remain `ready` after their vectors + # disappeared, causing the incremental sync to skip them forever. + # Outcome: the normal sync pipeline re-embeds every affected chunk. + if dimensions_changed: + await session.execute( + text( + "UPDATE search_vector_chunks SET embedding_status = 'pending' " + "WHERE vector_index = 'pgvector'" + ) + ) + await session.commit() + + self._initialized = True + + async def _existing_dimensions(self, session: AsyncSession) -> int | None: + exists = await session.execute( + text( + "SELECT 1 FROM information_schema.tables " + "WHERE table_name = 'search_vector_embeddings'" + ) + ) + if exists.fetchone() is None: + return None + + result = await session.execute( + text( + "SELECT atttypmod FROM pg_attribute " + "WHERE attrelid = 'search_vector_embeddings'::regclass " + "AND attname = 'embedding'" + ) + ) + value = result.scalar_one_or_none() + return int(value) if value is not None else None + + async def _chunk_ids_by_key( + self, + session: AsyncSession, + keys: Sequence[VectorKey], + ) -> dict[VectorKey, int]: + if not keys: + return {} + + params: dict[str, object] = {"project_id": self.scope.project_id} + predicates: list[str] = [] + for index, key in enumerate(keys): + params[f"entity_id_{index}"] = key.entity_id + params[f"chunk_key_{index}"] = key.chunk_key + predicates.append( + f"(entity_id = :entity_id_{index} AND chunk_key = :chunk_key_{index})" + ) + result = await session.execute( + text( + "SELECT id, entity_id, chunk_key FROM search_vector_chunks " + "WHERE project_id = :project_id AND (" + " OR ".join(predicates) + ")" + ), + params, + ) + return { + VectorKey(entity_id=int(row["entity_id"]), chunk_key=str(row["chunk_key"])): int( + row["id"] + ) + for row in result.mappings().all() + } + + async def upsert(self, records: Sequence[VectorRecord]) -> None: + if not records: + return + validate_vector_dimensions(self.scope, records) + await self.initialize() + + async with db.scoped_session(self._session_maker) as session: + ids_by_key = await self._chunk_ids_by_key(session, [record.key for record in records]) + missing = [record.key for record in records if record.key not in ids_by_key] + if missing: + raise RuntimeError(f"Vector manifest rows are missing for keys: {missing!r}") + + params: dict[str, object] = {"project_id": self.scope.project_id} + values: list[str] = [] + for index, record in enumerate(records): + params[f"chunk_id_{index}"] = ids_by_key[record.key] + params[f"embedding_{index}"] = self._format_vector(record.values) + params[f"dimensions_{index}"] = len(record.values) + values.append( + f"(:chunk_id_{index}, :project_id, " + f"CAST(:embedding_{index} AS vector), :dimensions_{index}, NOW())" + ) + await session.execute( + text(f""" + INSERT INTO search_vector_embeddings ( + chunk_id, project_id, embedding, embedding_dims, updated_at + ) VALUES {", ".join(values)} + ON CONFLICT (chunk_id) DO UPDATE SET + project_id = EXCLUDED.project_id, + embedding = EXCLUDED.embedding, + embedding_dims = EXCLUDED.embedding_dims, + updated_at = NOW() + """), + params, + ) + await session.commit() + + async def delete(self, keys: Sequence[VectorKey]) -> None: + if not keys: + return + await self.initialize() + async with db.scoped_session(self._session_maker) as session: + ids_by_key = await self._chunk_ids_by_key(session, keys) + chunk_ids = list(ids_by_key.values()) + if chunk_ids: + params = {f"chunk_id_{index}": value for index, value in enumerate(chunk_ids)} + placeholders = ", ".join(f":chunk_id_{index}" for index in range(len(chunk_ids))) + await session.execute( + text( + f"DELETE FROM search_vector_embeddings WHERE chunk_id IN ({placeholders})" + ), + params, + ) + await session.commit() + + async def delete_entity(self, entity_id: int) -> None: + await self.initialize() + async with db.scoped_session(self._session_maker) as session: + await session.execute( + text( + "DELETE FROM search_vector_embeddings WHERE chunk_id IN (" + "SELECT id FROM search_vector_chunks " + "WHERE project_id = :project_id AND entity_id = :entity_id)" + ), + {"project_id": self.scope.project_id, "entity_id": entity_id}, + ) + await session.commit() + + async def delete_orphans(self, _live_keys: Sequence[VectorKey]) -> None: + """Remove pgvector rows absent from the current ready manifest scope.""" + await self.initialize() + async with db.scoped_session(self._session_maker) as session: + await session.execute( + text( + "DELETE FROM search_vector_embeddings AS embeddings " + "WHERE embeddings.project_id = :project_id AND NOT EXISTS (" + "SELECT 1 FROM search_vector_chunks AS chunks " + "WHERE chunks.id = embeddings.chunk_id " + "AND chunks.project_id = :project_id " + "AND chunks.vector_index = 'pgvector' " + "AND chunks.embedding_model = :embedding_identity " + "AND chunks.embedding_status = 'ready')" + ), + { + "project_id": self.scope.project_id, + "embedding_identity": self.scope.embedding_identity, + }, + ) + await session.commit() + + async def search( + self, + query: Sequence[float], + *, + limit: int, + ) -> list[VectorMatch]: + if not query or limit <= 0: + return [] + validate_query_dimensions(self.scope, query) + await self.initialize() + async with db.scoped_session(self._session_maker) as session: + result = await session.execute( + text( + "SELECT c.entity_id, c.chunk_key, " + "1 - (e.embedding <=> CAST(:query AS vector)) AS similarity " + "FROM search_vector_embeddings e " + "JOIN search_vector_chunks c ON c.id = e.chunk_id " + "WHERE e.project_id = :project_id " + "AND e.embedding_dims = :dimensions " + "AND c.project_id = :project_id " + "AND c.vector_index = 'pgvector' " + "AND c.embedding_status = 'ready' " + "AND c.embedding_model = :embedding_identity " + "ORDER BY e.embedding <=> CAST(:query AS vector), " + "c.entity_id ASC, c.chunk_key ASC " + "LIMIT :limit" + ), + { + "query": self._format_vector(query), + "project_id": self.scope.project_id, + "dimensions": self.scope.dimensions, + "embedding_identity": self.scope.embedding_identity, + "limit": limit, + }, + ) + return [ + VectorMatch( + key=VectorKey( + entity_id=int(row["entity_id"]), + chunk_key=str(row["chunk_key"]), + ), + similarity=max(0.0, min(1.0, float(row["similarity"]))), + ) + for row in result.mappings().all() + ] diff --git a/src/basic_memory/repository/postgres_search_repository.py b/src/basic_memory/repository/postgres_search_repository.py index 2a17e4758..d370301ce 100644 --- a/src/basic_memory/repository/postgres_search_repository.py +++ b/src/basic_memory/repository/postgres_search_repository.py @@ -12,7 +12,7 @@ from sqlalchemy.ext.asyncio import AsyncSession from basic_memory import db -from basic_memory.config import BasicMemoryConfig, ConfigManager +from basic_memory.config import BasicMemoryConfig, ConfigManager, DatabaseBackend from basic_memory.repository.embedding_provider import EmbeddingProvider from basic_memory.repository.embedding_provider_factory import create_embedding_provider from basic_memory.repository.search_index_row import SearchIndexRow @@ -24,6 +24,12 @@ ) from basic_memory.repository.metadata_filters import parse_metadata_filters from basic_memory.repository.semantic_errors import SemanticDependenciesMissingError +from basic_memory.repository.semantic_vector_index import SemanticVectorIndex +from basic_memory.repository.semantic_vector_index_factory import ( + build_vector_index_scope, + resolve_semantic_vector_index_name, +) +from basic_memory.repository.pgvector_index import PgVectorIndex from basic_memory.schemas.search import SearchItemType, SearchRetrievalMode @@ -58,6 +64,8 @@ def __init__( project_id: int, app_config: BasicMemoryConfig | None = None, embedding_provider: EmbeddingProvider | None = None, + vector_index_name: str | None = None, + vector_index: SemanticVectorIndex | None = None, ): super().__init__(session_maker, project_id) self._app_config = app_config or ConfigManager().config @@ -71,6 +79,7 @@ def __init__( self._app_config.semantic_postgres_prepare_concurrency ) self._embedding_provider = embedding_provider + self._semantic_vector_index_name = vector_index_name or "pgvector" self._vector_dimensions = 384 self._vector_tables_initialized = False self._vector_tables_lock = asyncio.Lock() @@ -79,6 +88,26 @@ def __init__( self._embedding_provider = create_embedding_provider(self._app_config) if self._embedding_provider is not None: self._vector_dimensions = self._embedding_provider.dimensions + effective_name = vector_index_name or resolve_semantic_vector_index_name( + self._app_config, + DatabaseBackend.POSTGRES, + ) + if vector_index is None: + if effective_name != "pgvector": + raise SemanticDependenciesMissingError( + f"Semantic vector index '{effective_name}' must be created by the " + "search repository composition root." + ) + vector_index = PgVectorIndex( + session_maker, + build_vector_index_scope( + self._app_config, + self._embedding_provider, + project_id, + ), + ) + self._semantic_vector_index_name = effective_name + self._semantic_vector_index = vector_index async def init_search_index(self): """Create Postgres table with tsvector column and GIN indexes. @@ -275,23 +304,23 @@ def _prepare_single_term(self, term: str, is_prefix: bool = True) -> str: else: return cleaned_term - # ------------------------------------------------------------------ - # pgvector utility - # ------------------------------------------------------------------ - - @staticmethod - def _format_pgvector_literal(vector: list[float]) -> str: - if not vector: - return "[]" - values = ",".join(f"{float(value):.12g}" for value in vector) - return f"[{values}]" - # ------------------------------------------------------------------ # Abstract hook implementations (vector/semantic, Postgres-specific) # ------------------------------------------------------------------ async def _ensure_vector_tables(self) -> None: self._assert_semantic_available() + if not hasattr(self, "_semantic_vector_index"): + assert self._embedding_provider is not None + self._semantic_vector_index_name = "pgvector" + self._semantic_vector_index = PgVectorIndex( + self.session_maker, + build_vector_index_scope( + self._app_config, + self._embedding_provider, + self.project_id, + ), + ) if self._vector_tables_initialized: return @@ -302,13 +331,6 @@ async def _ensure_vector_tables(self) -> None: return async with db.scoped_session(self.session_maker) as session: - try: - await session.execute(text("CREATE EXTENSION IF NOT EXISTS vector")) - except Exception as exc: - raise SemanticDependenciesMissingError( - "pgvector extension is unavailable for this Postgres database." - ) from exc - # --- Chunks table (dimension-independent, may already exist via migration) --- # Trigger: fresh Postgres projects may not have vector chunk tables yet. # Why: runtime can bootstrap missing tables, but schema evolution must stay @@ -326,6 +348,9 @@ async def _ensure_vector_tables(self) -> None: source_hash TEXT NOT NULL, entity_fingerprint TEXT NOT NULL, embedding_model TEXT NOT NULL, + vector_index TEXT NOT NULL, + embedding_status TEXT NOT NULL + CHECK (embedding_status IN ('pending', 'ready')), updated_at TIMESTAMPTZ NOT NULL DEFAULT NOW(), UNIQUE (project_id, entity_id, chunk_key) ) @@ -341,136 +366,20 @@ async def _ensure_vector_tables(self) -> None: ) ) - # --- Embeddings table (dimension-dependent, created at runtime) --- - # Trigger: provider dimensions may differ from what was previously deployed. - # Why: the column type `vector(N)` is fixed at table creation; switching - # from FastEmbed (384) to OpenAI (1536) requires recreation. - # Outcome: mismatched table is dropped and recreated with correct dims. - # Embeddings are derived data — re-indexing will repopulate them. - existing_dims = await self._get_existing_embedding_dims(session) - if existing_dims is not None and existing_dims != self._vector_dimensions: - logger.warning( - f"Embedding dimension mismatch: table has {existing_dims}, " - f"provider expects {self._vector_dimensions}. " - "Dropping and recreating search_vector_embeddings." - ) - await session.execute(text("DROP TABLE IF EXISTS search_vector_embeddings")) - - await session.execute( - text( - f""" - CREATE TABLE IF NOT EXISTS search_vector_embeddings ( - chunk_id BIGINT PRIMARY KEY - REFERENCES search_vector_chunks(id) ON DELETE CASCADE, - project_id INTEGER NOT NULL, - embedding vector({self._vector_dimensions}) NOT NULL, - embedding_dims INTEGER NOT NULL, - updated_at TIMESTAMPTZ NOT NULL DEFAULT NOW() - ) - """ - ) - ) - await session.execute( - text( - """ - CREATE INDEX IF NOT EXISTS idx_search_vector_embeddings_project_dims - ON search_vector_embeddings (project_id, embedding_dims) - """ - ) - ) - # HNSW index for approximate nearest-neighbour search. - # Without this every vector query is a sequential scan. - await session.execute( - text( - """ - CREATE INDEX IF NOT EXISTS idx_search_vector_embeddings_hnsw - ON search_vector_embeddings - USING hnsw (embedding vector_cosine_ops) - WITH (m = 16, ef_construction = 64) - """ - ) - ) await session.commit() + await self._semantic_vector_index.initialize() + logger.debug(f"Postgres vector tables ready (dimensions={self._vector_dimensions})") self._vector_tables_initialized = True - async def _get_existing_embedding_dims(self, session: AsyncSession) -> int | None: - """Query the vector column dimension from an existing search_vector_embeddings table. - - Returns None when the table does not exist. - Uses information_schema to avoid regclass cast errors on missing tables, - then reads atttypmod from pg_attribute for the actual dimension value. - """ - # Check table existence via information_schema (no exception on missing) - exists_result = await session.execute( - text( - """ - SELECT 1 FROM information_schema.tables - WHERE table_name = 'search_vector_embeddings' - """ - ) - ) - if exists_result.fetchone() is None: - return None - - result = await session.execute( - text( - """ - SELECT atttypmod - FROM pg_attribute - WHERE attrelid = 'search_vector_embeddings'::regclass - AND attname = 'embedding' - """ - ) - ) - row = result.fetchone() - if row is None: - return None - # pgvector stores dimensions in atttypmod directly - return int(row[0]) - async def _run_vector_query( self, session: AsyncSession, query_embedding: list[float], candidate_limit: int, ) -> list[dict]: - if not query_embedding: - return [] - - embedding_dims = len(query_embedding) - query_embedding_literal = self._format_pgvector_literal(query_embedding) - - vector_result = await session.execute( - text( - """ - WITH vector_matches AS ( - SELECT - e.chunk_id, - (e.embedding <=> CAST(:query_embedding AS vector)) AS distance - FROM search_vector_embeddings e - WHERE e.project_id = :project_id - AND e.embedding_dims = :embedding_dims - ORDER BY e.embedding <=> CAST(:query_embedding AS vector) - LIMIT :vector_k - ) - SELECT c.entity_id, c.chunk_key, c.chunk_text, vector_matches.distance AS best_distance - FROM vector_matches - JOIN search_vector_chunks c ON c.id = vector_matches.chunk_id - WHERE c.project_id = :project_id - ORDER BY best_distance ASC - LIMIT :vector_k - """ - ), - { - "query_embedding": query_embedding_literal, - "project_id": self.project_id, - "embedding_dims": embedding_dims, - "vector_k": candidate_limit, - }, - ) - return [dict(row) for row in vector_result.mappings().all()] + return await super()._run_vector_query(session, query_embedding, candidate_limit) def _vector_prepare_window_size(self) -> int: """Use a bounded config-driven prepare window for Postgres vector sync.""" @@ -493,6 +402,7 @@ async def _upsert_scheduled_chunk_records( upsert_params: dict[str, object] = { "project_id": self.project_id, "entity_id": entity_id, + "vector_index": self._semantic_vector_index_name, } upsert_values: list[str] = [] # The SQL template is built from integer enumerate() indices only. @@ -507,7 +417,8 @@ async def _upsert_scheduled_chunk_records( "(" ":entity_id, :project_id, " f":chunk_key_{index}, :chunk_text_{index}, :source_hash_{index}, " - f":entity_fingerprint_{index}, :embedding_model_{index}, NOW()" + f":entity_fingerprint_{index}, :embedding_model_{index}, " + ":vector_index, 'pending', NOW()" ")" ) @@ -521,6 +432,8 @@ async def _upsert_scheduled_chunk_records( source_hash, entity_fingerprint, embedding_model, + vector_index, + embedding_status, updated_at ) VALUES {", ".join(upsert_values)} ON CONFLICT (project_id, entity_id, chunk_key) DO UPDATE SET @@ -528,6 +441,8 @@ async def _upsert_scheduled_chunk_records( source_hash = EXCLUDED.source_hash, entity_fingerprint = EXCLUDED.entity_fingerprint, embedding_model = EXCLUDED.embedding_model, + vector_index = EXCLUDED.vector_index, + embedding_status = EXCLUDED.embedding_status, updated_at = NOW() RETURNING id, chunk_key """), @@ -541,59 +456,12 @@ async def _upsert_scheduled_chunk_records( for record in scheduled_records ] - async def _write_embeddings( - self, - session: AsyncSession, - jobs: list[tuple[int, str]], - embeddings: list[list[float]], - ) -> None: - params: dict[str, object] = {"project_id": self.project_id} - value_rows: list[str] = [] - - # The SQL template is built from integer enumerate() indices only. - # No user-controlled text is interpolated into the statement. - for index, ((row_id, _), vector) in enumerate(zip(jobs, embeddings, strict=True)): - params[f"chunk_id_{index}"] = row_id - params[f"embedding_{index}"] = self._format_pgvector_literal(vector) - params[f"embedding_dims_{index}"] = len(vector) - value_rows.append( - "(" - f":chunk_id_{index}, :project_id, CAST(:embedding_{index} AS vector), " - f":embedding_dims_{index}, NOW()" - ")" - ) - - await session.execute( - text(f""" - INSERT INTO search_vector_embeddings ( - chunk_id, - project_id, - embedding, - embedding_dims, - updated_at - ) VALUES {", ".join(value_rows)} - ON CONFLICT (chunk_id) DO UPDATE SET - project_id = EXCLUDED.project_id, - embedding = EXCLUDED.embedding, - embedding_dims = EXCLUDED.embedding_dims, - updated_at = NOW() - """), - params, - ) - async def _delete_entity_chunks( self, session: AsyncSession, entity_id: int, ) -> None: - # Postgres has ON DELETE CASCADE from embeddings → chunks - await session.execute( - text( - "DELETE FROM search_vector_chunks " - "WHERE project_id = :project_id AND entity_id = :entity_id" - ), - {"project_id": self.project_id, "entity_id": entity_id}, - ) + await super()._delete_entity_chunks(session, entity_id) async def _delete_stale_chunks( self, @@ -601,21 +469,7 @@ async def _delete_stale_chunks( stale_ids: list[int], entity_id: int, ) -> None: - stale_placeholders = ", ".join(f":stale_id_{idx}" for idx in range(len(stale_ids))) - stale_params = { - "project_id": self.project_id, - "entity_id": entity_id, - **{f"stale_id_{idx}": row_id for idx, row_id in enumerate(stale_ids)}, - } - # CASCADE handles embedding deletion - await session.execute( - text( - "DELETE FROM search_vector_chunks " - f"WHERE id IN ({stale_placeholders}) " - "AND project_id = :project_id AND entity_id = :entity_id" - ), - stale_params, - ) + await super()._delete_stale_chunks(session, stale_ids, entity_id) def _distance_to_similarity(self, distance: float) -> float: """Convert pgvector cosine distance to cosine similarity. diff --git a/src/basic_memory/repository/search_repository.py b/src/basic_memory/repository/search_repository.py index ed6c05729..d15b7c36a 100644 --- a/src/basic_memory/repository/search_repository.py +++ b/src/basic_memory/repository/search_repository.py @@ -17,6 +17,7 @@ from basic_memory.repository.postgres_search_repository import PostgresSearchRepository from basic_memory.repository.search_index_row import SearchIndexRow from basic_memory.repository.search_repository_base import VectorSyncBatchResult +from basic_memory.repository.semantic_vector_index_factory import create_semantic_vector_index from basic_memory.repository.sqlite_search_repository import SQLiteSearchRepository from basic_memory.schemas.search import SearchItemType, SearchRetrievalMode @@ -96,6 +97,18 @@ async def delete_entity_vector_rows(self, entity_id: int) -> None: """Delete semantic vector chunks and embeddings for one entity.""" ... + async def delete_project_vector_rows(self) -> None: + """Delete all semantic vector chunks and embeddings for this project.""" + ... + + async def delete_stale_vector_rows(self) -> None: + """Delete semantic vectors whose source entities no longer exist.""" + ... + + async def reconcile_vector_index(self) -> None: + """Remove adapter vectors that have no current ready manifest row.""" + ... + async def sync_entity_vectors_batch( self, entity_ids: list[int], @@ -138,8 +151,17 @@ def create_search_repository( # Outcome: resolve the cached singleton here once and inject it, so the provider # is the single source of truth across all callers of this factory. embedding_provider = None + vector_index_name = None + vector_index = None if config.semantic_search_enabled: embedding_provider = create_embedding_provider(config) + vector_index_name, vector_index = create_semantic_vector_index( + session_maker=session_maker, + project_id=project_id, + app_config=config, + database_backend=database_backend, + embedding_provider=embedding_provider, + ) if database_backend == DatabaseBackend.POSTGRES: # pragma: no cover return PostgresSearchRepository( # pragma: no cover @@ -147,6 +169,8 @@ def create_search_repository( project_id=project_id, app_config=app_config, embedding_provider=embedding_provider, + vector_index_name=vector_index_name, + vector_index=vector_index, ) else: return SQLiteSearchRepository( @@ -154,6 +178,8 @@ def create_search_repository( project_id=project_id, app_config=app_config, embedding_provider=embedding_provider, + vector_index_name=vector_index_name, + vector_index=vector_index, ) diff --git a/src/basic_memory/repository/search_repository_base.py b/src/basic_memory/repository/search_repository_base.py index 5719e7774..f2c184496 100644 --- a/src/basic_memory/repository/search_repository_base.py +++ b/src/basic_memory/repository/search_repository_base.py @@ -32,6 +32,13 @@ SemanticDependenciesMissingError, SemanticSearchDisabledError, ) +from basic_memory.repository.semantic_vector_index import ( + SemanticVectorIndex, + SemanticVectorIndexReconciler, + VectorKey, + VectorMatch, + VectorRecord, +) from basic_memory.repository.semantic_vector_sync import ( EntitySyncRuntime as _EntitySyncRuntime, EntityVectorShardPlan as _EntityVectorShardPlan, @@ -80,6 +87,8 @@ class SearchRepositoryBase(ABC): _semantic_embedding_sync_batch_size: int _vector_dimensions: int _vector_tables_initialized: bool + _semantic_vector_index: SemanticVectorIndex + _semantic_vector_index_name: str = "" def __init__(self, session_maker: async_sessionmaker[AsyncSession], project_id: int): """Initialize with session maker and project_id filter. @@ -199,55 +208,240 @@ async def _ensure_vector_tables(self) -> None: """Create backend-specific vector chunk and embedding tables.""" pass - @abstractmethod async def _run_vector_query( self, session: AsyncSession, query_embedding: list[float], candidate_limit: int, ) -> list[dict]: - """Execute backend-specific nearest-neighbour vector query. + """Query the configured adapter and hydrate only live, ready manifest rows.""" + matches = await self._semantic_vector_index.search( + query_embedding, + limit=candidate_limit, + ) + return await self._hydrate_vector_matches(session, matches) - Returns list of mappings with keys ``entity_id`` and ``best_distance``. - """ - pass + async def _hydrate_vector_matches( + self, + session: AsyncSession, + matches: list[VectorMatch], + ) -> list[dict]: + """Resolve adapter matches through the authoritative ready manifest.""" + if not matches: + return [] + + params: dict[str, object] = { + "project_id": self.project_id, + "vector_index": self._semantic_vector_index_name, + "embedding_model": self._embedding_model_key(), + } + predicates: list[str] = [] + for index, match in enumerate(matches): + params[f"entity_id_{index}"] = match.key.entity_id + params[f"chunk_key_{index}"] = match.key.chunk_key + predicates.append( + f"(entity_id = :entity_id_{index} AND chunk_key = :chunk_key_{index})" + ) + + result = await session.execute( + text( + "SELECT entity_id, chunk_key, chunk_text FROM search_vector_chunks " + "WHERE project_id = :project_id " + "AND vector_index = :vector_index " + "AND embedding_model = :embedding_model " + "AND embedding_status = 'ready' " + "AND (" + " OR ".join(predicates) + ")" + ), + params, + ) + chunks_by_key = { + VectorKey( + entity_id=int(row["entity_id"]), + chunk_key=str(row["chunk_key"]), + ): str(row["chunk_text"]) + for row in result.mappings().all() + } + return [ + { + "entity_id": match.key.entity_id, + "chunk_key": match.key.chunk_key, + "chunk_text": chunks_by_key[match.key], + "best_similarity": match.similarity, + } + for match in matches + if match.key in chunks_by_key + ] - @abstractmethod async def _write_embeddings( self, session: AsyncSession, jobs: list[tuple[int, str]], embeddings: list[list[float]], ) -> None: - """Write embedding vectors for the given chunk row IDs. + """Legacy storage hook retained for focused pre-adapter test repositories.""" + raise NotImplementedError - ``jobs`` is a list of ``(chunk_row_id, chunk_text)`` pairs. - ``embeddings`` is the corresponding list of vectors. - """ - pass + async def _persist_embeddings( + self, + jobs: list[tuple[int, str]], + embeddings: list[list[float]], + ) -> None: + """Write vectors through the adapter, then make their manifest rows ready.""" + if not jobs: + return + + # Compatibility: focused orchestration tests and third-party subclasses + # from before the adapter contract may still override the private writer. + # Real repositories always configure `_semantic_vector_index` and take the + # manifest-safe path below. + if not hasattr(self, "_semantic_vector_index"): + async with db.scoped_session(self.session_maker) as session: + await self._prepare_vector_session(session) + await self._write_embeddings(session, jobs, embeddings) + await session.commit() + return + + row_ids = [row_id for row_id, _ in jobs] + params = {f"row_id_{index}": row_id for index, row_id in enumerate(row_ids)} + placeholders = ", ".join(f":row_id_{index}" for index in range(len(row_ids))) + async with db.scoped_session(self.session_maker) as session: + result = await session.execute( + text( + "SELECT id, entity_id, chunk_key FROM search_vector_chunks " + f"WHERE project_id = :project_id AND id IN ({placeholders})" + ), + {**params, "project_id": self.project_id}, + ) + keys_by_id = { + int(row["id"]): VectorKey( + entity_id=int(row["entity_id"]), + chunk_key=str(row["chunk_key"]), + ) + for row in result.mappings().all() + } + + missing_row_ids = [row_id for row_id in row_ids if row_id not in keys_by_id] + if missing_row_ids: + raise RuntimeError(f"Vector manifest rows disappeared before write: {missing_row_ids}") + + records = [ + VectorRecord(key=keys_by_id[row_id], values=tuple(embedding)) + for (row_id, _), embedding in zip(jobs, embeddings, strict=True) + ] + await self._semantic_vector_index.upsert(records) + + # Trigger: the adapter write completed successfully. + # Why: only SQL rows marked ready may hydrate search matches. If this commit + # fails, the adapter's idempotent upsert is retried while the row stays pending. + # Outcome: external partial failures fail closed without cross-store transactions. + async with db.scoped_session(self.session_maker) as session: + await session.execute( + text( + "UPDATE search_vector_chunks SET embedding_status = 'ready', " + f"updated_at = {self._timestamp_now_expr()} " + f"WHERE project_id = :project_id AND id IN ({placeholders}) " + "AND vector_index = :vector_index " + "AND embedding_model = :embedding_model" + ), + { + **params, + "project_id": self.project_id, + "vector_index": self._semantic_vector_index_name, + "embedding_model": self._embedding_model_key(), + }, + ) + await session.commit() - @abstractmethod async def _delete_entity_chunks( self, session: AsyncSession, entity_id: int, ) -> None: - """Delete all chunk + embedding rows for an entity. - - SQLite must explicitly delete embeddings first (no CASCADE). - Postgres relies on ON DELETE CASCADE from the FK. - """ - pass + """Stage an entity deletion by making its manifest rows non-searchable.""" + await session.execute( + text( + "UPDATE search_vector_chunks SET embedding_status = 'pending' " + "WHERE project_id = :project_id AND entity_id = :entity_id" + ), + {"project_id": self.project_id, "entity_id": entity_id}, + ) - @abstractmethod async def _delete_stale_chunks( self, session: AsyncSession, stale_ids: list[int], entity_id: int, ) -> None: - """Delete stale chunk rows (and their embeddings) by ID.""" - pass + """Stage stale chunk deletion by making manifest rows non-searchable.""" + if not stale_ids: + return + params = {f"stale_id_{index}": row_id for index, row_id in enumerate(stale_ids)} + placeholders = ", ".join(f":stale_id_{index}" for index in range(len(stale_ids))) + await session.execute( + text( + "UPDATE search_vector_chunks SET embedding_status = 'pending' " + f"WHERE project_id = :project_id AND entity_id = :entity_id " + f"AND id IN ({placeholders})" + ), + {**params, "project_id": self.project_id, "entity_id": entity_id}, + ) + + async def _finalize_prepared_vector_deletions( + self, + prepared: _PreparedEntityVectorSync, + ) -> None: + """Delete staged adapter records, then remove their SQL manifest rows. + + The prepare transaction commits `pending` first. If an external delete + fails, those rows remain non-searchable and the next sync retries the + idempotent delete instead of losing cleanup intent. + """ + if not prepared.delete_entity_vectors and not prepared.stale_chunk_ids: + return + if not hasattr(self, "_semantic_vector_index"): + return + + if prepared.delete_entity_vectors: + await self._semantic_vector_index.delete_entity(prepared.entity_id) + async with db.scoped_session(self.session_maker) as session: + await session.execute( + text( + "DELETE FROM search_vector_chunks " + "WHERE project_id = :project_id AND entity_id = :entity_id" + ), + {"project_id": self.project_id, "entity_id": prepared.entity_id}, + ) + await session.commit() + return + + row_ids = prepared.stale_chunk_ids + params = {f"row_id_{index}": row_id for index, row_id in enumerate(row_ids)} + placeholders = ", ".join(f":row_id_{index}" for index in range(len(row_ids))) + async with db.scoped_session(self.session_maker) as session: + result = await session.execute( + text( + "SELECT entity_id, chunk_key FROM search_vector_chunks " + f"WHERE project_id = :project_id AND entity_id = :entity_id " + f"AND id IN ({placeholders})" + ), + {**params, "project_id": self.project_id, "entity_id": prepared.entity_id}, + ) + keys = [ + VectorKey(entity_id=int(row["entity_id"]), chunk_key=str(row["chunk_key"])) + for row in result.mappings().all() + ] + + await self._semantic_vector_index.delete(keys) + async with db.scoped_session(self.session_maker) as session: + await session.execute( + text( + "DELETE FROM search_vector_chunks " + f"WHERE project_id = :project_id AND entity_id = :entity_id " + f"AND id IN ({placeholders})" + ), + {**params, "project_id": self.project_id, "entity_id": prepared.entity_id}, + ) + await session.commit() @abstractmethod def _distance_to_similarity(self, distance: float) -> float: @@ -403,9 +597,103 @@ async def delete_entity_vector_rows(self, entity_id: int) -> None: await self._ensure_vector_tables() async with db.scoped_session(self.session_maker) as session: - await self._prepare_vector_session(session) await self._delete_entity_chunks(session, entity_id) await session.commit() + await self._finalize_prepared_vector_deletions( + _PreparedEntityVectorSync( + entity_id=entity_id, + sync_start=time.perf_counter(), + source_rows_count=0, + embedding_jobs=[], + delete_entity_vectors=True, + ) + ) + + async def delete_project_vector_rows(self) -> None: + """Delete this project's vectors through the configured storage adapter. + + Core enumerates ownership from the SQL manifest because the adapter + contract intentionally has no project-wide listing or destructive reset. + Each entity deletion uses the same pending-first lifecycle as ordinary + sync cleanup, so an external failure remains retryable and fails closed. + """ + if not self._semantic_enabled: + return + + await self._ensure_vector_tables() + async with db.scoped_session(self.session_maker) as session: + result = await session.execute( + text( + "SELECT DISTINCT entity_id FROM search_vector_chunks " + "WHERE project_id = :project_id ORDER BY entity_id" + ), + {"project_id": self.project_id}, + ) + entity_ids = [int(entity_id) for entity_id in result.scalars().all()] + + for entity_id in entity_ids: + await self.delete_entity_vector_rows(entity_id) + + async def delete_stale_vector_rows(self) -> None: + """Delete vectors whose source entity no longer exists. + + The SQL manifest remains the source of truth for ownership. External + indexes receive stable entity deletes before their manifest rows are + removed, avoiding backend-specific cleanup in the service layer. + """ + if not self._semantic_enabled: + return + + await self._ensure_vector_tables() + async with db.scoped_session(self.session_maker) as session: + result = await session.execute( + text( + "SELECT DISTINCT entity_id FROM search_vector_chunks " + "WHERE project_id = :project_id AND entity_id NOT IN (" + "SELECT id FROM entity WHERE project_id = :project_id) " + "ORDER BY entity_id" + ), + {"project_id": self.project_id}, + ) + entity_ids = [int(entity_id) for entity_id in result.scalars().all()] + + for entity_id in entity_ids: + await self.delete_entity_vector_rows(entity_id) + + async def reconcile_vector_index(self) -> None: + """Let capable adapters prune records absent from the ready SQL manifest.""" + if not self._semantic_enabled: + return + + await self._ensure_vector_tables() + if not isinstance(self._semantic_vector_index, SemanticVectorIndexReconciler): + return + + async with db.scoped_session(self.session_maker) as session: + result = await session.execute( + text( + "SELECT entity_id, chunk_key FROM search_vector_chunks " + "WHERE project_id = :project_id " + "AND vector_index = :vector_index " + "AND embedding_model = :embedding_model " + "AND embedding_status = 'ready' " + "ORDER BY entity_id, chunk_key" + ), + { + "project_id": self.project_id, + "vector_index": self._semantic_vector_index_name, + "embedding_model": self._embedding_model_key(), + }, + ) + live_keys = [ + VectorKey( + entity_id=int(row["entity_id"]), + chunk_key=str(row["chunk_key"]), + ) + for row in result.mappings().all() + ] + + await self._semantic_vector_index.delete_orphans(live_keys) # ------------------------------------------------------------------ # Shared semantic search: guard, text processing, chunking @@ -867,9 +1155,23 @@ async def _search_vector_only( embed_ms = (time.perf_counter() - embed_start) * 1000 vector_query_start = time.perf_counter() - async with db.scoped_session(self.session_maker) as session: - await self._prepare_vector_session(session) - vector_rows = await self._run_vector_query(session, query_embedding, candidate_limit) + if hasattr(self, "_semantic_vector_index"): + matches = await self._semantic_vector_index.search( + query_embedding, + limit=candidate_limit, + ) + async with db.scoped_session(self.session_maker) as session: + vector_rows = await self._hydrate_vector_matches(session, matches) + else: + # Compatibility for focused test repositories that implement the + # pre-extension private query hook without configuring an adapter. + async with db.scoped_session(self.session_maker) as session: + await self._prepare_vector_session(session) + vector_rows = await self._run_vector_query( + session, + query_embedding, + candidate_limit, + ) vector_query_ms = (time.perf_counter() - vector_query_start) * 1000 vector_row_count = len(vector_rows) hydrate_ms = 0.0 @@ -910,8 +1212,12 @@ def _log_vector_summary() -> None: chunks_by_si_key: dict[SearchIndexKey, list[tuple[float, str]]] = {} for row in vector_rows: chunk_key = row.get("chunk_key", "") - distance = float(row["best_distance"]) - similarity = self._distance_to_similarity(distance) + if "best_similarity" in row: + similarity = float(row["best_similarity"]) + else: + # Compatibility: private test doubles may still return native distance. + distance = float(row["best_distance"]) + similarity = self._distance_to_similarity(distance) chunk_text = row.get("chunk_text", "") try: si_key = self._parse_chunk_key(chunk_key) diff --git a/src/basic_memory/repository/semantic_errors.py b/src/basic_memory/repository/semantic_errors.py index edc6fe9da..e8ddf006e 100644 --- a/src/basic_memory/repository/semantic_errors.py +++ b/src/basic_memory/repository/semantic_errors.py @@ -7,3 +7,7 @@ class SemanticSearchDisabledError(RuntimeError): class SemanticDependenciesMissingError(RuntimeError): """Raised when a semantic search dependency is unavailable or misconfigured.""" + + +class SemanticVectorIndexExtensionError(SemanticDependenciesMissingError): + """Raised when a configured external vector index cannot be loaded safely.""" diff --git a/src/basic_memory/repository/semantic_vector_index.py b/src/basic_memory/repository/semantic_vector_index.py new file mode 100644 index 000000000..ece213428 --- /dev/null +++ b/src/basic_memory/repository/semantic_vector_index.py @@ -0,0 +1,115 @@ +"""Storage contract for semantic vector index backends.""" + +from __future__ import annotations + +from collections.abc import Sequence +from dataclasses import dataclass +from typing import Protocol, runtime_checkable + + +SEMANTIC_VECTOR_INDEX_ENTRY_POINT_GROUP = "basic_memory.semantic_vector_indexes" + + +@dataclass(frozen=True, slots=True) +class VectorIndexScope: + """Stable isolation and embedding identity for one project's vectors.""" + + namespace: str + project_id: int + embedding_identity: str + dimensions: int + + +@dataclass(frozen=True, slots=True) +class VectorKey: + """Backend-independent identity for one semantic chunk vector.""" + + entity_id: int + chunk_key: str + + +@dataclass(frozen=True, slots=True) +class VectorRecord: + """One vector value to insert or replace idempotently.""" + + key: VectorKey + values: tuple[float, ...] + + +@dataclass(frozen=True, slots=True) +class VectorMatch: + """One nearest-neighbour match with normalized cosine similarity.""" + + key: VectorKey + similarity: float + + +@runtime_checkable +class SemanticVectorIndex(Protocol): + """Narrow storage boundary implemented by built-in and external indexes. + + The contract deliberately has no SQLAlchemy session and never calls an + embedding provider. Core owns the SQL manifest and embedding lifecycle; + implementations own only vector persistence and nearest-neighbour lookup. + """ + + scope: VectorIndexScope + + async def initialize(self) -> None: + """Create or validate backend storage for the configured scope.""" + ... + + async def upsert(self, records: Sequence[VectorRecord]) -> None: + """Insert or replace vectors by stable key.""" + ... + + async def delete(self, keys: Sequence[VectorKey]) -> None: + """Delete vectors by stable key; missing keys are successful no-ops.""" + ... + + async def delete_entity(self, entity_id: int) -> None: + """Delete every vector owned by an entity in this scope.""" + ... + + async def search( + self, + query: Sequence[float], + *, + limit: int, + ) -> list[VectorMatch]: + """Return nearest matches ordered by normalized cosine similarity.""" + ... + + +@runtime_checkable +class SemanticVectorIndexReconciler(Protocol): + """Optional cleanup capability for removing vectors absent from the live manifest.""" + + scope: VectorIndexScope + + async def delete_orphans(self, live_keys: Sequence[VectorKey]) -> None: + """Delete scoped vectors whose stable keys are not in ``live_keys``.""" + ... + + +def validate_vector_dimensions( + scope: VectorIndexScope, + records: Sequence[VectorRecord], +) -> None: + """Fail before a backend write when a vector has the wrong dimensions.""" + for record in records: + if len(record.values) != scope.dimensions: + raise ValueError( + "Vector dimensions do not match the configured index scope: " + f"expected {scope.dimensions}, got {len(record.values)} " + f"for {record.key.chunk_key}." + ) + + +def validate_query_dimensions(scope: VectorIndexScope, query: Sequence[float]) -> None: + """Fail before a backend query when the query vector has the wrong dimensions.""" + if len(query) != scope.dimensions: + raise ValueError( + "Query dimensions do not match the configured index scope: " + f"expected {scope.dimensions}, got {len(query)}." + ) diff --git a/src/basic_memory/repository/semantic_vector_index_factory.py b/src/basic_memory/repository/semantic_vector_index_factory.py new file mode 100644 index 000000000..3f36fe6bb --- /dev/null +++ b/src/basic_memory/repository/semantic_vector_index_factory.py @@ -0,0 +1,136 @@ +"""Composition-root factory for built-in and extension vector indexes.""" + +from __future__ import annotations + +import hashlib +from importlib.metadata import entry_points +from typing import Protocol + +from sqlalchemy.engine import make_url +from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker + +from basic_memory.config import BasicMemoryConfig, DatabaseBackend +from basic_memory.repository.embedding_provider import ( + EmbeddingProvider, + embedding_provider_identity, +) +from basic_memory.repository.semantic_errors import SemanticVectorIndexExtensionError +from basic_memory.repository.semantic_vector_index import ( + SEMANTIC_VECTOR_INDEX_ENTRY_POINT_GROUP, + SemanticVectorIndex, + VectorIndexScope, +) + + +class SemanticVectorIndexFactory(Protocol): + """Factory signature exposed to separately installed extension packages.""" + + def __call__( + self, + *, + scope: VectorIndexScope, + app_config: BasicMemoryConfig, + ) -> SemanticVectorIndex: ... + + +def resolve_semantic_vector_index_name( + app_config: BasicMemoryConfig, + database_backend: DatabaseBackend, +) -> str: + """Resolve the effective index while preserving sqlite-vec for local SQLite.""" + if database_backend == DatabaseBackend.SQLITE: + return "sqlite-vec" + return app_config.semantic_vector_index.strip().lower() + + +def semantic_embedding_identity(provider: EmbeddingProvider) -> str: + """Return the same model identity used by manifest invalidation.""" + return f"{type(provider).__name__}:{embedding_provider_identity(provider)}" + + +def _database_namespace(app_config: BasicMemoryConfig) -> str: + """Derive a stable, credential-free namespace from the authoritative database.""" + if app_config.database_url: + url = make_url(app_config.database_url) + locator = "|".join( + [ + url.get_backend_name(), + url.host or "", + str(url.port or ""), + url.database or "", + ] + ) + else: + locator = str((app_config.data_dir_path / "memory.db").resolve()) + digest = hashlib.sha256(locator.encode("utf-8")).hexdigest()[:24] + return f"basic-memory-{digest}" + + +def build_vector_index_scope( + app_config: BasicMemoryConfig, + provider: EmbeddingProvider, + project_id: int, +) -> VectorIndexScope: + """Build the explicit isolation contract handed to every vector adapter.""" + return VectorIndexScope( + namespace=_database_namespace(app_config), + project_id=project_id, + embedding_identity=semantic_embedding_identity(provider), + dimensions=provider.dimensions, + ) + + +def _load_extension_factory(name: str) -> SemanticVectorIndexFactory: + matches = list(entry_points(group=SEMANTIC_VECTOR_INDEX_ENTRY_POINT_GROUP, name=name)) + if not matches: + raise SemanticVectorIndexExtensionError( + f"Semantic vector index '{name}' is configured but no extension is installed. " + f"Install a package that provides the '{name}' entry point in " + f"'{SEMANTIC_VECTOR_INDEX_ENTRY_POINT_GROUP}'." + ) + if len(matches) > 1: + providers = ", ".join(sorted(entry_point.value for entry_point in matches)) + raise SemanticVectorIndexExtensionError( + f"Multiple semantic vector index extensions provide '{name}': {providers}." + ) + + loaded = matches[0].load() + if not callable(loaded): + raise SemanticVectorIndexExtensionError( + f"Semantic vector index entry point '{name}' must load a callable factory." + ) + return loaded + + +def create_semantic_vector_index( + *, + session_maker: async_sessionmaker[AsyncSession], + project_id: int, + app_config: BasicMemoryConfig, + database_backend: DatabaseBackend, + embedding_provider: EmbeddingProvider, +) -> tuple[str, SemanticVectorIndex]: + """Create the selected built-in adapter or load one external extension.""" + name = resolve_semantic_vector_index_name(app_config, database_backend) + scope = build_vector_index_scope(app_config, embedding_provider, project_id) + + if name == "sqlite-vec": + from basic_memory.repository.sqlite_vec_index import SQLiteVecIndex + + return name, SQLiteVecIndex(session_maker, scope) + if name == "pgvector": + from basic_memory.repository.pgvector_index import PgVectorIndex + + return name, PgVectorIndex(session_maker, scope) + + factory = _load_extension_factory(name) + index = factory(scope=scope, app_config=app_config) + if not isinstance(index, SemanticVectorIndex): + raise SemanticVectorIndexExtensionError( + f"Semantic vector index extension '{name}' returned an incompatible adapter." + ) + if index.scope != scope: + raise SemanticVectorIndexExtensionError( + f"Semantic vector index extension '{name}' returned an adapter for the wrong scope." + ) + return name, index diff --git a/src/basic_memory/repository/semantic_vector_sync.py b/src/basic_memory/repository/semantic_vector_sync.py index 851d9a985..7c5d3a871 100644 --- a/src/basic_memory/repository/semantic_vector_sync.py +++ b/src/basic_memory/repository/semantic_vector_sync.py @@ -63,6 +63,8 @@ class PreparedEntityVectorSync: remaining_jobs_after_shard: int = 0 prepare_seconds: float = 0.0 queue_start: float | None = None + delete_entity_vectors: bool = False + stale_chunk_ids: list[int] = field(default_factory=list) @dataclass(frozen=True, slots=True) @@ -152,6 +154,8 @@ class VectorChunkState: entity_fingerprint: str embedding_model: str has_embedding: bool + vector_index: str = "" + embedding_status: str = "" def plan_entity_vector_shard( @@ -621,9 +625,8 @@ def prepare_window_existing_rows_sql(placeholders: str) -> str: """Build SQL for existing chunk and embedding rows in one prepare window.""" return ( "SELECT c.entity_id, c.id, c.chunk_key, c.source_hash, c.entity_fingerprint, " - "c.embedding_model, (e.chunk_id IS NOT NULL) AS has_embedding " + "c.embedding_model, c.vector_index, c.embedding_status " "FROM search_vector_chunks c " - "LEFT JOIN search_vector_embeddings e ON e.chunk_id = c.id " f"WHERE c.project_id = :project_id AND c.entity_id IN ({placeholders}) " "ORDER BY c.entity_id ASC, c.chunk_key ASC" ) @@ -651,7 +654,12 @@ async def fetch_prepare_window_existing_rows( source_hash=str(row["source_hash"]), entity_fingerprint=str(row["entity_fingerprint"]), embedding_model=str(row["embedding_model"]), - has_embedding=bool(row["has_embedding"]), + vector_index=str(row["vector_index"]), + embedding_status=str(row["embedding_status"]), + has_embedding=( + str(row["embedding_status"]) == "ready" + and str(row["vector_index"]) == repository._semantic_vector_index_name + ), ) ) return grouped_rows @@ -718,6 +726,14 @@ async def prepare_entity_vector_jobs_window( plan, ) await session.commit() + + for index, _plan in mutation_plans: + prepared = prepared_by_index[index] + if isinstance(prepared, PreparedEntityVectorSync): + try: + await repository._finalize_prepared_vector_deletions(prepared) + except Exception as exc: + prepared_by_index[index] = exc except Exception as exc: # The mutation plans share one transaction, so a failed write # invalidates every entity whose result depended on that commit. @@ -761,7 +777,8 @@ async def prepare_entity_vector_jobs_prefetched( await repository._prepare_vector_session(session) prepared = await apply_entity_vector_prepare_plan(repository, session, planned) await session.commit() - return prepared + await repository._finalize_prepared_vector_deletions(prepared) + return prepared def plan_entity_vector_jobs_prefetched( @@ -795,6 +812,7 @@ def delete_entity_chunks() -> DeleteEntityVectorPreparePlan: current_entity_fingerprint = repository._build_entity_fingerprint(chunk_records) current_embedding_model = repository._embedding_model_key() + current_vector_index = repository._semantic_vector_index_name existing_by_key = {row.chunk_key: row for row in existing_rows} incoming_chunk_keys = {record["chunk_key"] for record in chunk_records} stale_ids = [ @@ -815,6 +833,7 @@ def delete_entity_chunks() -> DeleteEntityVectorPreparePlan: and all( row.entity_fingerprint == current_entity_fingerprint and row.embedding_model == current_embedding_model + and row.vector_index in {"", current_vector_index} for row in existing_rows ) ) @@ -843,8 +862,14 @@ def delete_entity_chunks() -> DeleteEntityVectorPreparePlan: same_source_hash = current.source_hash == record["source_hash"] same_entity_fingerprint = current.entity_fingerprint == current_entity_fingerprint same_embedding_model = current.embedding_model == current_embedding_model - - if same_source_hash and current.id not in orphan_ids and same_embedding_model: + same_vector_index = current.vector_index in {"", current_vector_index} + + if ( + same_source_hash + and current.id not in orphan_ids + and same_embedding_model + and same_vector_index + ): if not same_entity_fingerprint: metadata_update_ids.append(current.id) skipped_chunks_count += 1 @@ -896,6 +921,7 @@ async def apply_entity_vector_prepare_plan( source_rows_count=plan.source_rows_count, embedding_jobs=[], prepare_seconds=time.perf_counter() - plan.prepare_start, + delete_entity_vectors=True, ) timestamp_expr = repository._timestamp_now_expr() @@ -944,6 +970,7 @@ async def apply_entity_vector_prepare_plan( remaining_jobs_after_shard=plan.shard_plan.remaining_jobs_after_shard, prepare_seconds=prepare_seconds, queue_start=time.perf_counter(), + stale_chunk_ids=plan.stale_ids, ) @@ -967,6 +994,8 @@ async def upsert_scheduled_chunk_records( current.source_hash != record["source_hash"] or current.entity_fingerprint != entity_fingerprint or current.embedding_model != embedding_model + or current.vector_index != repository._semantic_vector_index_name + or not current.has_embedding ): await session.execute( text( @@ -974,6 +1003,8 @@ async def upsert_scheduled_chunk_records( "SET chunk_text = :chunk_text, source_hash = :source_hash, " "entity_fingerprint = :entity_fingerprint, " "embedding_model = :embedding_model, " + "vector_index = :vector_index, " + "embedding_status = 'pending', " f"updated_at = {timestamp_expr} " "WHERE id = :id" ), @@ -983,6 +1014,7 @@ async def upsert_scheduled_chunk_records( "source_hash": record["source_hash"], "entity_fingerprint": entity_fingerprint, "embedding_model": embedding_model, + "vector_index": repository._semantic_vector_index_name, }, ) embedding_jobs.append((current.id, record["chunk_text"])) @@ -993,10 +1025,11 @@ async def upsert_scheduled_chunk_records( "INSERT INTO search_vector_chunks (" "entity_id, project_id, chunk_key, chunk_text, source_hash, " "entity_fingerprint, embedding_model, updated_at" + ", vector_index, embedding_status" ") VALUES (" ":entity_id, :project_id, :chunk_key, :chunk_text, :source_hash, " ":entity_fingerprint, :embedding_model, " - f"{timestamp_expr}" + f"{timestamp_expr}, :vector_index, 'pending'" ") RETURNING id" ), { @@ -1007,6 +1040,7 @@ async def upsert_scheduled_chunk_records( "source_hash": record["source_hash"], "entity_fingerprint": entity_fingerprint, "embedding_model": embedding_model, + "vector_index": repository._semantic_vector_index_name, }, ) embedding_jobs.append((int(inserted.scalar_one()), record["chunk_text"])) @@ -1032,11 +1066,8 @@ async def flush_embedding_jobs( raise RuntimeError("Embedding provider returned an unexpected number of vectors.") write_start = time.perf_counter() - async with db.scoped_session(repository.session_maker) as session: - await repository._prepare_vector_session(session) - write_jobs = [(job.chunk_row_id, job.chunk_text) for job in flush_jobs] - await repository._write_embeddings(session, write_jobs, embeddings) - await session.commit() + write_jobs = [(job.chunk_row_id, job.chunk_text) for job in flush_jobs] + await repository._persist_embeddings(write_jobs, embeddings) write_seconds = time.perf_counter() - write_start flush_size = len(flush_jobs) diff --git a/src/basic_memory/repository/sqlite_search_repository.py b/src/basic_memory/repository/sqlite_search_repository.py index e918e6af6..bb97b293c 100644 --- a/src/basic_memory/repository/sqlite_search_repository.py +++ b/src/basic_memory/repository/sqlite_search_repository.py @@ -1,7 +1,6 @@ """SQLite FTS5-based search repository implementation.""" import asyncio -import json import re from contextlib import asynccontextmanager from datetime import datetime @@ -20,7 +19,6 @@ CREATE_SQLITE_SEARCH_VECTOR_CHUNKS, CREATE_SQLITE_SEARCH_VECTOR_CHUNKS_PROJECT_ENTITY, CREATE_SQLITE_SEARCH_VECTOR_CHUNKS_UNIQUE, - create_sqlite_search_vector_embeddings, ) from basic_memory.repository.embedding_provider import EmbeddingProvider from basic_memory.repository.embedding_provider_factory import create_embedding_provider @@ -29,6 +27,9 @@ from basic_memory.repository.search_repository_base import SearchRepositoryBase from basic_memory.repository.metadata_filters import parse_metadata_filters, build_sqlite_json_path from basic_memory.repository.semantic_errors import SemanticDependenciesMissingError +from basic_memory.repository.semantic_vector_index import SemanticVectorIndex +from basic_memory.repository.semantic_vector_index_factory import build_vector_index_scope +from basic_memory.repository.sqlite_vec_index import SQLiteVecIndex from basic_memory.schemas.search import SearchItemType, SearchRetrievalMode @@ -48,6 +49,8 @@ def __init__( project_id: int, app_config: BasicMemoryConfig | None = None, embedding_provider: EmbeddingProvider | None = None, + vector_index_name: str | None = None, + vector_index: SemanticVectorIndex | None = None, ): super().__init__(session_maker, project_id) self._entity_columns: set[str] | None = None @@ -59,6 +62,7 @@ def __init__( self._app_config.semantic_embedding_sync_batch_size ) self._embedding_provider = embedding_provider + self._semantic_vector_index_name = vector_index_name or "sqlite-vec" self._sqlite_vec_load_lock = asyncio.Lock() self._sqlite_prepare_write_lock = asyncio.Lock() self._vector_tables_initialized = False @@ -71,6 +75,14 @@ def __init__( self._embedding_provider = create_embedding_provider(self._app_config) if self._embedding_provider is not None: self._vector_dimensions = self._embedding_provider.dimensions + self._semantic_vector_index = vector_index or SQLiteVecIndex( + session_maker, + build_vector_index_scope( + self._app_config, + self._embedding_provider, + project_id, + ), + ) async def _get_entity_columns(self) -> set[str]: if self._entity_columns is None: @@ -441,6 +453,16 @@ async def _ensure_sqlite_vec_loaded(self, session) -> None: async def _ensure_vector_tables(self) -> None: self._assert_semantic_available() + if not hasattr(self, "_semantic_vector_index"): + assert self._embedding_provider is not None + self._semantic_vector_index = SQLiteVecIndex( + self.session_maker, + build_vector_index_scope( + self._app_config, + self._embedding_provider, + self.project_id, + ), + ) if self._vector_tables_initialized: return @@ -463,6 +485,8 @@ async def _ensure_vector_tables(self) -> None: "source_hash", "entity_fingerprint", "embedding_model", + "vector_index", + "embedding_status", "updated_at", } schema_mismatch = bool(chunks_columns) and set(chunks_columns) != expected_columns @@ -474,6 +498,8 @@ async def _ensure_vector_tables(self) -> None: logger.warning("search_vector_chunks schema mismatch, recreating vector tables") await session.execute(text("DROP TABLE IF EXISTS search_vector_embeddings")) await session.execute(text("DROP TABLE IF EXISTS search_vector_chunks")) + if isinstance(self._semantic_vector_index, SQLiteVecIndex): + self._semantic_vector_index.invalidate_initialization() await session.execute(CREATE_SQLITE_SEARCH_VECTOR_CHUNKS) await session.execute(CREATE_SQLITE_SEARCH_VECTOR_CHUNKS_PROJECT_ENTITY) @@ -484,25 +510,10 @@ async def _ensure_vector_tables(self) -> None: # Outcome: remove disposable derived data so chunk/vector schema is deterministic. await session.execute(text("DROP TABLE IF EXISTS search_vector_index")) - vector_sql_result = await session.execute( - text( - "SELECT sql FROM sqlite_master " - "WHERE type = 'table' AND name = 'search_vector_embeddings'" - ) - ) - vector_sql = vector_sql_result.scalar() - expected_dimension_sql = f"float[{self._vector_dimensions}]" - - if vector_sql and expected_dimension_sql not in vector_sql: - logger.warning( - f"Embedding dimension mismatch (expected {self._vector_dimensions}), " - "recreating search_vector_embeddings" - ) - await session.execute(text("DROP TABLE IF EXISTS search_vector_embeddings")) - - await session.execute(create_sqlite_search_vector_embeddings(self._vector_dimensions)) await session.commit() + await self._semantic_vector_index.initialize() + logger.debug(f"SQLite vector tables ready (dimensions={self._vector_dimensions})") self._vector_tables_initialized = True @@ -519,82 +530,14 @@ async def _run_vector_query( query_embedding: list[float], candidate_limit: int, ) -> list[dict]: - # Constraint: sqlite-vec enforces k <= 4096 for knn queries - vector_k = min(candidate_limit, self.SQLITE_VEC_MAX_K) - query_embedding_json = json.dumps(query_embedding) - vector_result = await session.execute( - text( - "WITH vector_matches AS (" - " SELECT rowid, distance " - " FROM search_vector_embeddings " - " WHERE embedding MATCH :query_embedding " - " AND k = :vector_k" - ") " - "SELECT c.entity_id, c.chunk_key, c.chunk_text, vector_matches.distance AS best_distance " - "FROM vector_matches " - "JOIN search_vector_chunks c ON c.id = vector_matches.rowid " - "WHERE c.project_id = :project_id " - "ORDER BY best_distance ASC " - "LIMIT :candidate_limit" - ), - { - "query_embedding": query_embedding_json, - "project_id": self.project_id, - "vector_k": vector_k, - "candidate_limit": candidate_limit, - }, - ) - return [dict(row) for row in vector_result.mappings().all()] - - async def _write_embeddings( - self, - session: AsyncSession, - jobs: list[tuple[int, str]], - embeddings: list[list[float]], - ) -> None: - rowids = [row_id for row_id, _ in jobs] - delete_params = {f"rowid_{idx}": rowid for idx, rowid in enumerate(rowids)} - delete_placeholders = ", ".join(f":rowid_{idx}" for idx in range(len(rowids))) - await session.execute( - text(f"DELETE FROM search_vector_embeddings WHERE rowid IN ({delete_placeholders})"), - delete_params, - ) - - insert_rows = [ - {"rowid": row_id, "embedding": json.dumps(embedding)} - for (row_id, _), embedding in zip(jobs, embeddings, strict=True) - ] - await session.execute( - text( - "INSERT INTO search_vector_embeddings (rowid, embedding) " - "VALUES (:rowid, :embedding)" - ), - insert_rows, - ) + return await super()._run_vector_query(session, query_embedding, candidate_limit) async def _delete_entity_chunks( self, session: AsyncSession, entity_id: int, ) -> None: - # sqlite-vec has no CASCADE — must delete embeddings before chunks - await session.execute( - text( - "DELETE FROM search_vector_embeddings " - "WHERE rowid IN (" - "SELECT id FROM search_vector_chunks " - "WHERE project_id = :project_id AND entity_id = :entity_id" - ")" - ), - {"project_id": self.project_id, "entity_id": entity_id}, - ) - await session.execute( - text( - "DELETE FROM search_vector_chunks " - "WHERE project_id = :project_id AND entity_id = :entity_id" - ), - {"project_id": self.project_id, "entity_id": entity_id}, - ) + await super()._delete_entity_chunks(session, entity_id) async def _delete_stale_chunks( self, @@ -602,47 +545,7 @@ async def _delete_stale_chunks( stale_ids: list[int], entity_id: int, ) -> None: - stale_params = { - "project_id": self.project_id, - "entity_id": entity_id, - **{f"row_{idx}": row_id for idx, row_id in enumerate(stale_ids)}, - } - stale_placeholders = ", ".join(f":row_{idx}" for idx in range(len(stale_ids))) - await session.execute( - text(f"DELETE FROM search_vector_embeddings WHERE rowid IN ({stale_placeholders})"), - stale_params, - ) - await session.execute( - text( - "DELETE FROM search_vector_chunks " - f"WHERE id IN ({stale_placeholders}) " - "AND project_id = :project_id AND entity_id = :entity_id" - ), - stale_params, - ) - - async def delete_project_vector_rows(self) -> None: - """Delete all vector rows for this project on a sqlite-vec-enabled connection.""" - await self._ensure_vector_tables() - - async with db.scoped_session(self.session_maker) as session: - await self._ensure_sqlite_vec_loaded(session) - - # Constraint: sqlite-vec stores embeddings separately with no cascade delete. - # Why: full rebuild must clear embeddings before chunk rows or stale vectors remain. - # Outcome: the next sync recreates the project's derived vectors from scratch. - await session.execute( - text( - "DELETE FROM search_vector_embeddings WHERE rowid IN (" - "SELECT id FROM search_vector_chunks WHERE project_id = :project_id)" - ), - {"project_id": self.project_id}, - ) - await session.execute( - text("DELETE FROM search_vector_chunks WHERE project_id = :project_id"), - {"project_id": self.project_id}, - ) - await session.commit() + await super()._delete_stale_chunks(session, stale_ids, entity_id) async def drop_vector_tables(self) -> None: """Drop SQLite vector tables on a sqlite-vec-enabled connection.""" @@ -663,38 +566,6 @@ async def drop_vector_tables(self) -> None: await session.commit() self._vector_tables_initialized = False - async def delete_stale_vector_rows(self) -> None: - """Delete vector rows whose source entities no longer exist.""" - await self._ensure_vector_tables() - - async with db.scoped_session(self.session_maker) as session: - await self._ensure_sqlite_vec_loaded(session) - - stale_entity_filter = ( - "entity_id NOT IN (SELECT id FROM entity WHERE project_id = :project_id)" - ) - params = {"project_id": self.project_id} - - # Trigger: deleted entities left behind derived vector rows. - # Why: sqlite-vec does not provide cascade cleanup from our chunk table. - # Outcome: stale vector state disappears before coverage stats or reindex runs. - await session.execute( - text( - "DELETE FROM search_vector_embeddings WHERE rowid IN (" - "SELECT id FROM search_vector_chunks " - f"WHERE project_id = :project_id AND {stale_entity_filter})" - ), - params, - ) - await session.execute( - text( - "DELETE FROM search_vector_chunks " - f"WHERE project_id = :project_id AND {stale_entity_filter}" - ), - params, - ) - await session.commit() - def _distance_to_similarity(self, distance: float) -> float: """Convert L2 distance to cosine similarity for normalized embeddings. @@ -715,15 +586,8 @@ async def _prepare_entity_write_scope(self): yield def _prepare_window_existing_rows_sql(self, placeholders: str) -> str: - """SQLite sqlite-vec stores embeddings by rowid rather than chunk_id.""" - return ( - "SELECT c.entity_id, c.id, c.chunk_key, c.source_hash, c.entity_fingerprint, " - "c.embedding_model, (e.rowid IS NOT NULL) AS has_embedding " - "FROM search_vector_chunks c " - "LEFT JOIN search_vector_embeddings e ON e.rowid = c.id " - f"WHERE c.project_id = :project_id AND c.entity_id IN ({placeholders}) " - "ORDER BY c.entity_id ASC, c.chunk_key ASC" - ) + """Use the authoritative SQL manifest for adapter-independent readiness.""" + return super()._prepare_window_existing_rows_sql(placeholders) # ------------------------------------------------------------------ # Index / bulk index overrides (FTS-only, no vector side-effects) diff --git a/src/basic_memory/repository/sqlite_vec_index.py b/src/basic_memory/repository/sqlite_vec_index.py new file mode 100644 index 000000000..cc8cfeb5d --- /dev/null +++ b/src/basic_memory/repository/sqlite_vec_index.py @@ -0,0 +1,286 @@ +"""Built-in sqlite-vec implementation of the semantic vector index contract.""" + +from __future__ import annotations + +import asyncio +import json +from collections.abc import Sequence + +from loguru import logger +from sqlalchemy import text +from sqlalchemy.exc import OperationalError as SAOperationalError +from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker + +from basic_memory import db +from basic_memory.models.search import create_sqlite_search_vector_embeddings +from basic_memory.repository.semantic_errors import SemanticDependenciesMissingError +from basic_memory.repository.semantic_vector_index import ( + SemanticVectorIndex, + VectorIndexScope, + VectorKey, + VectorMatch, + VectorRecord, + validate_query_dimensions, + validate_vector_dimensions, +) + + +SQLITE_VEC_MAX_K = 4096 + + +class SQLiteVecIndex(SemanticVectorIndex): + """Persist and query semantic vectors in SQLite with sqlite-vec.""" + + def __init__( + self, + session_maker: async_sessionmaker[AsyncSession], + scope: VectorIndexScope, + ) -> None: + self._session_maker = session_maker + self.scope = scope + self._initialized = False + self._initialize_lock = asyncio.Lock() + self._load_lock = asyncio.Lock() + + def invalidate_initialization(self) -> None: + """Require the next operation to revalidate lazily created vec storage.""" + self._initialized = False + + async def _ensure_loaded(self, session: AsyncSession) -> None: + try: + await session.execute(text("SELECT vec_version()")) + return + except SAOperationalError: + pass + + try: + import sqlite_vec + except ImportError as exc: + raise SemanticDependenciesMissingError( + "sqlite-vec package is missing. Install/update basic-memory to include " + "semantic dependencies: pip install -U basic-memory" + ) from exc + + async with self._load_lock: + try: + await session.execute(text("SELECT vec_version()")) + return + except SAOperationalError: + pass + + connection = await session.connection() + raw_connection = await connection.get_raw_connection() + driver_connection = raw_connection.driver_connection + if not hasattr(driver_connection, "enable_load_extension"): + raise SemanticDependenciesMissingError( + "This Python build does not support SQLite extension loading " + "(no enable_load_extension on sqlite3.Connection). Reinstall " + "basic-memory under uv-managed or Homebrew Python, or disable " + "semantic search." + ) + await driver_connection.enable_load_extension(True) + await driver_connection.load_extension(sqlite_vec.loadable_path()) + await driver_connection.enable_load_extension(False) + await session.execute(text("SELECT vec_version()")) + + async def initialize(self) -> None: + if self._initialized: + return + async with self._initialize_lock: + if self._initialized: + return + + async with db.scoped_session(self._session_maker) as session: + await self._ensure_loaded(session) + result = await session.execute( + text( + "SELECT sql FROM sqlite_master " + "WHERE type = 'table' AND name = 'search_vector_embeddings'" + ) + ) + vector_sql = result.scalar() + expected_dimensions = f"float[{self.scope.dimensions}]" + dimensions_changed = bool(vector_sql and expected_dimensions not in vector_sql) + if dimensions_changed: + logger.warning( + "Embedding dimension mismatch (expected {dimensions}); " + "recreating sqlite-vec storage", + dimensions=self.scope.dimensions, + ) + await session.execute(text("DROP TABLE IF EXISTS search_vector_embeddings")) + + await session.execute(create_sqlite_search_vector_embeddings(self.scope.dimensions)) + if dimensions_changed: + await session.execute( + text( + "UPDATE search_vector_chunks SET embedding_status = 'pending' " + "WHERE vector_index = 'sqlite-vec'" + ) + ) + await session.commit() + self._initialized = True + + async def _rowids_by_key( + self, + session: AsyncSession, + keys: Sequence[VectorKey], + ) -> dict[VectorKey, int]: + if not keys: + return {} + params: dict[str, object] = {"project_id": self.scope.project_id} + predicates: list[str] = [] + for index, key in enumerate(keys): + params[f"entity_id_{index}"] = key.entity_id + params[f"chunk_key_{index}"] = key.chunk_key + predicates.append( + f"(entity_id = :entity_id_{index} AND chunk_key = :chunk_key_{index})" + ) + result = await session.execute( + text( + "SELECT id, entity_id, chunk_key FROM search_vector_chunks " + "WHERE project_id = :project_id AND (" + " OR ".join(predicates) + ")" + ), + params, + ) + return { + VectorKey(entity_id=int(row["entity_id"]), chunk_key=str(row["chunk_key"])): int( + row["id"] + ) + for row in result.mappings().all() + } + + async def upsert(self, records: Sequence[VectorRecord]) -> None: + if not records: + return + validate_vector_dimensions(self.scope, records) + await self.initialize() + async with db.scoped_session(self._session_maker) as session: + await self._ensure_loaded(session) + rowids_by_key = await self._rowids_by_key(session, [record.key for record in records]) + missing = [record.key for record in records if record.key not in rowids_by_key] + if missing: + raise RuntimeError(f"Vector manifest rows are missing for keys: {missing!r}") + + rowids = [rowids_by_key[record.key] for record in records] + params = {f"rowid_{index}": rowid for index, rowid in enumerate(rowids)} + placeholders = ", ".join(f":rowid_{index}" for index in range(len(rowids))) + await session.execute( + text(f"DELETE FROM search_vector_embeddings WHERE rowid IN ({placeholders})"), + params, + ) + await session.execute( + text( + "INSERT INTO search_vector_embeddings (rowid, embedding) " + "VALUES (:rowid, :embedding)" + ), + [ + { + "rowid": rowids_by_key[record.key], + "embedding": json.dumps(record.values), + } + for record in records + ], + ) + await session.commit() + + async def delete(self, keys: Sequence[VectorKey]) -> None: + if not keys: + return + await self.initialize() + async with db.scoped_session(self._session_maker) as session: + await self._ensure_loaded(session) + rowids = list((await self._rowids_by_key(session, keys)).values()) + if rowids: + params = {f"rowid_{index}": rowid for index, rowid in enumerate(rowids)} + placeholders = ", ".join(f":rowid_{index}" for index in range(len(rowids))) + await session.execute( + text(f"DELETE FROM search_vector_embeddings WHERE rowid IN ({placeholders})"), + params, + ) + await session.commit() + + async def delete_entity(self, entity_id: int) -> None: + await self.initialize() + async with db.scoped_session(self._session_maker) as session: + await self._ensure_loaded(session) + await session.execute( + text( + "DELETE FROM search_vector_embeddings WHERE rowid IN (" + "SELECT id FROM search_vector_chunks " + "WHERE project_id = :project_id AND entity_id = :entity_id)" + ), + {"project_id": self.scope.project_id, "entity_id": entity_id}, + ) + await session.commit() + + async def delete_orphans(self, _live_keys: Sequence[VectorKey]) -> None: + """Remove sqlite-vec rows absent from the current ready manifest scope.""" + await self.initialize() + async with db.scoped_session(self._session_maker) as session: + await self._ensure_loaded(session) + await session.execute( + text( + "DELETE FROM search_vector_embeddings WHERE rowid IN (" + "SELECT id FROM search_vector_chunks " + "WHERE project_id = :project_id AND NOT (" + "vector_index = 'sqlite-vec' " + "AND embedding_model = :embedding_identity " + "AND embedding_status = 'ready'))" + ), + { + "project_id": self.scope.project_id, + "embedding_identity": self.scope.embedding_identity, + }, + ) + await session.commit() + + async def search( + self, + query: Sequence[float], + *, + limit: int, + ) -> list[VectorMatch]: + if not query or limit <= 0: + return [] + validate_query_dimensions(self.scope, query) + await self.initialize() + vector_k = min(limit, SQLITE_VEC_MAX_K) + async with db.scoped_session(self._session_maker) as session: + await self._ensure_loaded(session) + result = await session.execute( + text( + "WITH vector_matches AS (" + " SELECT rowid, distance FROM search_vector_embeddings " + " WHERE 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 " + "WHERE c.project_id = :project_id " + "AND c.vector_index = 'sqlite-vec' " + "AND c.embedding_status = 'ready' " + "AND c.embedding_model = :embedding_identity " + "ORDER BY vector_matches.distance ASC, " + "c.entity_id ASC, c.chunk_key ASC LIMIT :limit" + ), + { + "query": json.dumps(list(query)), + "vector_k": vector_k, + "project_id": self.scope.project_id, + "embedding_identity": self.scope.embedding_identity, + "limit": limit, + }, + ) + return [ + VectorMatch( + key=VectorKey( + entity_id=int(row["entity_id"]), + chunk_key=str(row["chunk_key"]), + ), + similarity=max( + 0.0, + min(1.0, 1.0 - (float(row["distance"]) ** 2) / 2.0), + ), + ) + for row in result.mappings().all() + ] diff --git a/src/basic_memory/services/project_service.py b/src/basic_memory/services/project_service.py index f112c92e5..71589032e 100644 --- a/src/basic_memory/services/project_service.py +++ b/src/basic_memory/services/project_service.py @@ -11,12 +11,16 @@ from loguru import logger from sqlalchemy import text -from sqlalchemy.exc import OperationalError as SAOperationalError from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker from basic_memory import db from basic_memory.models import Project from basic_memory.repository.project_repository import ProjectRepository +from basic_memory.repository.embedding_provider_factory import create_embedding_provider +from basic_memory.repository.semantic_vector_index_factory import ( + resolve_semantic_vector_index_name, + semantic_embedding_identity, +) from basic_memory.schemas import ( ActivityMetrics, EmbeddingStatus, @@ -1040,24 +1044,26 @@ async def get_embedding_status(self, project_id: int) -> EmbeddingStatus: query_prefix_set = bool(config.semantic_embedding_query_prefix) is_postgres = config.database_backend == DatabaseBackend.POSTGRES + vector_index = resolve_semantic_vector_index_name(config, config.database_backend) + embedding_identity = semantic_embedding_identity(create_embedding_provider(config)) - # --- Check vector table existence --- - # Both search_vector_chunks and search_vector_embeddings must exist - # for the detailed stats queries (JOINs between them) to work. + # --- Check vector manifest existence --- + # The SQL manifest is authoritative even when vector values live outside + # the database, so status never probes backend-specific storage tables. if is_postgres: table_check_sql = text( "SELECT COUNT(*) FROM information_schema.tables " - "WHERE table_name IN ('search_vector_chunks', 'search_vector_embeddings')" + "WHERE table_name = 'search_vector_chunks'" ) else: table_check_sql = text( "SELECT COUNT(*) FROM sqlite_master " - "WHERE type = 'table' AND name IN ('search_vector_chunks', 'search_vector_embeddings')" + "WHERE type = 'table' AND name = 'search_vector_chunks'" ) async with db.scoped_session(self.session_maker) as session: table_result = await self.repository.execute_query(session, table_check_sql, {}) - vector_tables_exist = (table_result.scalar() or 0) == 2 + vector_tables_exist = (table_result.scalar() or 0) == 1 if not vector_tables_exist: # Count distinct entities in search index for the recommendation message @@ -1081,19 +1087,17 @@ async def get_embedding_status(self, project_id: int) -> EmbeddingStatus: total_indexed_entities=total_indexed_entities, vector_tables_exist=False, reindex_recommended=True, - reindex_reason=("Vector tables not initialized — run: bm reindex --embeddings"), + reindex_reason=( + "Vector manifest not initialized — run: bm reindex --embeddings" + ), ) - # --- Count queries (tables exist) --- + # --- Count queries (manifest exists) --- # Filter by entity existence to exclude stale rows from deleted entities # that remain in derived search tables (search_index, search_vector_chunks) entity_exists = ( "AND entity_id IN (SELECT id FROM entity WHERE project_id = :project_id)" ) - # Same filter for aliased chunks table (used in JOIN queries below) - chunk_entity_exists = ( - "AND c.entity_id IN (SELECT id FROM entity WHERE project_id = :project_id)" - ) si_result = await self.repository.execute_query( session, @@ -1105,105 +1109,56 @@ async def get_embedding_status(self, project_id: int) -> EmbeddingStatus: ) total_indexed_entities = si_result.scalar() or 0 - try: - chunks_result = await self.repository.execute_query( - session, - text( - "SELECT COUNT(*) FROM search_vector_chunks " - f"WHERE project_id = :project_id {entity_exists}" - ), - {"project_id": project_id}, - ) - total_chunks = chunks_result.scalar() or 0 - - entities_with_chunks_result = await self.repository.execute_query( - session, - text( - "SELECT COUNT(DISTINCT entity_id) FROM search_vector_chunks " - f"WHERE project_id = :project_id {entity_exists}" - ), - {"project_id": project_id}, - ) - total_entities_with_chunks = entities_with_chunks_result.scalar() or 0 - - # Embeddings count — join pattern differs between SQLite and Postgres - if is_postgres: - embeddings_sql = text( - "SELECT COUNT(*) FROM search_vector_chunks c " - "JOIN search_vector_embeddings e ON e.chunk_id = c.id " - f"WHERE c.project_id = :project_id {chunk_entity_exists}" - ) - else: - embeddings_sql = text( - "SELECT COUNT(*) FROM search_vector_chunks c " - "JOIN search_vector_embeddings e ON e.rowid = c.id " - f"WHERE c.project_id = :project_id {chunk_entity_exists}" - ) + manifest_params = { + "project_id": project_id, + "vector_index": vector_index, + "embedding_identity": embedding_identity, + } + current_ready = ( + "vector_index = :vector_index " + "AND embedding_model = :embedding_identity " + "AND embedding_status = 'ready'" + ) - # The embeddings/orphan JOINs read search_vector_embeddings, a vec0 - # virtual table. On SQLite that table is only visible on a connection - # that loaded sqlite-vec, so route these through scalar_vec_query which - # loads the extension first. Postgres has no per-connection extension - # and uses the bare pooled session. - async def _vec_scalar(vec_sql) -> int: - if is_postgres: - result = await self.repository.execute_query( - session, vec_sql, {"project_id": project_id} - ) - return result.scalar() or 0 - count = await self.repository.scalar_vec_query( - session, vec_sql, {"project_id": project_id} - ) - # Trigger: sqlite-vec genuinely can't load on this Python build. - # Why: without the extension the vec0 JOIN can't run at all. - # Outcome: raise the canonical error so the except block emits the - # true "sqlite-vec unavailable" fallback instead of reporting 0. - if count is None: - raise SAOperationalError( - str(vec_sql), {}, Exception("no such module: vec0") - ) - return count + chunks_result = await self.repository.execute_query( + session, + text( + "SELECT COUNT(*) FROM search_vector_chunks " + f"WHERE project_id = :project_id {entity_exists}" + ), + {"project_id": project_id}, + ) + total_chunks = chunks_result.scalar() or 0 - total_embeddings = await _vec_scalar(embeddings_sql) + entities_with_chunks_result = await self.repository.execute_query( + session, + text( + "SELECT COUNT(DISTINCT entity_id) FROM search_vector_chunks " + f"WHERE project_id = :project_id {entity_exists}" + ), + {"project_id": project_id}, + ) + total_entities_with_chunks = entities_with_chunks_result.scalar() or 0 - # Orphaned chunks (chunks without embeddings — indicates interrupted indexing) - if is_postgres: - orphan_sql = text( - "SELECT COUNT(*) FROM search_vector_chunks c " - "LEFT JOIN search_vector_embeddings e ON e.chunk_id = c.id " - f"WHERE c.project_id = :project_id AND e.chunk_id IS NULL {chunk_entity_exists}" - ) - else: - orphan_sql = text( - "SELECT COUNT(*) FROM search_vector_chunks c " - "LEFT JOIN search_vector_embeddings e ON e.rowid = c.id " - f"WHERE c.project_id = :project_id AND e.rowid IS NULL {chunk_entity_exists}" - ) - orphaned_chunks = await _vec_scalar(orphan_sql) - except SAOperationalError as exc: - # Trigger: sqlite_master can list vec0 virtual tables even when sqlite-vec - # is not loaded in the current Python runtime. - # Why: project info should degrade gracefully instead of crashing on stats queries. - # Outcome: report vector tables as unavailable and point the user to install the - # missing dependency before rebuilding embeddings. - if is_postgres or "no such module: vec0" not in str(exc).lower(): - raise + embeddings_result = await self.repository.execute_query( + session, + text( + "SELECT COUNT(*) FROM search_vector_chunks " + f"WHERE project_id = :project_id AND {current_ready} {entity_exists}" + ), + manifest_params, + ) + total_embeddings = embeddings_result.scalar() or 0 - return EmbeddingStatus( - semantic_search_enabled=True, - embedding_provider=provider, - embedding_model=model, - embedding_dimensions=dimensions, - embedding_document_prefix_set=document_prefix_set, - embedding_query_prefix_set=query_prefix_set, - total_indexed_entities=total_indexed_entities, - vector_tables_exist=False, - reindex_recommended=True, - reindex_reason=( - "SQLite vector tables exist but sqlite-vec is unavailable in this Python " - "environment — install/update basic-memory, then run: bm reindex --embeddings" - ), - ) + orphaned_result = await self.repository.execute_query( + session, + text( + "SELECT COUNT(*) FROM search_vector_chunks " + f"WHERE project_id = :project_id AND NOT ({current_ready}) {entity_exists}" + ), + manifest_params, + ) + orphaned_chunks = orphaned_result.scalar() or 0 # --- Reindex recommendation logic (priority order) --- reindex_recommended = False @@ -1215,7 +1170,7 @@ async def _vec_scalar(vec_sql) -> int: elif orphaned_chunks > 0: reindex_recommended = True reindex_reason = ( - f"{orphaned_chunks} orphaned chunks found (interrupted indexing) " + f"{orphaned_chunks} chunks need vector indexing (pending or stale) " "— run: bm reindex --embeddings" ) elif total_indexed_entities > total_entities_with_chunks: diff --git a/src/basic_memory/services/search_service.py b/src/basic_memory/services/search_service.py index e16ebbaf9..34b82f3eb 100644 --- a/src/basic_memory/services/search_service.py +++ b/src/basic_memory/services/search_service.py @@ -86,23 +86,13 @@ async def init_search_index(self): async def reindex_all(self, background_tasks: Optional[BackgroundTasks] = None) -> None: """Reindex all content from database.""" - from basic_memory.repository.sqlite_search_repository import SQLiteSearchRepository - logger.info("Starting full reindex") - # Clear and recreate search index + # Trigger: a full search rebuild removes every derived row. + # Why: vector storage may live outside SQL, so cleanup must cross the + # repository boundary before the full-text table is recreated. + # Outcome: built-in and extension indexes follow the same lifecycle. + await self.repository.delete_project_vector_rows() await self.repository.execute_query(text("DROP TABLE IF EXISTS search_index"), params={}) - if isinstance(self.repository, SQLiteSearchRepository): - await self.repository.drop_vector_tables() - else: - await self.repository.execute_query( - text("DROP TABLE IF EXISTS search_vector_embeddings"), params={} - ) - await self.repository.execute_query( - text("DROP TABLE IF EXISTS search_vector_chunks"), params={} - ) - await self.repository.execute_query( - text("DROP TABLE IF EXISTS search_vector_index"), params={} - ) await self.init_search_index() # Reindex all entities @@ -584,6 +574,7 @@ async def reindex_vectors(self, progress_callback=None, force_full: bool = False entity_ids, progress_callback=progress_callback, ) + await self.repository.reconcile_vector_index() stats = { "total_entities": batch_result.entities_total, "embedded": batch_result.entities_synced, @@ -605,20 +596,8 @@ async def _clear_project_vectors_for_full_reindex(self) -> None: we need to clear the derived vector state first to force fresh embeddings. Outcome: the next batch sync recreates every eligible entity's vectors. """ - from basic_memory.repository.sqlite_search_repository import SQLiteSearchRepository - project_id = self.repository.project_id - params = {"project_id": project_id} - - # Constraint: sqlite-vec stores embeddings in a separate rowid table with - # no cascade delete, so embeddings must be removed before chunk rows. - if isinstance(self.repository, SQLiteSearchRepository): - await self.repository.delete_project_vector_rows() - else: - await self.repository.execute_query( - text("DELETE FROM search_vector_chunks WHERE project_id = :project_id"), - params, - ) + await self.repository.delete_project_vector_rows() logger.info("Cleared project vectors for full reindex", project_id=project_id) async def _purge_stale_search_rows(self) -> None: @@ -628,9 +607,6 @@ async def _purge_stale_search_rows(self) -> None: Why: stale rows inflate embedding coverage stats in project info Outcome: search tables only contain rows for entities that still exist """ - from basic_memory.repository.sqlite_search_repository import SQLiteSearchRepository - from sqlalchemy import text - project_id = self.repository.project_id stale_entity_filter = ( "entity_id NOT IN (SELECT id FROM entity WHERE project_id = :project_id)" @@ -645,18 +621,7 @@ async def _purge_stale_search_rows(self) -> None: params, ) - # SQLite vec has no CASCADE — must delete embeddings before chunks - if isinstance(self.repository, SQLiteSearchRepository): - await self.repository.delete_stale_vector_rows() - else: - # Postgres CASCADE handles embedding deletion automatically - await self.repository.execute_query( - text( - f"DELETE FROM search_vector_chunks " - f"WHERE project_id = :project_id AND {stale_entity_filter}" - ), - params, - ) + await self.repository.delete_stale_vector_rows() logger.info("Purged stale search rows for deleted entities", project_id=project_id) diff --git a/test-int/test_embedding_status_vec0.py b/test-int/test_embedding_status_vec0.py index e56d4b0f6..510fe5809 100644 --- a/test-int/test_embedding_status_vec0.py +++ b/test-int/test_embedding_status_vec0.py @@ -1,18 +1,12 @@ -"""Integration regression test for get_embedding_status against a real vec0 table. +"""Integration regression test for manifest-backed status with a real vec0 table. Regression for #658: after a successful `bm reindex --embeddings`, `bm project info` still reported "sqlite-vec is unavailable", "Indexed 0/N", and "Chunks 0", and recommended an unnecessary reindex. -Root cause: get_embedding_status() ran the vec0 JOIN count queries on a bare pooled -ProjectRepository session that never loaded the sqlite-vec extension, so SQLite raised -"no such module: vec0", which the except block mis-reported as "unavailable". - -This test exercises the real failure path: it builds a REAL vec0 virtual table, writes a -real embedding into it via the search repository, then queries get_embedding_status through -a ProjectRepository session that did NOT pre-load the extension (mirroring the bug). The -healthy unit test substitutes a plain regular table for vec0 and therefore does not cover -this path. +Status now reads the authoritative SQL manifest instead of reopening backend-specific +vector storage. This test builds and writes a real vec0 table, then proves a fresh project +status connection can report the ready manifest without loading sqlite-vec. """ import os @@ -26,6 +20,7 @@ from basic_memory.config import BasicMemoryConfig, ConfigManager, DatabaseBackend from basic_memory.repository.entity_repository import EntityRepository from basic_memory.repository.project_repository import ProjectRepository +from basic_memory.repository.semantic_vector_index import VectorKey, VectorRecord from basic_memory.repository.sqlite_search_repository import SQLiteSearchRepository from basic_memory.services.project_service import ProjectService @@ -45,12 +40,7 @@ def _unit_vector(dimensions: int) -> list[float]: @pytest.mark.asyncio async def test_embedding_status_reads_real_vec0_table(engine_factory, test_project, config_manager): - """get_embedding_status must report a populated real vec0 table as healthy. - - Before the fix, the vec0 JOIN ran on a session without sqlite-vec loaded and - raised "no such module: vec0", which the except block mapped to - vector_tables_exist=False + reindex_recommended=True. - """ + """A ready vec0-backed manifest remains healthy on a fresh SQL connection.""" # Trigger: Postgres test matrix executes the same suite. # Why: vec0 + per-connection sqlite-vec loading is SQLite-specific. # Outcome: keep the regression on the backend that can actually hit this path. @@ -121,28 +111,42 @@ async def test_embedding_status_reads_real_vec0_table(engine_factory, test_proje ) await session.commit() - # --- Insert a chunk + a real embedding into the vec0 table --- - # _write_embeddings writes the embedding into the vec0 virtual table keyed by - # rowid == chunk id, exactly like the reindex path. + # --- Insert a pending manifest row, then write the real vec0 value --- + # The pending row commits first because the vector adapter owns a separate + # transaction and must be able to resolve the stable key. async with db.scoped_session(session_maker) as session: - await search_repo._ensure_sqlite_vec_loaded(session) chunk_result = await session.execute( text( "INSERT INTO search_vector_chunks " "(entity_id, project_id, chunk_key, chunk_text, source_hash, " - "entity_fingerprint, embedding_model) " + "entity_fingerprint, embedding_model, vector_index, embedding_status) " "VALUES (:eid, :pid, 'chunk-1', 'vec content', 'hash', " - "'fp-hash', 'bge-small-en-v1.5') " + "'fp-hash', :embedding_model, :vector_index, 'pending') " "RETURNING id" ), - {"eid": entity_id, "pid": project_id}, + { + "eid": entity_id, + "pid": project_id, + "embedding_model": search_repo._embedding_model_key(), + "vector_index": search_repo._semantic_vector_index_name, + }, ) chunk_id = chunk_result.scalar_one() + await session.commit() - await search_repo._write_embeddings( - session, - [(chunk_id, "vec content")], - [_unit_vector(dimensions)], + await search_repo._semantic_vector_index.upsert( + [ + VectorRecord( + key=VectorKey(entity_id=entity_id, chunk_key="chunk-1"), + values=tuple(_unit_vector(dimensions)), + ) + ] + ) + + async with db.scoped_session(session_maker) as session: + await session.execute( + text("UPDATE search_vector_chunks SET embedding_status = 'ready' WHERE id = :chunk_id"), + {"chunk_id": chunk_id}, ) await session.commit() @@ -171,7 +175,7 @@ def _config_manager_semantic_enabled() -> ConfigManager: status = await project_service.get_embedding_status(project_id) assert status.semantic_search_enabled is True - # The vec0 JOIN must succeed, so the table is reported as present and healthy. + # Status reads the ready manifest, so it never needs to reopen the vec0 table. assert status.vector_tables_exist is True assert status.reindex_recommended is False assert status.reindex_reason is None diff --git a/tests/repository/test_pgvector_index.py b/tests/repository/test_pgvector_index.py new file mode 100644 index 000000000..5e08e6247 --- /dev/null +++ b/tests/repository/test_pgvector_index.py @@ -0,0 +1,253 @@ +"""Unit tests for the built-in pgvector semantic index adapter.""" + +from __future__ import annotations + +from collections.abc import AsyncIterator +from contextlib import asynccontextmanager +from unittest.mock import MagicMock + +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.semantic_errors import SemanticDependenciesMissingError +from basic_memory.repository.semantic_vector_index import ( + VectorIndexScope, + VectorKey, + VectorRecord, +) + + +class FakeResult: + """Small SQLAlchemy result stand-in for adapter boundary tests.""" + + def __init__( + self, + *, + rows: list[dict[str, object]] | None = None, + scalar: object | None = None, + fetchone: object | None = None, + ) -> None: + self._rows = rows or [] + self._scalar = scalar + self._fetchone = fetchone + + def fetchone(self) -> object | None: + return self._fetchone + + def scalar_one_or_none(self) -> object | None: + return self._scalar + + def mappings(self) -> FakeResult: + return self + + def all(self) -> list[dict[str, object]]: + return self._rows + + +class FakeSession: + """Query-aware session that records adapter SQL without PostgreSQL.""" + + def __init__( + self, + *, + table_exists: bool = False, + dimensions: int | None = None, + chunk_rows: list[dict[str, object]] | None = None, + search_rows: list[dict[str, object]] | None = None, + fail_extension: bool = False, + ) -> None: + self.table_exists = table_exists + self.dimensions = dimensions + self.chunk_rows = chunk_rows or [] + self.search_rows = search_rows or [] + self.fail_extension = fail_extension + self.calls: list[tuple[str, dict[str, object] | None]] = [] + self.commit_count = 0 + + async def execute( + self, + statement: object, + params: dict[str, object] | None = None, + ) -> FakeResult: + sql = str(statement) + self.calls.append((sql, params)) + if "CREATE EXTENSION" in sql and self.fail_extension: + raise RuntimeError("extension unavailable") + if "information_schema.tables" in sql: + return FakeResult(fetchone=(1,) if self.table_exists else None) + if "SELECT atttypmod" in sql: + return FakeResult(scalar=self.dimensions) + if "SELECT id, entity_id, chunk_key" in sql: + return FakeResult(rows=self.chunk_rows) + if "AS similarity" in sql: + return FakeResult(rows=self.search_rows) + return FakeResult() + + async def commit(self) -> None: + self.commit_count += 1 + + +def _scope(dimensions: int = 4) -> VectorIndexScope: + return VectorIndexScope( + namespace="basic-memory-test", + project_id=7, + embedding_identity="stub:4", + dimensions=dimensions, + ) + + +def _install_session(monkeypatch: pytest.MonkeyPatch, session: FakeSession) -> None: + @asynccontextmanager + async def fake_scoped_session(_session_maker: object) -> AsyncIterator[FakeSession]: + yield session + + monkeypatch.setattr(pgvector_index_module.db, "scoped_session", fake_scoped_session) + + +def _sql_calls(session: FakeSession) -> list[str]: + return [sql for sql, _params in session.calls] + + +@pytest.mark.asyncio +async def test_initialize_creates_storage_once(monkeypatch) -> None: + session = FakeSession() + _install_session(monkeypatch, session) + index = PgVectorIndex(MagicMock(), _scope()) + + await index.initialize() + await index.initialize() + + sql_calls = _sql_calls(session) + assert sum("CREATE EXTENSION" in sql for sql in sql_calls) == 1 + assert any("embedding vector(4)" in sql for sql in sql_calls) + assert any("USING hnsw" in sql for sql in sql_calls) + assert not any("embedding_status = 'pending'" in sql for sql in sql_calls) + assert session.commit_count == 1 + + +@pytest.mark.asyncio +async def test_initialize_rebuilds_dimension_mismatch_and_invalidates_manifest(monkeypatch) -> None: + session = FakeSession(table_exists=True, dimensions=8) + _install_session(monkeypatch, session) + index = PgVectorIndex(MagicMock(), _scope()) + + await index.initialize() + + sql_calls = _sql_calls(session) + assert any("DROP TABLE IF EXISTS search_vector_embeddings" in sql for sql in sql_calls) + assert any("embedding_status = 'pending'" in sql for sql in sql_calls) + + +@pytest.mark.asyncio +async def test_initialize_reports_missing_pgvector_extension(monkeypatch) -> None: + session = FakeSession(fail_extension=True) + _install_session(monkeypatch, session) + index = PgVectorIndex(MagicMock(), _scope()) + + with pytest.raises(SemanticDependenciesMissingError, match="pgvector extension"): + await index.initialize() + + +@pytest.mark.asyncio +async def test_upsert_resolves_stable_keys_and_writes_one_batch(monkeypatch) -> None: + key_a = VectorKey(entity_id=11, chunk_key="entity:11:0") + key_b = VectorKey(entity_id=11, chunk_key="entity:11:1") + session = FakeSession( + chunk_rows=[ + {"id": 101, "entity_id": 11, "chunk_key": key_a.chunk_key}, + {"id": 102, "entity_id": 11, "chunk_key": key_b.chunk_key}, + ] + ) + _install_session(monkeypatch, session) + index = PgVectorIndex(MagicMock(), _scope()) + index._initialized = True + + await index.upsert( + [ + VectorRecord(key=key_a, values=(1.0, 0.0, 0.0, 0.0)), + VectorRecord(key=key_b, values=(0.0, 1.0, 0.0, 0.0)), + ] + ) + + insert_call = next(call for call in session.calls if "INSERT INTO" in call[0]) + assert insert_call[1] == { + "project_id": 7, + "chunk_id_0": 101, + "embedding_0": "[1,0,0,0]", + "dimensions_0": 4, + "chunk_id_1": 102, + "embedding_1": "[0,1,0,0]", + "dimensions_1": 4, + } + assert session.commit_count == 1 + + +@pytest.mark.asyncio +async def test_upsert_rejects_missing_manifest_key(monkeypatch) -> None: + key = VectorKey(entity_id=12, chunk_key="entity:12:0") + session = FakeSession() + _install_session(monkeypatch, session) + index = PgVectorIndex(MagicMock(), _scope()) + index._initialized = True + + with pytest.raises(RuntimeError, match="manifest rows are missing"): + await index.upsert([VectorRecord(key=key, values=(1.0, 0.0, 0.0, 0.0))]) + + +@pytest.mark.asyncio +async def test_delete_stable_keys_and_entity(monkeypatch) -> None: + key = VectorKey(entity_id=13, chunk_key="entity:13:0") + session = FakeSession(chunk_rows=[{"id": 103, "entity_id": 13, "chunk_key": key.chunk_key}]) + _install_session(monkeypatch, session) + index = PgVectorIndex(MagicMock(), _scope()) + index._initialized = True + + await index.delete([]) + await index.delete([key]) + await index.delete_entity(13) + await index.delete_orphans([key]) + + delete_calls = [call for call in session.calls if "DELETE FROM" in call[0]] + assert len(delete_calls) == 3 + assert delete_calls[0][1] == {"chunk_id_0": 103} + assert delete_calls[1][1] == {"project_id": 7, "entity_id": 13} + assert delete_calls[2][1] == { + "project_id": 7, + "embedding_identity": "stub:4", + } + assert session.commit_count == 3 + + +@pytest.mark.asyncio +async def test_search_returns_normalized_stable_matches(monkeypatch) -> None: + session = FakeSession( + search_rows=[ + {"entity_id": 14, "chunk_key": "entity:14:0", "similarity": 1.4}, + {"entity_id": 15, "chunk_key": "entity:15:0", "similarity": -0.2}, + ] + ) + _install_session(monkeypatch, session) + index = PgVectorIndex(MagicMock(), _scope()) + index._initialized = True + + assert await index.search([], limit=5) == [] + assert await index.search([1.0, 0.0, 0.0, 0.0], limit=0) == [] + with pytest.raises(ValueError, match="expected 4, got 2"): + await index.search([1.0, 0.0], limit=5) + + matches = await index.search([1.0, 0.0, 0.0, 0.0], limit=5) + + assert [(match.key.entity_id, match.similarity) for match in matches] == [ + (14, 1.0), + (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 search_call[1] == { + "query": "[1,0,0,0]", + "project_id": 7, + "dimensions": 4, + "embedding_identity": "stub:4", + "limit": 5, + } diff --git a/tests/repository/test_postgres_search_repository_unit.py b/tests/repository/test_postgres_search_repository_unit.py index a119129e3..f6d4bdc90 100644 --- a/tests/repository/test_postgres_search_repository_unit.py +++ b/tests/repository/test_postgres_search_repository_unit.py @@ -12,6 +12,7 @@ import basic_memory.repository.search_repository_base as search_repository_base_module from basic_memory.config import BasicMemoryConfig, DatabaseBackend +from basic_memory.repository.pgvector_index import PgVectorIndex from basic_memory.repository.postgres_search_repository import PostgresSearchRepository from basic_memory.repository.search_repository_base import _PreparedEntityVectorSync from basic_memory.repository.semantic_errors import ( @@ -60,21 +61,21 @@ def _make_repo( ) -# --- _format_pgvector_literal tests (lines 248-252) ----------------------- +# --- PgVectorIndex vector literal formatting ------------------------------ class TestFormatPgvectorLiteral: - """Cover PostgresSearchRepository._format_pgvector_literal.""" + """Cover the pgvector adapter's wire-format helper.""" def test_empty_vector(self): - assert PostgresSearchRepository._format_pgvector_literal([]) == "[]" + assert PgVectorIndex._format_vector([]) == "[]" def test_single_value(self): - result = PostgresSearchRepository._format_pgvector_literal([1.0]) + result = PgVectorIndex._format_vector([1.0]) assert result == "[1]" def test_multiple_values(self): - result = PostgresSearchRepository._format_pgvector_literal([0.1, 0.2, 0.3]) + result = PgVectorIndex._format_vector([0.1, 0.2, 0.3]) assert result.startswith("[") assert result.endswith("]") parts = result.strip("[]").split(",") @@ -82,15 +83,15 @@ def test_multiple_values(self): def test_high_precision(self): """Verify that 12-significant-digit formatting is used.""" - result = PostgresSearchRepository._format_pgvector_literal([1.23456789012345]) + result = PgVectorIndex._format_vector([1.23456789012345]) assert "1.23456789012" in result def test_integers_formatted_without_trailing_zeros(self): - result = PostgresSearchRepository._format_pgvector_literal([1.0, 2.0, 3.0]) + result = PgVectorIndex._format_vector([1.0, 2.0, 3.0]) assert result == "[1,2,3]" def test_negative_values(self): - result = PostgresSearchRepository._format_pgvector_literal([-0.5, 0.5]) + result = PgVectorIndex._format_vector([-0.5, 0.5]) assert "-0.5" in result assert "0.5" in result @@ -188,7 +189,17 @@ async def fake_scoped_session(session_maker): "basic_memory.repository.postgres_search_repository.db.scoped_session", fake_scoped_session, ) - monkeypatch.setattr(repo, "_get_existing_embedding_dims", AsyncMock(return_value=None)) + missing_table = MagicMock() + missing_table.fetchone.return_value = None + session.execute.side_effect = [ + MagicMock(), + MagicMock(), + MagicMock(), + missing_table, + MagicMock(), + MagicMock(), + MagicMock(), + ] await repo._ensure_vector_tables() @@ -199,7 +210,7 @@ async def fake_scoped_session(session_maker): "CREATE TABLE IF NOT EXISTS search_vector_embeddings" in sql for sql in executed_sql ) assert not any("ALTER TABLE search_vector_chunks" in sql for sql in executed_sql) - session.commit.assert_awaited_once() + assert session.commit.await_count == 2 assert repo._vector_tables_initialized is True @@ -261,28 +272,6 @@ async def test_delete_entity_chunks_executes_sql(self): assert params["entity_id"] == 42 -# --- _write_embeddings (lines 437-439) ------------------------------------- - - -class TestWriteEmbeddings: - """Cover _write_embeddings upsert logic.""" - - @pytest.mark.asyncio - async def test_write_embeddings_executes_single_bulk_upsert(self): - repo = _make_repo() - session = AsyncMock() - jobs = [(100, "chunk text A"), (200, "chunk text B")] - embeddings = [[0.1, 0.2, 0.3, 0.4], [0.5, 0.6, 0.7, 0.8]] - await repo._write_embeddings(session, jobs, embeddings) - assert session.execute.call_count == 1 - params = session.execute.call_args[0][1] - assert params["chunk_id_0"] == 100 - assert params["chunk_id_1"] == 200 - assert params["project_id"] == repo.project_id - assert params["embedding_dims_0"] == 4 - assert params["embedding_dims_1"] == 4 - - class TestBatchPrepareWindow: """Cover the shared batched prepare window used by Postgres.""" diff --git a/tests/repository/test_semantic_vector_index.py b/tests/repository/test_semantic_vector_index.py new file mode 100644 index 000000000..1c30bd125 --- /dev/null +++ b/tests/repository/test_semantic_vector_index.py @@ -0,0 +1,296 @@ +"""Contract and discovery tests for pluggable semantic vector indexes.""" + +from __future__ import annotations + +from collections.abc import Sequence +from dataclasses import dataclass +from typing import Any +from unittest.mock import MagicMock + +import pytest + +from basic_memory.config import BasicMemoryConfig, DatabaseBackend +from basic_memory.repository.embedding_provider import EmbeddingProvider +from basic_memory.repository.postgres_search_repository import PostgresSearchRepository +from basic_memory.repository.search_repository import create_search_repository +from basic_memory.repository.semantic_errors import SemanticVectorIndexExtensionError +from basic_memory.repository.semantic_vector_index import ( + SEMANTIC_VECTOR_INDEX_ENTRY_POINT_GROUP, + SemanticVectorIndex, + VectorIndexScope, + VectorKey, + VectorMatch, + VectorRecord, + validate_query_dimensions, + validate_vector_dimensions, +) +from basic_memory.repository.semantic_vector_index_factory import ( + build_vector_index_scope, + create_semantic_vector_index, + resolve_semantic_vector_index_name, +) + + +class StubEmbeddingProvider: + """Small embedding provider used only to build deterministic scopes.""" + + model_name = "stub-model" + dimensions = 3 + + async def embed_query(self, text: str) -> list[float]: + return [1.0, 0.0, 0.0] + + async def embed_documents(self, texts: list[str]) -> list[list[float]]: + return [[1.0, 0.0, 0.0] for _ in texts] + + def runtime_log_attrs(self) -> dict[str, Any]: + return {} + + +class StubVectorIndex: + """Structurally complete adapter for runtime protocol checks.""" + + def __init__(self, scope: VectorIndexScope): + self.scope = scope + + async def initialize(self) -> None: + return None + + async def upsert(self, records: Sequence[VectorRecord]) -> None: + return None + + async def delete(self, keys: Sequence[VectorKey]) -> None: + return None + + async def delete_entity(self, entity_id: int) -> None: + return None + + async def search( + self, + query: Sequence[float], + *, + limit: int, + ) -> list[VectorMatch]: + return [] + + +@dataclass(frozen=True) +class StubEntryPoint: + value: str + loaded: object + + def load(self) -> object: + return self.loaded + + +def _postgres_config(**overrides: object) -> BasicMemoryConfig: + values: dict[str, object] = { + "env": "test", + "database_backend": DatabaseBackend.POSTGRES, + "database_url": "postgresql+asyncpg://user:secret@db.example.test:5432/memory", + "semantic_search_enabled": True, + "semantic_vector_index": "milvus", + } + values.update(overrides) + return BasicMemoryConfig(**values) + + +def test_vector_contract_values_and_dimension_validation() -> None: + scope = VectorIndexScope( + namespace="basic-memory-test", + project_id=7, + embedding_identity="stub:3", + dimensions=3, + ) + key = VectorKey(entity_id=11, chunk_key="entity:11:0") + record = VectorRecord(key=key, values=(1.0, 0.0, 0.0)) + + assert isinstance(StubVectorIndex(scope), SemanticVectorIndex) + validate_vector_dimensions(scope, [record]) + validate_query_dimensions(scope, [1.0, 0.0, 0.0]) + + with pytest.raises(ValueError, match="expected 3, got 2"): + validate_vector_dimensions(scope, [VectorRecord(key=key, values=(1.0, 0.0))]) + with pytest.raises(ValueError, match="expected 3, got 1"): + validate_query_dimensions(scope, [1.0]) + + +def test_selector_defaults_to_pgvector_and_sqlite_remains_automatic() -> None: + default_config = BasicMemoryConfig(env="test") + milvus_config = _postgres_config() + + assert default_config.semantic_vector_index == "pgvector" + assert ( + resolve_semantic_vector_index_name(default_config, DatabaseBackend.POSTGRES) == "pgvector" + ) + assert resolve_semantic_vector_index_name(milvus_config, DatabaseBackend.SQLITE) == "sqlite-vec" + + +def test_scope_is_stable_credential_free_and_project_isolated() -> None: + provider: EmbeddingProvider = StubEmbeddingProvider() + first = build_vector_index_scope(_postgres_config(), provider, project_id=7) + rotated_password = build_vector_index_scope( + _postgres_config( + database_url=( + "postgresql+asyncpg://rotated-user:new-secret@db.example.test:5432/memory" + "?sslmode=require" + ) + ), + provider, + project_id=7, + ) + other_project = build_vector_index_scope(_postgres_config(), provider, project_id=8) + + assert first.namespace == rotated_password.namespace + assert "secret" not in first.namespace + assert first.project_id != other_project.project_id + assert first.embedding_identity == "StubEmbeddingProvider:stub-model:3" + assert first.dimensions == 3 + + +def test_missing_configured_extension_fails_without_fallback(monkeypatch) -> None: + monkeypatch.setattr( + "basic_memory.repository.semantic_vector_index_factory.entry_points", + lambda **_kwargs: (), + ) + + with pytest.raises( + SemanticVectorIndexExtensionError, + match="configured but no extension is installed", + ): + create_semantic_vector_index( + session_maker=MagicMock(), + project_id=7, + app_config=_postgres_config(), + database_backend=DatabaseBackend.POSTGRES, + embedding_provider=StubEmbeddingProvider(), + ) + + +def test_extension_factory_receives_explicit_scope_and_config(monkeypatch) -> None: + captured: dict[str, object] = {} + + def factory(*, scope: VectorIndexScope, app_config: BasicMemoryConfig) -> StubVectorIndex: + captured.update(scope=scope, app_config=app_config) + return StubVectorIndex(scope) + + monkeypatch.setattr( + "basic_memory.repository.semantic_vector_index_factory.entry_points", + lambda **kwargs: ( + ( + StubEntryPoint( + value="basic_memory_milvus:create_index", + loaded=factory, + ), + ) + if kwargs + == { + "group": SEMANTIC_VECTOR_INDEX_ENTRY_POINT_GROUP, + "name": "milvus", + } + else () + ), + ) + config = _postgres_config() + + name, index = create_semantic_vector_index( + session_maker=MagicMock(), + project_id=7, + app_config=config, + database_backend=DatabaseBackend.POSTGRES, + embedding_provider=StubEmbeddingProvider(), + ) + + assert name == "milvus" + assert isinstance(index, StubVectorIndex) + assert captured["app_config"] is config + assert captured["scope"] == index.scope + + +@pytest.mark.parametrize( + ("entry_points", "message"), + [ + ( + ( + StubEntryPoint("first:create", lambda **_kwargs: None), + StubEntryPoint("second:create", lambda **_kwargs: None), + ), + "Multiple semantic vector index extensions", + ), + ((StubEntryPoint("invalid:value", object()),), "must load a callable factory"), + ( + (StubEntryPoint("incompatible:create", lambda **_kwargs: object()),), + "returned an incompatible adapter", + ), + ], +) +def test_invalid_extension_registration_fails_explicitly( + monkeypatch, + entry_points: tuple[StubEntryPoint, ...], + message: str, +) -> None: + monkeypatch.setattr( + "basic_memory.repository.semantic_vector_index_factory.entry_points", + lambda **_kwargs: entry_points, + ) + + with pytest.raises(SemanticVectorIndexExtensionError, match=message): + create_semantic_vector_index( + session_maker=MagicMock(), + project_id=7, + app_config=_postgres_config(), + database_backend=DatabaseBackend.POSTGRES, + embedding_provider=StubEmbeddingProvider(), + ) + + +def test_extension_cannot_replace_the_required_scope(monkeypatch) -> None: + wrong_scope = VectorIndexScope( + namespace="other-installation", + project_id=999, + embedding_identity="other-model", + dimensions=3, + ) + monkeypatch.setattr( + "basic_memory.repository.semantic_vector_index_factory.entry_points", + lambda **_kwargs: ( + StubEntryPoint( + "wrong-scope:create", + lambda **_factory_kwargs: StubVectorIndex(wrong_scope), + ), + ), + ) + + with pytest.raises(SemanticVectorIndexExtensionError, match="wrong scope"): + create_semantic_vector_index( + session_maker=MagicMock(), + project_id=7, + app_config=_postgres_config(), + database_backend=DatabaseBackend.POSTGRES, + embedding_provider=StubEmbeddingProvider(), + ) + + +def test_search_repository_composition_root_injects_selected_adapter(monkeypatch) -> None: + provider = StubEmbeddingProvider() + scope = build_vector_index_scope(_postgres_config(), provider, project_id=7) + index = StubVectorIndex(scope) + monkeypatch.setattr( + "basic_memory.repository.search_repository.create_embedding_provider", + lambda _config: provider, + ) + monkeypatch.setattr( + "basic_memory.repository.search_repository.create_semantic_vector_index", + lambda **_kwargs: ("milvus", index), + ) + + repository = create_search_repository( + MagicMock(), + project_id=7, + app_config=_postgres_config(), + database_backend=DatabaseBackend.POSTGRES, + ) + + assert isinstance(repository, PostgresSearchRepository) + assert repository._semantic_vector_index_name == "milvus" + assert repository._semantic_vector_index is index diff --git a/tests/repository/test_semantic_vector_sync.py b/tests/repository/test_semantic_vector_sync.py index 8be7da334..986361aaa 100644 --- a/tests/repository/test_semantic_vector_sync.py +++ b/tests/repository/test_semantic_vector_sync.py @@ -304,9 +304,9 @@ async def test_prepare_window_read_helpers_handle_empty_inputs() -> None: assert source_rows == {} assert existing_rows == {} - assert "LEFT JOIN search_vector_embeddings" in ( - semantic_vector_sync.prepare_window_existing_rows_sql(":entity_id_0") - ) + manifest_sql = semantic_vector_sync.prepare_window_existing_rows_sql(":entity_id_0") + assert "embedding_status" in manifest_sql + assert "search_vector_embeddings" not in manifest_sql session.execute.assert_not_awaited() diff --git a/tests/repository/test_sqlite_vector_search_repository.py b/tests/repository/test_sqlite_vector_search_repository.py index f6efa5376..c4ba40934 100644 --- a/tests/repository/test_sqlite_vector_search_repository.py +++ b/tests/repository/test_sqlite_vector_search_repository.py @@ -2,9 +2,10 @@ import asyncio import hashlib +from collections.abc import Sequence from contextlib import asynccontextmanager from datetime import datetime, timezone -from typing import Any, cast +from typing import cast from unittest.mock import AsyncMock, MagicMock import pytest @@ -15,8 +16,17 @@ from basic_memory.repository.embedding_provider import EmbeddingProvider from basic_memory.repository.litellm_provider import LiteLLMEmbeddingProvider from basic_memory.repository.prefixing_provider import PrefixingEmbeddingProvider +from basic_memory.repository import search_repository_base as search_repository_base_module from basic_memory.repository.search_index_row import SearchIndexRow +from basic_memory.repository.semantic_vector_index import ( + VectorIndexScope, + VectorKey, + VectorMatch, + VectorRecord, +) from basic_memory.repository.sqlite_search_repository import SQLiteSearchRepository +from basic_memory.repository import sqlite_vec_index as sqlite_vec_index_module +from basic_memory.repository.sqlite_vec_index import SQLITE_VEC_MAX_K, SQLiteVecIndex from basic_memory.schemas.search import SearchItemType, SearchRetrievalMode @@ -53,6 +63,61 @@ class StubEmbeddingProviderV2(StubEmbeddingProvider): model_name = "stub-v2" +class RecordingVectorIndex: + """In-memory adapter with injectable write/delete failures.""" + + def __init__(self) -> None: + self.scope = VectorIndexScope( + namespace="test", + project_id=1, + embedding_identity="test", + dimensions=4, + ) + self.records: dict[VectorKey, tuple[float, ...]] = {} + self.upsert_calls: list[list[VectorRecord]] = [] + self.deleted_entities: list[int] = [] + self.reconcile_calls: list[list[VectorKey]] = [] + self.fail_upsert = False + self.fail_delete_entity = False + self.fail_search = False + + async def initialize(self) -> None: + return None + + async def upsert(self, records: Sequence[VectorRecord]) -> None: + self.upsert_calls.append(list(records)) + if self.fail_upsert: + raise RuntimeError("adapter write failed") + self.records.update({record.key: record.values for record in records}) + + async def delete(self, keys: Sequence[VectorKey]) -> None: + for key in keys: + self.records.pop(key, None) + + async def delete_entity(self, entity_id: int) -> None: + self.deleted_entities.append(entity_id) + if self.fail_delete_entity: + raise RuntimeError("adapter delete failed") + self.records = { + key: values for key, values in self.records.items() if key.entity_id != entity_id + } + + async def delete_orphans(self, live_keys: Sequence[VectorKey]) -> None: + self.reconcile_calls.append(list(live_keys)) + live_key_set = set(live_keys) + self.records = {key: values for key, values in self.records.items() if key in live_key_set} + + async def search( + self, + query: Sequence[float], + *, + limit: int, + ) -> list[VectorMatch]: + if self.fail_search: + raise RuntimeError("adapter query failed") + return [VectorMatch(key=key, similarity=1.0) for key in list(self.records)[:limit]] + + def _entity_row( *, project_id: int, @@ -172,6 +237,8 @@ async def test_sqlite_vec_tables_are_created_and_rebuilt(search_repository): "source_hash", "entity_fingerprint", "embedding_model", + "vector_index", + "embedding_status", "updated_at", } @@ -184,6 +251,76 @@ async def test_sqlite_vec_tables_are_created_and_rebuilt(search_repository): assert table_result.scalar_one() == "search_vector_embeddings" +@pytest.mark.asyncio +async def test_sqlite_vec_reconciliation_is_project_scoped(search_repository): + """Reconciliation removes non-ready local rows without touching another project.""" + if not isinstance(search_repository, SQLiteSearchRepository): + pytest.skip("sqlite-vec reconciliation behavior is local SQLite-only.") + + _enable_semantic(search_repository) + await search_repository.init_search_index() + index = cast(SQLiteVecIndex, search_repository._semantic_vector_index) + 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, :entity_id, :project_id, :chunk_key, 'text', 'hash', " + "'fingerprint', :embedding_model, 'sqlite-vec', :embedding_status)" + ), + [ + { + "id": 901, + "entity_id": 901, + "project_id": search_repository.project_id, + "chunk_key": "entity:901:0", + "embedding_model": embedding_identity, + "embedding_status": "pending", + }, + { + "id": 902, + "entity_id": 902, + "project_id": search_repository.project_id, + "chunk_key": "entity:902:0", + "embedding_model": embedding_identity, + "embedding_status": "ready", + }, + { + "id": 903, + "entity_id": 903, + "project_id": search_repository.project_id + 1, + "chunk_key": "entity:903:0", + "embedding_model": embedding_identity, + "embedding_status": "pending", + }, + ], + ) + await session.execute( + text( + "INSERT INTO search_vector_embeddings (rowid, embedding) " + "VALUES (:rowid, :embedding)" + ), + [{"rowid": rowid, "embedding": "[1,0,0,0]"} for rowid in (901, 902, 903)], + ) + await session.commit() + + await index.delete_orphans([]) + + async with db.scoped_session(search_repository.session_maker) as session: + remaining = await session.execute( + text( + "SELECT rowid FROM search_vector_embeddings " + "WHERE rowid IN (901, 902, 903) ORDER BY rowid" + ) + ) + assert remaining.scalars().all() == [902, 903] + + @pytest.mark.asyncio async def test_sqlite_chunk_upsert_and_delete_lifecycle(search_repository): """sync_entity_vectors updates changed chunks and clears vectors when source rows disappear.""" @@ -252,6 +389,312 @@ async def test_sqlite_chunk_upsert_and_delete_lifecycle(search_repository): assert int(embedding_count.scalar_one()) == 0 +@pytest.mark.asyncio +async def test_adapter_write_failure_stays_pending_and_retries_idempotently(search_repository): + """A partial external write must never make an uncommitted vector searchable.""" + if not isinstance(search_repository, SQLiteSearchRepository): + pytest.skip("Semantic manifest behavior is exercised through local SQLite.") + + _enable_semantic(search_repository) + adapter = RecordingVectorIndex() + adapter.fail_upsert = True + search_repository._semantic_vector_index = adapter + search_repository._semantic_vector_index_name = "recording" + await search_repository.init_search_index() + await search_repository.index_item( + _entity_row( + project_id=search_repository.project_id, + row_id=111, + entity_id=111, + title="Retryable Adapter Write", + permalink="specs/retryable-adapter-write", + content_stems="auth token retry", + ) + ) + + with pytest.raises(RuntimeError, match="adapter write failed"): + await search_repository.sync_entity_vectors(111) + + async with db.scoped_session(search_repository.session_maker) as session: + failed_state = await session.execute( + text( + "SELECT DISTINCT embedding_status FROM search_vector_chunks " + "WHERE project_id = :project_id AND entity_id = :entity_id" + ), + {"project_id": search_repository.project_id, "entity_id": 111}, + ) + assert failed_state.scalars().all() == ["pending"] + + adapter.fail_upsert = False + await search_repository.sync_entity_vectors(111) + + async with db.scoped_session(search_repository.session_maker) as session: + recovered_state = await session.execute( + text( + "SELECT DISTINCT vector_index, embedding_status FROM search_vector_chunks " + "WHERE project_id = :project_id AND entity_id = :entity_id" + ), + {"project_id": search_repository.project_id, "entity_id": 111}, + ) + assert recovered_state.all() == [("recording", "ready")] + + assert len(adapter.upsert_calls) == 2 + assert [record.key for record in adapter.upsert_calls[0]] == [ + record.key for record in adapter.upsert_calls[1] + ] + + +@pytest.mark.asyncio +async def test_ready_commit_failure_retries_same_stable_adapter_key( + search_repository, + monkeypatch, +): + """An adapter success followed by SQL failure must remain a safe idempotent retry.""" + if not isinstance(search_repository, SQLiteSearchRepository): + pytest.skip("Semantic manifest behavior is exercised through local SQLite.") + + _enable_semantic(search_repository) + adapter = RecordingVectorIndex() + search_repository._semantic_vector_index = adapter + search_repository._semantic_vector_index_name = "recording" + await search_repository.init_search_index() + + async with db.scoped_session(search_repository.session_maker) as session: + inserted = await session.execute( + text( + "INSERT INTO search_vector_chunks (" + "entity_id, project_id, chunk_key, chunk_text, source_hash, " + "entity_fingerprint, embedding_model, vector_index, embedding_status" + ") VALUES (" + ":entity_id, :project_id, :chunk_key, :chunk_text, :source_hash, " + ":entity_fingerprint, :embedding_model, :vector_index, 'pending'" + ") RETURNING id" + ), + { + "entity_id": 115, + "project_id": search_repository.project_id, + "chunk_key": "entity:115:0", + "chunk_text": "ready commit retry", + "source_hash": "hash", + "entity_fingerprint": "fingerprint", + "embedding_model": search_repository._embedding_model_key(), + "vector_index": "recording", + }, + ) + row_id = int(inserted.scalar_one()) + await session.commit() + + original_scoped_session = search_repository_base_module.db.scoped_session + context_count = 0 + + @asynccontextmanager + async def fail_ready_commit(session_maker): + nonlocal context_count + context_count += 1 + async with original_scoped_session(session_maker) as session: + if context_count == 2: + + async def fail_commit() -> None: + raise RuntimeError("ready commit failed") + + monkeypatch.setattr(session, "commit", fail_commit) + yield session + + monkeypatch.setattr( + search_repository_base_module.db, + "scoped_session", + fail_ready_commit, + ) + with pytest.raises(RuntimeError, match="ready commit failed"): + await search_repository._persist_embeddings( + [(row_id, "ready commit retry")], + [[1.0, 0.0, 0.0, 0.0]], + ) + monkeypatch.setattr( + search_repository_base_module.db, + "scoped_session", + original_scoped_session, + ) + + async with db.scoped_session(search_repository.session_maker) as session: + failed_status = await session.execute( + text("SELECT embedding_status FROM search_vector_chunks WHERE id = :row_id"), + {"row_id": row_id}, + ) + assert failed_status.scalar_one() == "pending" + + await search_repository._persist_embeddings( + [(row_id, "ready commit retry")], + [[1.0, 0.0, 0.0, 0.0]], + ) + + async with db.scoped_session(search_repository.session_maker) as session: + recovered_status = await session.execute( + text("SELECT embedding_status FROM search_vector_chunks WHERE id = :row_id"), + {"row_id": row_id}, + ) + assert recovered_status.scalar_one() == "ready" + + assert len(adapter.upsert_calls) == 2 + assert adapter.upsert_calls[0][0].key == adapter.upsert_calls[1][0].key + + +@pytest.mark.asyncio +async def test_adapter_delete_failure_stays_pending_until_retry(search_repository): + """External delete failure preserves non-searchable manifest intent for retry.""" + if not isinstance(search_repository, SQLiteSearchRepository): + pytest.skip("Semantic manifest behavior is exercised through local SQLite.") + + _enable_semantic(search_repository) + adapter = RecordingVectorIndex() + search_repository._semantic_vector_index = adapter + search_repository._semantic_vector_index_name = "recording" + await search_repository.init_search_index() + await search_repository.index_item( + _entity_row( + project_id=search_repository.project_id, + row_id=112, + entity_id=112, + title="Retryable Adapter Delete", + permalink="specs/retryable-adapter-delete", + content_stems="schema migration retry", + ) + ) + await search_repository.sync_entity_vectors(112) + + adapter.fail_delete_entity = True + with pytest.raises(RuntimeError, match="adapter delete failed"): + await search_repository.delete_entity_vector_rows(112) + + async with db.scoped_session(search_repository.session_maker) as session: + failed_state = await session.execute( + text( + "SELECT DISTINCT embedding_status FROM search_vector_chunks " + "WHERE project_id = :project_id AND entity_id = :entity_id" + ), + {"project_id": search_repository.project_id, "entity_id": 112}, + ) + assert failed_state.scalars().all() == ["pending"] + + adapter.fail_delete_entity = False + await search_repository.delete_entity_vector_rows(112) + + async with db.scoped_session(search_repository.session_maker) as session: + row_count = await session.execute( + text( + "SELECT COUNT(*) FROM search_vector_chunks " + "WHERE project_id = :project_id AND entity_id = :entity_id" + ), + {"project_id": search_repository.project_id, "entity_id": 112}, + ) + assert row_count.scalar_one() == 0 + + assert adapter.deleted_entities == [112, 112] + + +@pytest.mark.asyncio +async def test_adapter_matches_hydrate_only_current_ready_manifest_rows(search_repository): + """Stale external matches fail closed unless SQL says the current row is ready.""" + if not isinstance(search_repository, SQLiteSearchRepository): + pytest.skip("Semantic manifest behavior is exercised through local SQLite.") + + _enable_semantic(search_repository) + adapter = RecordingVectorIndex() + search_repository._semantic_vector_index = adapter + search_repository._semantic_vector_index_name = "recording" + await search_repository.init_search_index() + await search_repository.index_item( + _entity_row( + project_id=search_repository.project_id, + row_id=113, + entity_id=113, + title="Manifest Authority", + permalink="specs/manifest-authority", + content_stems="queue worker task", + ) + ) + await search_repository.sync_entity_vectors(113) + adapter.records[VectorKey(entity_id=999, chunk_key="foreign:999:0")] = ( + 1.0, + 0.0, + 0.0, + 0.0, + ) + + ready_results = await search_repository.search( + search_text="queue worker", + retrieval_mode=SearchRetrievalMode.VECTOR, + ) + assert {result.entity_id for result in ready_results} == {113} + + await search_repository.reconcile_vector_index() + assert set(adapter.records) == set(adapter.reconcile_calls[0]) + assert {key.entity_id for key in adapter.records} == {113} + + adapter.fail_search = True + with pytest.raises(RuntimeError, match="adapter query failed"): + await search_repository.search( + search_text="queue worker", + retrieval_mode=SearchRetrievalMode.VECTOR, + ) + adapter.fail_search = False + + async with db.scoped_session(search_repository.session_maker) as session: + await session.execute( + text( + "UPDATE search_vector_chunks SET embedding_status = 'pending' " + "WHERE project_id = :project_id AND entity_id = :entity_id" + ), + {"project_id": search_repository.project_id, "entity_id": 113}, + ) + await session.commit() + + pending_results = await search_repository.search( + search_text="queue worker", + retrieval_mode=SearchRetrievalMode.VECTOR, + ) + assert pending_results == [] + + +@pytest.mark.asyncio +async def test_vector_index_switch_reembeds_manifest_rows(search_repository): + """Changing the selected adapter identity invalidates otherwise unchanged chunks.""" + if not isinstance(search_repository, SQLiteSearchRepository): + pytest.skip("Semantic manifest behavior is exercised through local SQLite.") + + _enable_semantic(search_repository) + adapter = RecordingVectorIndex() + search_repository._semantic_vector_index = adapter + search_repository._semantic_vector_index_name = "recording-a" + await search_repository.init_search_index() + await search_repository.index_item( + _entity_row( + project_id=search_repository.project_id, + row_id=114, + entity_id=114, + title="Index Switch", + permalink="specs/index-switch", + content_stems="database semantic index switch", + ) + ) + await search_repository.sync_entity_vectors(114) + + search_repository._semantic_vector_index_name = "recording-b" + await search_repository.sync_entity_vectors(114) + + async with db.scoped_session(search_repository.session_maker) as session: + state = await session.execute( + text( + "SELECT DISTINCT vector_index, embedding_status FROM search_vector_chunks " + "WHERE project_id = :project_id AND entity_id = :entity_id" + ), + {"project_id": search_repository.project_id, "entity_id": 114}, + ) + assert state.all() == [("recording-b", "ready")] + + assert len(adapter.upsert_calls) == 2 + + @pytest.mark.asyncio async def test_sqlite_vector_sync_skips_unchanged_and_reembeds_changed_content(search_repository): """SQLite vector sync tracks new, changed, unchanged, and model-changed entities.""" @@ -693,8 +1136,8 @@ async def test_sqlite_hybrid_search_combines_fts_and_vector(search_repository): @pytest.mark.asyncio -async def test_run_vector_query_caps_k_at_sqlite_vec_limit(search_repository): - """_run_vector_query must cap the knn k param at SQLITE_VEC_MAX_K (4096). +async def test_run_vector_query_caps_k_at_sqlite_vec_limit(search_repository, monkeypatch): + """The sqlite-vec adapter caps k while preserving the requested outer limit. sqlite-vec raises OperationalError when k > 4096. The candidate_limit passed from the base class can exceed this for large projects, so @@ -706,34 +1149,39 @@ async def test_run_vector_query_caps_k_at_sqlite_vec_limit(search_repository): _enable_semantic(search_repository) await search_repository.init_search_index() - # Track the parameters passed to session.execute + index = cast(SQLiteVecIndex, search_repository._semantic_vector_index) captured_params: list[dict] = [] + session = AsyncMock() async def capturing_execute(stmt, params=None): if params and "vector_k" in params: captured_params.append(dict(params)) - # Return empty result set mock_result = MagicMock() mock_result.mappings.return_value.all.return_value = [] return mock_result - async with db.scoped_session(search_repository.session_maker) as session: - await search_repository._prepare_vector_session(session) - cast(Any, session).execute = capturing_execute - - query_embedding = [0.1] * search_repository._vector_dimensions + @asynccontextmanager + async def fake_scoped_session(_session_maker): + yield session - # candidate_limit exceeds sqlite-vec limit - await search_repository._run_vector_query(session, query_embedding, 10000) + session.execute = capturing_execute + monkeypatch.setattr(sqlite_vec_index_module.db, "scoped_session", fake_scoped_session) + monkeypatch.setattr(index, "_ensure_loaded", AsyncMock()) + query_embedding = [0.1] * search_repository._vector_dimensions - assert len(captured_params) == 1 - assert captured_params[0]["vector_k"] == SQLiteSearchRepository.SQLITE_VEC_MAX_K - assert captured_params[0]["candidate_limit"] == 10000 + await index.search(query_embedding, limit=10000) - # candidate_limit within limit should pass through unchanged - captured_params.clear() - await search_repository._run_vector_query(session, query_embedding, 500) + assert captured_params == [ + { + "query": "[0.1, 0.1, 0.1, 0.1]", + "vector_k": SQLITE_VEC_MAX_K, + "project_id": search_repository.project_id, + "embedding_identity": search_repository._embedding_model_key(), + "limit": 10000, + } + ] - assert len(captured_params) == 1 - assert captured_params[0]["vector_k"] == 500 - assert captured_params[0]["candidate_limit"] == 500 + captured_params.clear() + await index.search(query_embedding, limit=500) + assert captured_params[0]["vector_k"] == 500 + assert captured_params[0]["limit"] == 500 diff --git a/tests/services/test_project_removal_bug.py b/tests/services/test_project_removal_bug.py index 06cd08019..46f82a5b9 100644 --- a/tests/services/test_project_removal_bug.py +++ b/tests/services/test_project_removal_bug.py @@ -230,9 +230,10 @@ async def test_remove_project_purges_search_rows(project_service: ProjectService text( "INSERT INTO search_vector_chunks " "(entity_id, project_id, chunk_key, chunk_text, source_hash, " - " entity_fingerprint, embedding_model) " + " entity_fingerprint, embedding_model, vector_index, embedding_status) " "VALUES (:entity_id, :project_id, :chunk_key, :chunk_text, " - " :source_hash, :entity_fingerprint, :embedding_model)" + " :source_hash, :entity_fingerprint, :embedding_model, " + " 'sqlite-vec', 'pending')" ), { "entity_id": 999_001, @@ -331,9 +332,10 @@ async def test_remove_project_purges_vector_embeddings( text( "INSERT INTO search_vector_chunks " "(id, entity_id, project_id, chunk_key, chunk_text, source_hash, " - " entity_fingerprint, embedding_model) " + " entity_fingerprint, embedding_model, vector_index, embedding_status) " "VALUES (:id, :entity_id, :project_id, :chunk_key, :chunk_text, " - " :source_hash, :entity_fingerprint, :embedding_model)" + " :source_hash, :entity_fingerprint, :embedding_model, " + " 'sqlite-vec', 'ready')" ), { "id": 999_201, diff --git a/tests/services/test_project_service_embedding_status.py b/tests/services/test_project_service_embedding_status.py index 2b12aa3c7..4bb643c36 100644 --- a/tests/services/test_project_service_embedding_status.py +++ b/tests/services/test_project_service_embedding_status.py @@ -7,6 +7,11 @@ from sqlalchemy import text from basic_memory import db +from basic_memory.repository.embedding_provider_factory import create_embedding_provider +from basic_memory.repository.semantic_vector_index_factory import ( + resolve_semantic_vector_index_name, + semantic_embedding_identity, +) from basic_memory.schemas.project_info import EmbeddingStatus from basic_memory.services.project_service import ProjectService @@ -20,27 +25,6 @@ async def _execute(project_service: ProjectService, query, params=None): return await project_service.repository.execute_query(session, query, params or {}) -async def _create_embeddings_stub(project_service: ProjectService) -> None: - """Create a minimal search_vector_embeddings stub so vector_tables_exist is True. - - Test fixtures run with semantic search disabled, so the real vec0/pgvector - embeddings table is never created. get_embedding_status only probes table - existence and joins on chunk_id (rowid on SQLite), so a plain table suffices. - """ - await _execute( - project_service, - text( - "CREATE TABLE IF NOT EXISTS search_vector_embeddings ( chunk_id INTEGER PRIMARY KEY)" - ), - {}, - ) - - -async def _drop_embeddings_stub(project_service: ProjectService) -> None: - """Drop the stub table to avoid polluting subsequent tests.""" - await _execute(project_service, text("DROP TABLE IF EXISTS search_vector_embeddings"), {}) - - @pytest.mark.asyncio async def test_embedding_status_semantic_disabled(project_service: ProjectService, test_project): """When semantic search is disabled, return minimal status with zero counts.""" @@ -88,7 +72,7 @@ async def test_embedding_status_vector_tables_missing( assert status.embedding_model == "bge-small-en-v1.5" assert status.vector_tables_exist is False assert status.reindex_recommended is True - assert "Vector tables not initialized" in (status.reindex_reason or "") + assert "Vector manifest not initialized" in (status.reindex_reason or "") @pytest.mark.asyncio @@ -96,9 +80,6 @@ async def test_embedding_status_entities_without_chunks( project_service: ProjectService, test_graph, test_project ): """When entities have search_index rows but no chunks, recommend reindex.""" - # search_vector_chunks comes from Base.metadata; the embeddings table needs a stub - # because fixtures run with semantic search disabled. - await _create_embeddings_stub(project_service) with patch.object( type(project_service), "config_manager", @@ -108,8 +89,6 @@ async def test_embedding_status_entities_without_chunks( ): status = await project_service.get_embedding_status(test_project.id) - await _drop_embeddings_stub(project_service) - assert status.semantic_search_enabled is True assert status.vector_tables_exist is True # test_graph creates entities indexed in search_index, but no vector chunks @@ -123,8 +102,7 @@ async def test_embedding_status_entities_without_chunks( async def test_embedding_status_orphaned_chunks( project_service: ProjectService, test_graph, test_project ): - """When chunks exist without matching embeddings, recommend reindex.""" - # Insert a chunk row (no matching embedding = orphan) + """When a manifest row is pending, recommend reindex.""" # Get a real entity_id from the test graph entity_result = await _execute( project_service, @@ -133,32 +111,12 @@ async def test_embedding_status_orphaned_chunks( ) entity_id = entity_result.scalar() - await _execute( + await _insert_manifest_chunk( project_service, - text( - "INSERT INTO search_vector_chunks " - "(" - "entity_id, project_id, chunk_key, chunk_text, source_hash, " - "entity_fingerprint, embedding_model" - ") " - "VALUES (" - ":entity_id, :project_id, 'chunk-1', 'test text', 'abc123', " - "'fp-abc123', 'bge-small-en-v1.5'" - ")" - ), - {"entity_id": entity_id, "project_id": test_project.id}, - ) - - # Create a minimal search_vector_embeddings stub (not a real vector table) - # so the LEFT JOIN works and finds the orphan. - # Uses chunk_id as PK — Postgres queries join on chunk_id, - # SQLite queries join on rowid which aliases INTEGER PRIMARY KEY. - await _execute( - project_service, - text( - "CREATE TABLE IF NOT EXISTS search_vector_embeddings ( chunk_id INTEGER PRIMARY KEY)" - ), - {}, + entity_id=entity_id, + project_id=test_project.id, + chunk_key="chunk-1", + embedding_status="pending", ) with patch.object( @@ -170,35 +128,52 @@ async def test_embedding_status_orphaned_chunks( ): status = await project_service.get_embedding_status(test_project.id) - # Clean up stub table to avoid polluting subsequent tests - await _execute(project_service, text("DROP TABLE IF EXISTS search_vector_embeddings"), {}) - assert status.vector_tables_exist is True assert status.total_chunks == 1 assert status.orphaned_chunks == 1 assert status.reindex_recommended is True - assert "orphaned chunks" in (status.reindex_reason or "") + assert "need vector indexing" in (status.reindex_reason or "") @pytest.mark.asyncio -async def test_embedding_status_handles_sqlite_vec_unavailable( +async def test_embedding_status_counts_only_current_ready_manifest_rows( project_service: ProjectService, test_graph, test_project ): - """When sqlite-vec can't load at all, degrade to unavailable status instead of crashing.""" - # Trigger: Postgres test matrix executes the same unit suite. - # Why: sqlite-vec loading failures are specific to SQLite virtual tables, not Postgres joins. - # Outcome: keep the regression focused on the backend that can actually hit this path. - if _is_postgres(): - pytest.skip("sqlite-vec unavailable handling is SQLite-specific.") - - # Both vector tables must exist so the status check reaches the vec query; - # fixtures run with semantic search disabled, so stub the embeddings table. - await _create_embeddings_stub(project_service) - - # scalar_vec_query returns None when the extension can't be loaded on this - # Python build (e.g. the python.org macOS interpreter). Simulate that here. - async def _vec_query_unavailable(_session, query, params=None): - return None + """Status is backend-neutral and treats pending or mismatched rows as stale.""" + entity_result = await _execute( + project_service, + text("SELECT id FROM entity WHERE project_id = :project_id LIMIT 1"), + {"project_id": test_project.id}, + ) + entity_id = entity_result.scalar() + + await _insert_manifest_chunk( + project_service, + entity_id=entity_id, + project_id=test_project.id, + chunk_key="ready", + ) + await _insert_manifest_chunk( + project_service, + entity_id=entity_id, + project_id=test_project.id, + chunk_key="pending", + embedding_status="pending", + ) + await _insert_manifest_chunk( + project_service, + entity_id=entity_id, + project_id=test_project.id, + chunk_key="wrong-index", + vector_index="milvus", + ) + await _insert_manifest_chunk( + project_service, + entity_id=entity_id, + project_id=test_project.id, + chunk_key="wrong-model", + embedding_identity="OtherProvider:other:384", + ) with patch.object( type(project_service), @@ -210,17 +185,17 @@ async def _vec_query_unavailable(_session, query, params=None): with patch.object( project_service.repository, "scalar_vec_query", - side_effect=_vec_query_unavailable, + side_effect=AssertionError("status must not query vector storage"), ): status = await project_service.get_embedding_status(test_project.id) - await _drop_embeddings_stub(project_service) - assert status.semantic_search_enabled is True - assert status.total_indexed_entities > 0 - assert status.vector_tables_exist is False + assert status.vector_tables_exist is True + assert status.total_chunks == 4 + assert status.total_embeddings == 1 + assert status.orphaned_chunks == 3 assert status.reindex_recommended is True - assert "sqlite-vec is unavailable" in (status.reindex_reason or "") + assert "pending or stale" in (status.reindex_reason or "") @pytest.mark.asyncio @@ -229,18 +204,7 @@ async def test_embedding_status_healthy(project_service: ProjectService, test_gr # Clear any leftover data from prior tests await _execute(project_service, text("DELETE FROM search_vector_chunks"), {}) - # Drop any existing virtual table (may have been created by search_service init) - # and recreate as a simple regular table for testing the join logic. - # Uses chunk_id as PK — Postgres queries join on chunk_id, - # SQLite queries join on rowid which aliases INTEGER PRIMARY KEY. - await _execute(project_service, text("DROP TABLE IF EXISTS search_vector_embeddings"), {}) - await _execute( - project_service, - text("CREATE TABLE search_vector_embeddings ( chunk_id INTEGER PRIMARY KEY)"), - {}, - ) - - # Insert a chunk + matching embedding for every search_index entity + # Insert a current, ready manifest row for every search_index entity. entity_result = await _execute( project_service, text("SELECT DISTINCT entity_id FROM search_index WHERE project_id = :project_id"), @@ -250,30 +214,11 @@ async def test_embedding_status_healthy(project_service: ProjectService, test_gr chunk_id = 1 for eid in entity_ids: - await _execute( + await _insert_manifest_chunk( project_service, - text( - "INSERT INTO search_vector_chunks " - "(" - "id, entity_id, project_id, chunk_key, chunk_text, source_hash, " - "entity_fingerprint, embedding_model" - ") " - "VALUES (" - ":id, :entity_id, :project_id, :key, 'text', 'hash', " - "'fp-hash', 'bge-small-en-v1.5'" - ")" - ), - { - "id": chunk_id, - "entity_id": eid, - "project_id": test_project.id, - "key": f"chunk-{chunk_id}", - }, - ) - await _execute( - project_service, - text("INSERT INTO search_vector_embeddings (chunk_id) VALUES (:chunk_id)"), - {"chunk_id": chunk_id}, + entity_id=eid, + project_id=test_project.id, + chunk_key=f"chunk-{chunk_id}", ) chunk_id += 1 @@ -286,9 +231,6 @@ async def test_embedding_status_healthy(project_service: ProjectService, test_gr ): status = await project_service.get_embedding_status(test_project.id) - # Clean up stub table to avoid polluting subsequent tests - await _execute(project_service, text("DROP TABLE IF EXISTS search_vector_embeddings"), {}) - assert status.vector_tables_exist is True assert status.total_chunks > 0 assert status.total_embeddings == status.total_chunks @@ -310,9 +252,6 @@ async def test_embedding_status_excludes_stale_entity_ids( # Include 'id' column — required NOT NULL on Postgres (regular table), # ignored on SQLite (FTS5 virtual table where id is UNINDEXED). stale_entity_id = 999999 - # Both vector tables must exist to reach the stale-filtered count queries; - # fixtures run with semantic search disabled, so stub the embeddings table. - await _create_embeddings_stub(project_service) await _execute( project_service, text( @@ -372,3 +311,37 @@ def _config_manager_with(semantic_search_enabled: bool): # Patch the config object in-place cm.config.semantic_search_enabled = semantic_search_enabled return cm + + +async def _insert_manifest_chunk( + project_service: ProjectService, + *, + entity_id: int, + project_id: int, + chunk_key: str, + vector_index: str | None = None, + embedding_identity: str | None = None, + embedding_status: str = "ready", +) -> None: + """Insert one manifest row with explicit backend and readiness identity.""" + config = _config_manager_with(semantic_search_enabled=True).config + active_vector_index = resolve_semantic_vector_index_name(config, config.database_backend) + active_embedding_identity = semantic_embedding_identity(create_embedding_provider(config)) + await _execute( + project_service, + text( + "INSERT INTO search_vector_chunks " + "(entity_id, project_id, chunk_key, chunk_text, source_hash, " + "entity_fingerprint, embedding_model, vector_index, embedding_status) " + "VALUES (:entity_id, :project_id, :chunk_key, 'test text', 'hash', " + "'fingerprint', :embedding_identity, :vector_index, :embedding_status)" + ), + { + "entity_id": entity_id, + "project_id": project_id, + "chunk_key": chunk_key, + "embedding_identity": embedding_identity or active_embedding_identity, + "vector_index": vector_index or active_vector_index, + "embedding_status": embedding_status, + }, + ) diff --git a/tests/services/test_semantic_search.py b/tests/services/test_semantic_search.py index a4ff6e6e0..cd978e3aa 100644 --- a/tests/services/test_semantic_search.py +++ b/tests/services/test_semantic_search.py @@ -247,11 +247,14 @@ async def test_reindex_vectors_respects_embed_opt_out(search_service, monkeypatc ) monkeypatch.setattr(search_service, "_purge_stale_search_rows", purge_stale_rows) monkeypatch.setattr(search_service, "sync_entity_vectors_batch", sync_batch) + reconcile = AsyncMock() + monkeypatch.setattr(search_service.repository, "reconcile_vector_index", reconcile) stats = await search_service.reindex_vectors() purge_stale_rows.assert_awaited_once() sync_batch.assert_awaited_once_with([41, 42], progress_callback=None) + reconcile.assert_awaited_once() assert stats == { "total_entities": 2, "embedded": 1, @@ -284,10 +287,13 @@ async def sync_entity_vectors_batch(entity_ids, progress_callback=None): monkeypatch.setattr(repository, "delete_stale_vector_rows", delete_stale_vector_rows) monkeypatch.setattr(search_service, "sync_entity_vectors_batch", sync_entity_vectors_batch) + reconcile = AsyncMock() + monkeypatch.setattr(repository, "reconcile_vector_index", reconcile) stats = await search_service.reindex_vectors() assert calls == ["purge", "sync"] + reconcile.assert_awaited_once() assert stats == { "total_entities": 1, "embedded": 1, @@ -297,8 +303,8 @@ async def sync_entity_vectors_batch(entity_ids, progress_callback=None): @pytest.mark.asyncio -async def test_reindex_all_uses_sqlite_vec_aware_drop(search_service, monkeypatch): - """Full service reindex should not drop vec0 tables through a raw connection.""" +async def test_reindex_all_uses_vector_adapter_cleanup(search_service, monkeypatch): + """Full service reindex should clean vectors through the repository boundary.""" repository = _sqlite_repo(search_service) executed_sql: list[str] = [] calls: list[str] = [] @@ -306,17 +312,17 @@ async def test_reindex_all_uses_sqlite_vec_aware_drop(search_service, monkeypatc async def execute_query(query, params=None): executed_sql.append(str(query)) - async def drop_vector_tables(): - calls.append("drop_vector_tables") + async def delete_project_vector_rows(): + calls.append("delete_project_vector_rows") monkeypatch.setattr(repository, "execute_query", execute_query) - monkeypatch.setattr(repository, "drop_vector_tables", drop_vector_tables) + monkeypatch.setattr(repository, "delete_project_vector_rows", delete_project_vector_rows) monkeypatch.setattr(search_service, "init_search_index", AsyncMock()) monkeypatch.setattr(search_service.entity_repository, "find_all", AsyncMock(return_value=[])) await search_service.reindex_all() - assert calls == ["drop_vector_tables"] + assert calls == ["delete_project_vector_rows"] assert all("search_vector_embeddings" not in sql for sql in executed_sql) @@ -381,6 +387,8 @@ async def test_reindex_vectors_force_full_clears_project_vectors_before_resync( ) monkeypatch.setattr(search_service, "_purge_stale_search_rows", purge_stale_rows) monkeypatch.setattr(repository, "delete_project_vector_rows", delete_project_vectors) + reconcile = AsyncMock() + monkeypatch.setattr(repository, "reconcile_vector_index", reconcile) monkeypatch.setattr(search_service, "sync_entity_vectors_batch", sync_batch) stats = await search_service.reindex_vectors(force_full=True) @@ -388,6 +396,7 @@ async def test_reindex_vectors_force_full_clears_project_vectors_before_resync( purge_stale_rows.assert_awaited_once() delete_project_vectors.assert_awaited_once() sync_batch.assert_awaited_once_with([41, 42], progress_callback=None) + reconcile.assert_awaited_once() assert stats == { "total_entities": 2, "embedded": 2, diff --git a/tests/test_semantic_vector_index_migration.py b/tests/test_semantic_vector_index_migration.py new file mode 100644 index 000000000..8fcc97c9a --- /dev/null +++ b/tests/test_semantic_vector_index_migration.py @@ -0,0 +1,100 @@ +"""Tests for semantic vector manifest state migration behavior.""" + +from __future__ import annotations + +from types import SimpleNamespace +from unittest.mock import MagicMock + +from basic_memory.alembic.versions import ( + o8j9k0l1m2n3_add_vector_index_manifest_state as migration, +) + + +def _connection(dialect: str) -> SimpleNamespace: + return SimpleNamespace(dialect=SimpleNamespace(name=dialect)) + + +def test_upgrade_backfills_ready_state_from_existing_pgvector_rows(monkeypatch) -> None: + connection = _connection("postgresql") + inspector = MagicMock() + inspector.get_table_names.return_value = [ + "search_vector_chunks", + "search_vector_embeddings", + ] + execute = MagicMock() + create_check_constraint = MagicMock() + monkeypatch.setattr(migration.op, "get_bind", lambda: connection) + monkeypatch.setattr(migration, "inspect", lambda _connection: inspector) + monkeypatch.setattr(migration.op, "execute", execute) + monkeypatch.setattr(migration.op, "create_check_constraint", create_check_constraint) + + migration.upgrade() + + statements = [call.args[0] for call in execute.call_args_list] + assert any("ADD COLUMN IF NOT EXISTS vector_index" in sql for sql in statements) + assert any("ADD COLUMN IF NOT EXISTS embedding_status" in sql for sql in statements) + assert any("SET vector_index = 'pgvector'" in sql for sql in statements) + assert any("WHEN EXISTS" in sql and "THEN 'ready'" in sql for sql in statements) + assert any("ALTER COLUMN vector_index SET NOT NULL" in sql for sql in statements) + assert any("ALTER COLUMN embedding_status SET NOT NULL" in sql for sql in statements) + create_check_constraint.assert_called_once_with( + "ck_search_vector_chunks_embedding_status", + "search_vector_chunks", + "embedding_status IN ('pending', 'ready')", + ) + + +def test_upgrade_marks_pending_without_embedding_table(monkeypatch) -> None: + connection = _connection("postgresql") + inspector = MagicMock() + inspector.get_table_names.return_value = ["search_vector_chunks"] + execute = MagicMock() + monkeypatch.setattr(migration.op, "get_bind", lambda: connection) + monkeypatch.setattr(migration, "inspect", lambda _connection: inspector) + monkeypatch.setattr(migration.op, "execute", execute) + monkeypatch.setattr(migration.op, "create_check_constraint", MagicMock()) + + migration.upgrade() + + statements = [call.args[0] for call in execute.call_args_list] + assert "UPDATE search_vector_chunks SET embedding_status = 'pending'" in statements + + +def test_upgrade_is_noop_for_sqlite_or_missing_manifest(monkeypatch) -> None: + execute = MagicMock() + monkeypatch.setattr(migration.op, "execute", execute) + monkeypatch.setattr(migration.op, "create_check_constraint", MagicMock()) + + monkeypatch.setattr(migration.op, "get_bind", lambda: _connection("sqlite")) + migration.upgrade() + + inspector = MagicMock() + inspector.get_table_names.return_value = [] + monkeypatch.setattr(migration.op, "get_bind", lambda: _connection("postgresql")) + monkeypatch.setattr(migration, "inspect", lambda _connection: inspector) + migration.upgrade() + + execute.assert_not_called() + + +def test_downgrade_removes_manifest_state(monkeypatch) -> None: + connection = _connection("postgresql") + inspector = MagicMock() + inspector.get_table_names.return_value = ["search_vector_chunks"] + execute = MagicMock() + drop_constraint = MagicMock() + monkeypatch.setattr(migration.op, "get_bind", lambda: connection) + monkeypatch.setattr(migration, "inspect", lambda _connection: inspector) + monkeypatch.setattr(migration.op, "execute", execute) + monkeypatch.setattr(migration.op, "drop_constraint", drop_constraint) + + migration.downgrade() + + drop_constraint.assert_called_once_with( + "ck_search_vector_chunks_embedding_status", + "search_vector_chunks", + type_="check", + ) + statements = [call.args[0] for call in execute.call_args_list] + assert any("DROP COLUMN IF EXISTS embedding_status" in sql for sql in statements) + assert any("DROP COLUMN IF EXISTS vector_index" in sql for sql in statements) From 367a2f782aaaf53f47fd954077371e9c652ef65c Mon Sep 17 00:00:00 2001 From: phernandez Date: Wed, 22 Jul 2026 08:20:57 -0500 Subject: [PATCH 02/24] test(core): align semantic diagnostics with vector adapters Signed-off-by: phernandez --- test-int/semantic/test_search_diagnostics.py | 29 ++++++++++---------- 1 file changed, 15 insertions(+), 14 deletions(-) diff --git a/test-int/semantic/test_search_diagnostics.py b/test-int/semantic/test_search_diagnostics.py index b924b0b5c..12d544ac4 100644 --- a/test-int/semantic/test_search_diagnostics.py +++ b/test-int/semantic/test_search_diagnostics.py @@ -313,12 +313,10 @@ async def test_score_fusion_preserves_strong_vector_match(sqlite_engine_factory, @pytest.mark.semantic @pytest.mark.benchmark async def test_similarity_formula_analysis(sqlite_engine_factory, tmp_path): - """Analyze the raw distance-to-similarity mapping for real queries. + """Analyze normalized similarities returned by the configured vector index. - Production formulas are backend-specific: - - SQLite: similarity = max(0, 1 - L2²/2) for normalized embeddings - - Postgres: similarity = max(0, 1 - cosine_distance) - This test compares old and new mappings for diagnostics. + Distance metrics are backend-specific and remain inside each vector adapter. + The repository boundary exposes normalized cosine similarity in ``[0, 1]``. """ skip_if_needed(DIAG_COMBO) provider = _create_fastembed_provider() @@ -335,7 +333,8 @@ async def test_similarity_formula_analysis(sqlite_engine_factory, tmp_path): ] for query_text in queries: - # Get raw vector distances by querying at the repository level + # Query through the repository boundary so this diagnostic verifies the + # normalized contract rather than reaching into sqlite-vec internals. query_embedding = await provider.embed_query(query_text.strip()) from basic_memory import db as bm_db @@ -343,16 +342,18 @@ async def test_similarity_formula_analysis(sqlite_engine_factory, tmp_path): repo = cast(Any, service.repository) async with bm_db.scoped_session(repo.session_maker) as session: await repo._prepare_vector_session(session) - raw_rows = await repo._run_vector_query(session, query_embedding, candidate_limit=20) + vector_rows = await repo._run_vector_query( + session, + query_embedding, + candidate_limit=20, + ) print(f"\nQuery: '{query_text}'") - print(f" {'chunk_key':<40} {'distance':>10} {'sim_old':>12} {'sim_new':>12}") - for row in raw_rows[:10]: - dist = float(row["best_distance"]) - sim_old = 1.0 / (1.0 + max(dist, 0.0)) - # New formula: L2 distance → cosine similarity for normalized embeddings - sim_new = repo._distance_to_similarity(dist) - print(f" {row['chunk_key']:<40} {dist:>10.4f} {sim_old:>12.4f} {sim_new:>12.4f}") + print(f" {'chunk_key':<40} {'similarity':>12}") + for row in vector_rows[:10]: + similarity = float(row["best_similarity"]) + assert 0.0 <= similarity <= 1.0 + print(f" {row['chunk_key']:<40} {similarity:>12.4f}") # --- Test: min_similarity threshold effectiveness --- From d780c16f2894d6ad41fb09ba810e56605360e811 Mon Sep 17 00:00:00 2001 From: phernandez Date: Wed, 22 Jul 2026 08:29:59 -0500 Subject: [PATCH 03/24] fix(core): harden semantic reindex lifecycle Signed-off-by: phernandez --- .../repository/embedding_provider_factory.py | 64 ++++++++++ .../repository/litellm_provider.py | 39 ++++-- .../repository/prefixing_provider.py | 25 +++- .../repository/search_repository_base.py | 54 +++++++-- src/basic_memory/services/project_service.py | 7 +- tests/repository/test_semantic_search_base.py | 113 +++++++++++++++++- .../test_project_service_embedding_status.py | 59 ++++++++- 7 files changed, 330 insertions(+), 31 deletions(-) diff --git a/src/basic_memory/repository/embedding_provider_factory.py b/src/basic_memory/repository/embedding_provider_factory.py index dcee3616d..3a472a306 100644 --- a/src/basic_memory/repository/embedding_provider_factory.py +++ b/src/basic_memory/repository/embedding_provider_factory.py @@ -12,6 +12,7 @@ PrefixingEmbeddingProvider, embedding_prefix_digest, normalize_embedding_prefix, + prefixing_embedding_identity, ) # Cache key fields are limited to values that change the *identity* of the loaded @@ -143,6 +144,69 @@ def reset_embedding_provider_cache() -> None: _EMBEDDING_PROVIDER_CACHE.clear() +def configured_embedding_provider_identity(app_config: BasicMemoryConfig) -> str: + """Resolve the persisted embedding identity without constructing a provider.""" + provider_name = app_config.semantic_embedding_provider.strip().lower() + configured_dimensions = app_config.semantic_embedding_dimensions + + if provider_name == "fastembed": + provider_type_name = "FastEmbedEmbeddingProvider" + model_name = app_config.semantic_embedding_model + dimensions = configured_dimensions or 384 + provider_identity = f"{model_name}:{dimensions}" + elif provider_name == "openai": + provider_type_name = "OpenAIEmbeddingProvider" + model_name = app_config.semantic_embedding_model or "text-embedding-3-small" + if model_name == "bge-small-en-v1.5": + model_name = "text-embedding-3-small" + dimensions = configured_dimensions or 1536 + provider_identity = f"{model_name}:{dimensions}" + elif provider_name == "litellm": + from basic_memory.repository.litellm_provider import ( + _default_input_types, + litellm_embedding_identity, + ) + + provider_type_name = "LiteLLMEmbeddingProvider" + model_name = app_config.semantic_embedding_model or "openai/text-embedding-3-small" + if model_name == "bge-small-en-v1.5": + model_name = "openai/text-embedding-3-small" + if configured_dimensions is None and model_name != "openai/text-embedding-3-small": + raise ValueError( + "semantic_embedding_dimensions must be set when " + "semantic_embedding_provider='litellm' uses a non-default model. " + f"Configured model: {model_name!r}." + ) + dimensions = configured_dimensions or 1536 + default_document_input_type, default_query_input_type = _default_input_types(model_name) + provider_identity = litellm_embedding_identity( + model_name=model_name, + dimensions=dimensions, + document_input_type=( + app_config.semantic_embedding_document_input_type or default_document_input_type + ), + query_input_type=( + app_config.semantic_embedding_query_input_type or default_query_input_type + ), + forward_dimensions=app_config.semantic_embedding_forward_dimensions, + ) + else: + raise ValueError(f"Unsupported semantic embedding provider: {provider_name}") + + document_prefix = normalize_embedding_prefix(app_config.semantic_embedding_document_prefix) + query_prefix = normalize_embedding_prefix(app_config.semantic_embedding_query_prefix) + if document_prefix is None and query_prefix is None: + return f"{provider_type_name}:{provider_identity}" + + prefixed_identity = prefixing_embedding_identity( + provider_type_name=provider_type_name, + provider_identity=provider_identity, + document_prefix=document_prefix, + query_prefix=query_prefix, + ) + return f"PrefixingEmbeddingProvider:{prefixed_identity}" + + def create_embedding_provider(app_config: BasicMemoryConfig) -> EmbeddingProvider: """Create an embedding provider based on semantic config. diff --git a/src/basic_memory/repository/litellm_provider.py b/src/basic_memory/repository/litellm_provider.py index 2e8a7d384..a9055fce2 100644 --- a/src/basic_memory/repository/litellm_provider.py +++ b/src/basic_memory/repository/litellm_provider.py @@ -64,6 +64,28 @@ def _should_forward_dimensions(model_name: str, forward_dimensions: bool | None) return "text-embedding-3" in normalized +def litellm_embedding_identity( + *, + model_name: str, + dimensions: int, + document_input_type: str | None, + query_input_type: str | None, + forward_dimensions: bool | None, +) -> str: + """Return the LiteLLM request semantics that invalidate stored vectors.""" + resolved_document_input_type = document_input_type or "-" + resolved_query_input_type = query_input_type or "-" + resolved_forward_dimensions = str( + _should_forward_dimensions(model_name, forward_dimensions) + ).lower() + return ( + f"{model_name}:{dimensions}:" + f"document_input_type={resolved_document_input_type}:" + f"query_input_type={resolved_query_input_type}:" + f"forward_dimensions={resolved_forward_dimensions}" + ) + + def _import_litellm() -> Any: """Import LiteLLM without letting its import-time dotenv hook read cwd secrets.""" # Constraint: LiteLLM 1.85.0 loads .env files at import time when @@ -127,18 +149,13 @@ def runtime_log_attrs(self) -> dict[str, Any]: def identity_key(self) -> str: """Return the embedding semantics that should invalidate stored vectors.""" - document_input_type = self.document_input_type or "-" - query_input_type = self.query_input_type or "-" - forward_dimensions = str( - _should_forward_dimensions(self.model_name, self.forward_dimensions) - ).lower() - identity = ( - f"{self.model_name}:{self.dimensions}:" - f"document_input_type={document_input_type}:" - f"query_input_type={query_input_type}:" - f"forward_dimensions={forward_dimensions}" + return litellm_embedding_identity( + model_name=self.model_name, + dimensions=self.dimensions, + document_input_type=self.document_input_type, + query_input_type=self.query_input_type, + forward_dimensions=self.forward_dimensions, ) - return identity async def _embed(self, texts: list[str], *, input_type: str | None) -> list[list[float]]: if not texts: diff --git a/src/basic_memory/repository/prefixing_provider.py b/src/basic_memory/repository/prefixing_provider.py index 6234f9d62..89755a385 100644 --- a/src/basic_memory/repository/prefixing_provider.py +++ b/src/basic_memory/repository/prefixing_provider.py @@ -26,6 +26,21 @@ def embedding_prefix_digest(value: str | None) -> str: return hashlib.sha256(normalized.encode("utf-8")).hexdigest() +def prefixing_embedding_identity( + *, + provider_type_name: str, + provider_identity: str, + document_prefix: str | None, + query_prefix: str | None, +) -> str: + """Return prefix semantics without exposing literal prefix content.""" + return ( + f"{provider_type_name}:{provider_identity}:" + f"document_prefix_sha256={embedding_prefix_digest(document_prefix)}:" + f"query_prefix_sha256={embedding_prefix_digest(query_prefix)}" + ) + + class PrefixingEmbeddingProvider(EmbeddingProvider): """Apply document/query text prefixes before delegating to an embedding provider.""" @@ -74,9 +89,9 @@ def runtime_log_attrs(self) -> dict[str, Any]: def identity_key(self) -> str: """Return embedding semantics without exposing literal prefix content.""" - provider_identity = embedding_provider_identity(self.provider) - return ( - f"{type(self.provider).__name__}:{provider_identity}:" - f"document_prefix_sha256={embedding_prefix_digest(self.document_prefix)}:" - f"query_prefix_sha256={embedding_prefix_digest(self.query_prefix)}" + return prefixing_embedding_identity( + provider_type_name=type(self.provider).__name__, + provider_identity=embedding_provider_identity(self.provider), + document_prefix=self.document_prefix, + query_prefix=self.query_prefix, ) diff --git a/src/basic_memory/repository/search_repository_base.py b/src/basic_memory/repository/search_repository_base.py index f2c184496..a9bfaaae5 100644 --- a/src/basic_memory/repository/search_repository_base.py +++ b/src/basic_memory/repository/search_repository_base.py @@ -10,7 +10,7 @@ import logfire as logfire from loguru import logger -from sqlalchemy import Executable, Result, text +from sqlalchemy import Executable, Result, inspect, text from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker from basic_memory import db @@ -614,14 +614,19 @@ async def delete_project_vector_rows(self) -> None: Core enumerates ownership from the SQL manifest because the adapter contract intentionally has no project-wide listing or destructive reset. - Each entity deletion uses the same pending-first lifecycle as ordinary - sync cleanup, so an external failure remains retryable and fails closed. + A full reindex clears the manifest even when semantic search is disabled, + so stale ready rows cannot become current if the feature is re-enabled. + Adapter cleanup is best-effort because the manifest remains the search + authority and an unavailable extension must not block an FTS rebuild. """ - if not self._semantic_enabled: - return - - await self._ensure_vector_tables() async with db.scoped_session(self.session_maker) as session: + connection = await session.connection() + manifest_exists = await connection.run_sync( + lambda sync_connection: inspect(sync_connection).has_table("search_vector_chunks") + ) + if not manifest_exists: + return + result = await session.execute( text( "SELECT DISTINCT entity_id FROM search_vector_chunks " @@ -630,9 +635,40 @@ async def delete_project_vector_rows(self) -> None: {"project_id": self.project_id}, ) entity_ids = [int(entity_id) for entity_id in result.scalars().all()] + await session.execute( + text( + "UPDATE search_vector_chunks SET embedding_status = 'pending' " + "WHERE project_id = :project_id" + ), + {"project_id": self.project_id}, + ) + await session.commit() - for entity_id in entity_ids: - await self.delete_entity_vector_rows(entity_id) + if hasattr(self, "_semantic_vector_index"): + try: + await self._semantic_vector_index.initialize() + for entity_id in entity_ids: + await self._semantic_vector_index.delete_entity(entity_id) + except Exception as exc: + # Trigger: a configured external adapter cannot initialize or delete. + # Why: SQL is the authoritative hydration manifest, so clearing it + # prevents stale search results even when external cleanup is unavailable. + # Outcome: finish the FTS rebuild and leave a visible warning about + # storage that may need later adapter reconciliation. + logger.warning( + "Could not clean semantic vector adapter during full reindex: " + "project_id={project_id} vector_index={vector_index} error={error}", + project_id=self.project_id, + vector_index=self._semantic_vector_index_name, + error=exc, + ) + + async with db.scoped_session(self.session_maker) as session: + await session.execute( + text("DELETE FROM search_vector_chunks WHERE project_id = :project_id"), + {"project_id": self.project_id}, + ) + await session.commit() async def delete_stale_vector_rows(self) -> None: """Delete vectors whose source entity no longer exists. diff --git a/src/basic_memory/services/project_service.py b/src/basic_memory/services/project_service.py index 71589032e..64e047cde 100644 --- a/src/basic_memory/services/project_service.py +++ b/src/basic_memory/services/project_service.py @@ -16,10 +16,11 @@ from basic_memory import db from basic_memory.models import Project from basic_memory.repository.project_repository import ProjectRepository -from basic_memory.repository.embedding_provider_factory import create_embedding_provider +from basic_memory.repository.embedding_provider_factory import ( + configured_embedding_provider_identity, +) from basic_memory.repository.semantic_vector_index_factory import ( resolve_semantic_vector_index_name, - semantic_embedding_identity, ) from basic_memory.schemas import ( ActivityMetrics, @@ -1045,7 +1046,7 @@ async def get_embedding_status(self, project_id: int) -> EmbeddingStatus: is_postgres = config.database_backend == DatabaseBackend.POSTGRES vector_index = resolve_semantic_vector_index_name(config, config.database_backend) - embedding_identity = semantic_embedding_identity(create_embedding_provider(config)) + embedding_identity = configured_embedding_provider_identity(config) # --- Check vector manifest existence --- # The SQL manifest is authoritative even when vector values live outside diff --git a/tests/repository/test_semantic_search_base.py b/tests/repository/test_semantic_search_base.py index cdd79f9a2..69b966c1a 100644 --- a/tests/repository/test_semantic_search_base.py +++ b/tests/repository/test_semantic_search_base.py @@ -5,7 +5,7 @@ from datetime import datetime from types import SimpleNamespace from typing import Any -from unittest.mock import AsyncMock +from unittest.mock import AsyncMock, Mock import pytest @@ -140,6 +140,117 @@ async def fail_if_search_runs(**kwargs): assert search_calls == [] +@pytest.mark.asyncio +async def test_project_vector_cleanup_skips_missing_manifest(monkeypatch): + """A full reindex should remain safe before semantic tables ever existed.""" + repo = _ConcreteRepo() + session = AsyncMock() + connection = AsyncMock() + connection.run_sync.return_value = False + session.connection.return_value = connection + + @asynccontextmanager + async def fake_scoped_session(_session_maker): + yield session + + monkeypatch.setattr(search_repository_base_module.db, "scoped_session", fake_scoped_session) + + await repo.delete_project_vector_rows() + + session.execute.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_project_vector_cleanup_clears_disabled_manifest(monkeypatch): + """Disabled semantic search must not preserve stale ready manifest rows.""" + repo = _ConcreteRepo() + session = AsyncMock() + connection = AsyncMock() + connection.run_sync.return_value = True + session.connection.return_value = connection + entity_result = SimpleNamespace(scalars=lambda: SimpleNamespace(all=lambda: [41, 42])) + session.execute.return_value = entity_result + + @asynccontextmanager + async def fake_scoped_session(_session_maker): + yield session + + monkeypatch.setattr(search_repository_base_module.db, "scoped_session", fake_scoped_session) + + await repo.delete_project_vector_rows() + + statements = [str(call.args[0]) for call in session.execute.await_args_list] + assert any(statement.startswith("UPDATE search_vector_chunks") for statement in statements) + assert any(statement.startswith("DELETE FROM search_vector_chunks") for statement in statements) + + +@pytest.mark.asyncio +async def test_project_vector_cleanup_uses_available_adapter(monkeypatch): + """An available adapter should receive every manifest-owned entity deletion.""" + repo = _ConcreteRepo() + adapter: Any = SimpleNamespace( + initialize=AsyncMock(), + delete_entity=AsyncMock(), + ) + repo._semantic_vector_index = adapter + repo._semantic_vector_index_name = "milvus" + session = AsyncMock() + connection = AsyncMock() + connection.run_sync.return_value = True + session.connection.return_value = connection + session.execute.return_value = SimpleNamespace( + scalars=lambda: SimpleNamespace(all=lambda: [41, 42]) + ) + + @asynccontextmanager + async def fake_scoped_session(_session_maker): + yield session + + monkeypatch.setattr(search_repository_base_module.db, "scoped_session", fake_scoped_session) + + await repo.delete_project_vector_rows() + + adapter.initialize.assert_awaited_once() + assert adapter.delete_entity.await_args_list == [ + ((41,), {}), + ((42,), {}), + ] + + +@pytest.mark.asyncio +async def test_project_vector_cleanup_clears_manifest_after_adapter_failure(monkeypatch): + """An extension failure should fail closed by removing searchable manifests.""" + repo = _ConcreteRepo() + adapter: Any = SimpleNamespace( + initialize=AsyncMock(side_effect=RuntimeError("adapter unavailable")), + delete_entity=AsyncMock(), + ) + repo._semantic_vector_index = adapter + repo._semantic_vector_index_name = "milvus" + session = AsyncMock() + connection = AsyncMock() + connection.run_sync.return_value = True + session.connection.return_value = connection + session.execute.return_value = SimpleNamespace( + scalars=lambda: SimpleNamespace(all=lambda: [41]) + ) + + @asynccontextmanager + async def fake_scoped_session(_session_maker): + yield session + + warning = Mock() + monkeypatch.setattr(search_repository_base_module.db, "scoped_session", fake_scoped_session) + monkeypatch.setattr(search_repository_base_module.logger, "warning", warning) + + await repo.delete_project_vector_rows() + + adapter.delete_entity.assert_not_awaited() + warning.assert_called_once() + statements = [str(call.args[0]) for call in session.execute.await_args_list] + assert any(statement.startswith("DELETE FROM search_vector_chunks") for statement in statements) + + @pytest.mark.asyncio async def test_sync_entity_vectors_batch_flushes_at_configured_threshold(monkeypatch): """Batch sync should flush queued jobs at semantic_embedding_sync_batch_size boundaries.""" diff --git a/tests/services/test_project_service_embedding_status.py b/tests/services/test_project_service_embedding_status.py index 4bb643c36..3b197ae94 100644 --- a/tests/services/test_project_service_embedding_status.py +++ b/tests/services/test_project_service_embedding_status.py @@ -7,7 +7,11 @@ from sqlalchemy import text from basic_memory import db -from basic_memory.repository.embedding_provider_factory import create_embedding_provider +from basic_memory.config import BasicMemoryConfig +from basic_memory.repository.embedding_provider_factory import ( + configured_embedding_provider_identity, + create_embedding_provider, +) from basic_memory.repository.semantic_vector_index_factory import ( resolve_semantic_vector_index_name, semantic_embedding_identity, @@ -44,6 +48,57 @@ async def test_embedding_status_semantic_disabled(project_service: ProjectServic assert status.total_embeddings == 0 +@pytest.mark.parametrize( + "config", + [ + BasicMemoryConfig(), + BasicMemoryConfig( + semantic_embedding_provider="openai", + semantic_embedding_model="text-embedding-3-large", + semantic_embedding_dimensions=1024, + ), + BasicMemoryConfig( + semantic_embedding_provider="litellm", + semantic_embedding_model="cohere/embed-english-v3.0", + semantic_embedding_dimensions=1024, + semantic_embedding_document_prefix="passage: ", + semantic_embedding_query_prefix="query: ", + ), + ], +) +def test_configured_embedding_identity_matches_runtime_provider( + config: BasicMemoryConfig, +) -> None: + """Status identity must exactly match the provider used by vector sync.""" + assert configured_embedding_provider_identity(config) == semantic_embedding_identity( + create_embedding_provider(config) + ) + + +@pytest.mark.asyncio +async def test_embedding_status_does_not_construct_provider( + project_service: ProjectService, + test_project, +) -> None: + """Project status should remain a metadata query, not provider initialization.""" + with ( + patch.object( + type(project_service), + "config_manager", + new_callable=lambda: property( + lambda self: _config_manager_with(semantic_search_enabled=True) + ), + ), + patch( + "basic_memory.repository.embedding_provider_factory.create_embedding_provider", + side_effect=AssertionError("status must not construct an embedding provider"), + ), + ): + status = await project_service.get_embedding_status(test_project.id) + + assert status.semantic_search_enabled is True + + @pytest.mark.asyncio async def test_embedding_status_vector_tables_missing( project_service: ProjectService, test_graph, test_project @@ -326,7 +381,7 @@ async def _insert_manifest_chunk( """Insert one manifest row with explicit backend and readiness identity.""" config = _config_manager_with(semantic_search_enabled=True).config active_vector_index = resolve_semantic_vector_index_name(config, config.database_backend) - active_embedding_identity = semantic_embedding_identity(create_embedding_provider(config)) + active_embedding_identity = configured_embedding_provider_identity(config) await _execute( project_service, text( From 91d5411d8c8cde8f3d29f86be87a044afddf451f Mon Sep 17 00:00:00 2001 From: phernandez Date: Wed, 22 Jul 2026 08:42:01 -0500 Subject: [PATCH 04/24] fix(core): reconcile orphaned sqlite vectors Signed-off-by: phernandez --- .../repository/sqlite_vec_index.py | 21 +++++++++++++++++++ .../test_sqlite_vector_search_repository.py | 6 +++--- 2 files changed, 24 insertions(+), 3 deletions(-) diff --git a/src/basic_memory/repository/sqlite_vec_index.py b/src/basic_memory/repository/sqlite_vec_index.py index cc8cfeb5d..2779fc073 100644 --- a/src/basic_memory/repository/sqlite_vec_index.py +++ b/src/basic_memory/repository/sqlite_vec_index.py @@ -218,6 +218,27 @@ async def delete_orphans(self, _live_keys: Sequence[VectorKey]) -> None: await self.initialize() async with db.scoped_session(self._session_maker) as session: await self._ensure_loaded(session) + # A vec row without any manifest has no remaining project owner. + # Remove these globally before the project-scoped stale-state pass; + # otherwise they can occupy sqlite-vec's top-k window forever. + orphan_result = await session.execute( + text( + "SELECT rowid FROM search_vector_embeddings " + "EXCEPT SELECT id FROM search_vector_chunks" + ) + ) + orphan_rowids = [int(rowid) for rowid in orphan_result.scalars().all()] + if orphan_rowids: + params = { + f"orphan_rowid_{index}": rowid for index, rowid in enumerate(orphan_rowids) + } + placeholders = ", ".join( + f":orphan_rowid_{index}" for index in range(len(orphan_rowids)) + ) + await session.execute( + text(f"DELETE FROM search_vector_embeddings WHERE rowid IN ({placeholders})"), + params, + ) await session.execute( text( "DELETE FROM search_vector_embeddings WHERE rowid IN (" diff --git a/tests/repository/test_sqlite_vector_search_repository.py b/tests/repository/test_sqlite_vector_search_repository.py index c4ba40934..4345767cc 100644 --- a/tests/repository/test_sqlite_vector_search_repository.py +++ b/tests/repository/test_sqlite_vector_search_repository.py @@ -253,7 +253,7 @@ async def test_sqlite_vec_tables_are_created_and_rebuilt(search_repository): @pytest.mark.asyncio async def test_sqlite_vec_reconciliation_is_project_scoped(search_repository): - """Reconciliation removes non-ready local rows without touching another project.""" + """Reconciliation removes orphan/local stale rows without touching another project.""" if not isinstance(search_repository, SQLiteSearchRepository): pytest.skip("sqlite-vec reconciliation behavior is local SQLite-only.") @@ -305,7 +305,7 @@ async def test_sqlite_vec_reconciliation_is_project_scoped(search_repository): "INSERT INTO search_vector_embeddings (rowid, embedding) " "VALUES (:rowid, :embedding)" ), - [{"rowid": rowid, "embedding": "[1,0,0,0]"} for rowid in (901, 902, 903)], + [{"rowid": rowid, "embedding": "[1,0,0,0]"} for rowid in (901, 902, 903, 904)], ) await session.commit() @@ -315,7 +315,7 @@ async def test_sqlite_vec_reconciliation_is_project_scoped(search_repository): remaining = await session.execute( text( "SELECT rowid FROM search_vector_embeddings " - "WHERE rowid IN (901, 902, 903) ORDER BY rowid" + "WHERE rowid IN (901, 902, 903, 904) ORDER BY rowid" ) ) assert remaining.scalars().all() == [902, 903] From e398410dff44d570c8a01bcf8ae60067e533d3e1 Mon Sep 17 00:00:00 2001 From: phernandez Date: Wed, 22 Jul 2026 09:06:30 -0500 Subject: [PATCH 05/24] fix(core): invalidate manifests when vector storage resets Signed-off-by: phernandez --- src/basic_memory/repository/pgvector_index.py | 7 ++-- .../repository/sqlite_vec_index.py | 5 ++- tests/repository/test_pgvector_index.py | 17 ++++++++- .../test_sqlite_vector_search_repository.py | 38 +++++++++++++++++++ 4 files changed, 61 insertions(+), 6 deletions(-) diff --git a/src/basic_memory/repository/pgvector_index.py b/src/basic_memory/repository/pgvector_index.py index 7225d8087..9f0693bb6 100644 --- a/src/basic_memory/repository/pgvector_index.py +++ b/src/basic_memory/repository/pgvector_index.py @@ -57,6 +57,7 @@ async def initialize(self) -> None: ) from exc existing_dimensions = await self._existing_dimensions(session) + storage_missing = existing_dimensions is None dimensions_changed = ( existing_dimensions is not None and existing_dimensions != self.scope.dimensions ) @@ -97,11 +98,11 @@ async def initialize(self) -> None: ) ) - # Trigger: pgvector's fixed-width column was recreated for a new model. + # Trigger: pgvector storage was created or its fixed-width column changed. # Why: SQL manifest rows can otherwise remain `ready` after their vectors - # disappeared, causing the incremental sync to skip them forever. + # disappeared, causing incremental sync to skip them forever. # Outcome: the normal sync pipeline re-embeds every affected chunk. - if dimensions_changed: + if storage_missing or dimensions_changed: await session.execute( text( "UPDATE search_vector_chunks SET embedding_status = 'pending' " diff --git a/src/basic_memory/repository/sqlite_vec_index.py b/src/basic_memory/repository/sqlite_vec_index.py index 2779fc073..cacc2a195 100644 --- a/src/basic_memory/repository/sqlite_vec_index.py +++ b/src/basic_memory/repository/sqlite_vec_index.py @@ -99,6 +99,7 @@ async def initialize(self) -> None: ) ) vector_sql = result.scalar() + storage_missing = not vector_sql expected_dimensions = f"float[{self.scope.dimensions}]" dimensions_changed = bool(vector_sql and expected_dimensions not in vector_sql) if dimensions_changed: @@ -110,7 +111,9 @@ async def initialize(self) -> None: await session.execute(text("DROP TABLE IF EXISTS search_vector_embeddings")) await session.execute(create_sqlite_search_vector_embeddings(self.scope.dimensions)) - if dimensions_changed: + # Missing or dimension-rebuilt vec storage has no vectors, so ready + # manifests must become pending before incremental sync inspects them. + if storage_missing or dimensions_changed: await session.execute( text( "UPDATE search_vector_chunks SET embedding_status = 'pending' " diff --git a/tests/repository/test_pgvector_index.py b/tests/repository/test_pgvector_index.py index 5e08e6247..f52301733 100644 --- a/tests/repository/test_pgvector_index.py +++ b/tests/repository/test_pgvector_index.py @@ -110,7 +110,7 @@ def _sql_calls(session: FakeSession) -> list[str]: @pytest.mark.asyncio -async def test_initialize_creates_storage_once(monkeypatch) -> None: +async def test_initialize_creates_storage_once_and_invalidates_manifest(monkeypatch) -> None: session = FakeSession() _install_session(monkeypatch, session) index = PgVectorIndex(MagicMock(), _scope()) @@ -122,10 +122,23 @@ async def test_initialize_creates_storage_once(monkeypatch) -> None: assert sum("CREATE EXTENSION" in sql for sql in sql_calls) == 1 assert any("embedding vector(4)" in sql for sql in sql_calls) assert any("USING hnsw" in sql for sql in sql_calls) - assert not any("embedding_status = 'pending'" in sql for sql in sql_calls) + assert any("embedding_status = 'pending'" in sql for sql in sql_calls) assert session.commit_count == 1 +@pytest.mark.asyncio +async def test_initialize_preserves_manifest_when_storage_is_unchanged(monkeypatch) -> None: + session = FakeSession(table_exists=True, dimensions=4) + _install_session(monkeypatch, session) + index = PgVectorIndex(MagicMock(), _scope()) + + await index.initialize() + + sql_calls = _sql_calls(session) + assert not any("DROP TABLE IF EXISTS search_vector_embeddings" in sql for sql in sql_calls) + assert not any("embedding_status = 'pending'" in sql for sql in sql_calls) + + @pytest.mark.asyncio async def test_initialize_rebuilds_dimension_mismatch_and_invalidates_manifest(monkeypatch) -> None: session = FakeSession(table_exists=True, dimensions=8) diff --git a/tests/repository/test_sqlite_vector_search_repository.py b/tests/repository/test_sqlite_vector_search_repository.py index 4345767cc..e92605e3a 100644 --- a/tests/repository/test_sqlite_vector_search_repository.py +++ b/tests/repository/test_sqlite_vector_search_repository.py @@ -251,6 +251,44 @@ async def test_sqlite_vec_tables_are_created_and_rebuilt(search_repository): assert table_result.scalar_one() == "search_vector_embeddings" +@pytest.mark.asyncio +async def test_sqlite_vec_recreated_storage_invalidates_ready_manifest(search_repository): + """Recreated empty vec storage must force unchanged chunks back to pending.""" + if not isinstance(search_repository, SQLiteSearchRepository): + pytest.skip("sqlite-vec storage recovery is local SQLite-only.") + + _enable_semantic(search_repository) + await search_repository.init_search_index() + index = cast(SQLiteVecIndex, search_repository._semantic_vector_index) + + async with db.scoped_session(search_repository.session_maker) as 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 (" + "905, 905, :project_id, 'entity:905:0', 'text', 'hash', " + "'fingerprint', :embedding_model, 'sqlite-vec', 'ready')" + ), + { + "project_id": search_repository.project_id, + "embedding_model": search_repository._embedding_model_key(), + }, + ) + await session.execute(text("DROP TABLE search_vector_embeddings")) + await session.commit() + + index.invalidate_initialization() + await index.initialize() + + async with db.scoped_session(search_repository.session_maker) as session: + result = await session.execute( + text("SELECT embedding_status FROM search_vector_chunks WHERE id = 905") + ) + assert result.scalar_one() == "pending" + + @pytest.mark.asyncio async def test_sqlite_vec_reconciliation_is_project_scoped(search_repository): """Reconciliation removes orphan/local stale rows without touching another project.""" From 8fa2182b0ca87ee34c59201b54ac28043ef01b41 Mon Sep 17 00:00:00 2001 From: phernandez Date: Wed, 22 Jul 2026 09:14:25 -0500 Subject: [PATCH 06/24] test(core): cover vector storage invalidation query Signed-off-by: phernandez --- tests/repository/test_postgres_search_repository_unit.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/repository/test_postgres_search_repository_unit.py b/tests/repository/test_postgres_search_repository_unit.py index f6d4bdc90..cbc12bb5f 100644 --- a/tests/repository/test_postgres_search_repository_unit.py +++ b/tests/repository/test_postgres_search_repository_unit.py @@ -199,6 +199,7 @@ async def fake_scoped_session(session_maker): MagicMock(), MagicMock(), MagicMock(), + MagicMock(), ] await repo._ensure_vector_tables() From 29b738901ab0bdc3a836097f3590f8ba21337d1c Mon Sep 17 00:00:00 2001 From: phernandez Date: Wed, 22 Jul 2026 09:41:49 -0500 Subject: [PATCH 07/24] fix(core): harden project vector cleanup Signed-off-by: phernandez --- src/basic_memory/deps/services.py | 10 +++- .../repository/search_repository.py | 2 +- .../repository/search_repository_base.py | 45 +++++++++----- src/basic_memory/services/project_service.py | 19 +++++- tests/repository/test_semantic_search_base.py | 58 +++++++++++++++++++ .../test_project_service_operations.py | 49 ++++++++++++++++ 6 files changed, 166 insertions(+), 17 deletions(-) diff --git a/src/basic_memory/deps/services.py b/src/basic_memory/deps/services.py index 7b972ad38..0f7e63edb 100644 --- a/src/basic_memory/deps/services.py +++ b/src/basic_memory/deps/services.py @@ -42,6 +42,7 @@ LocalDirectoryFileDeleteEnqueuer, ) from basic_memory.repository.accepted_note_repositories import AcceptedNoteRepositories +from basic_memory.repository.search_repository import create_search_repository from basic_memory.index.local_project import ( LocalProjectIndexCommand, LocalProjectIndexRunner, @@ -519,7 +520,14 @@ async def get_project_service( markdown_processor = MarkdownProcessor(entity_parser, app_config=app_config) file_service = FileService(Path.home(), markdown_processor, app_config=app_config) return ProjectService( - repository=project_repository, session_maker=session_maker, file_service=file_service + repository=project_repository, + session_maker=session_maker, + file_service=file_service, + search_repository_factory=lambda project_id: create_search_repository( + session_maker=session_maker, + project_id=project_id, + app_config=app_config, + ), ) diff --git a/src/basic_memory/repository/search_repository.py b/src/basic_memory/repository/search_repository.py index d15b7c36a..5139c481a 100644 --- a/src/basic_memory/repository/search_repository.py +++ b/src/basic_memory/repository/search_repository.py @@ -97,7 +97,7 @@ async def delete_entity_vector_rows(self, entity_id: int) -> None: """Delete semantic vector chunks and embeddings for one entity.""" ... - async def delete_project_vector_rows(self) -> None: + async def delete_project_vector_rows(self, *, strict_adapter_cleanup: bool = False) -> None: """Delete all semantic vector chunks and embeddings for this project.""" ... diff --git a/src/basic_memory/repository/search_repository_base.py b/src/basic_memory/repository/search_repository_base.py index a9bfaaae5..8d9ba1fba 100644 --- a/src/basic_memory/repository/search_repository_base.py +++ b/src/basic_memory/repository/search_repository_base.py @@ -609,7 +609,7 @@ async def delete_entity_vector_rows(self, entity_id: int) -> None: ) ) - async def delete_project_vector_rows(self) -> None: + async def delete_project_vector_rows(self, *, strict_adapter_cleanup: bool = False) -> None: """Delete this project's vectors through the configured storage adapter. Core enumerates ownership from the SQL manifest because the adapter @@ -618,6 +618,8 @@ async def delete_project_vector_rows(self) -> None: so stale ready rows cannot become current if the feature is re-enabled. Adapter cleanup is best-effort because the manifest remains the search authority and an unavailable extension must not block an FTS rebuild. + Project deletion opts into strict cleanup so external data ownership is + preserved for a retry instead of being discarded after an adapter failure. """ async with db.scoped_session(self.session_maker) as session: connection = await session.connection() @@ -627,6 +629,13 @@ async def delete_project_vector_rows(self) -> None: if not manifest_exists: return + manifest_has_embedding_status = await connection.run_sync( + lambda sync_connection: any( + column["name"] == "embedding_status" + for column in inspect(sync_connection).get_columns("search_vector_chunks") + ) + ) + result = await session.execute( text( "SELECT DISTINCT entity_id FROM search_vector_chunks " @@ -635,14 +644,20 @@ async def delete_project_vector_rows(self) -> None: {"project_id": self.project_id}, ) entity_ids = [int(entity_id) for entity_id in result.scalars().all()] - await session.execute( - text( - "UPDATE search_vector_chunks SET embedding_status = 'pending' " - "WHERE project_id = :project_id" - ), - {"project_id": self.project_id}, - ) - await session.commit() + + # Trigger: the manifest predates embedding lifecycle state. + # Why: legacy SQLite schemas must reach cleanup before lazy schema repair runs. + # Outcome: skip staging only for that obsolete schema; the manifest is still + # deleted below, and later semantic initialization recreates the current schema. + if manifest_has_embedding_status: + await session.execute( + text( + "UPDATE search_vector_chunks SET embedding_status = 'pending' " + "WHERE project_id = :project_id" + ), + {"project_id": self.project_id}, + ) + await session.commit() if hasattr(self, "_semantic_vector_index"): try: @@ -651,17 +666,19 @@ async def delete_project_vector_rows(self) -> None: await self._semantic_vector_index.delete_entity(entity_id) except Exception as exc: # Trigger: a configured external adapter cannot initialize or delete. - # Why: SQL is the authoritative hydration manifest, so clearing it - # prevents stale search results even when external cleanup is unavailable. - # Outcome: finish the FTS rebuild and leave a visible warning about - # storage that may need later adapter reconciliation. + # Why: reindex may discard derived state, but project deletion must + # not discard the only ownership manifest for external data. + # Outcome: strict callers stop for a retry; reindex logs the adapter + # failure and continues with a non-searchable empty manifest. logger.warning( - "Could not clean semantic vector adapter during full reindex: " + "Could not clean semantic vector adapter: " "project_id={project_id} vector_index={vector_index} error={error}", project_id=self.project_id, vector_index=self._semantic_vector_index_name, error=exc, ) + if strict_adapter_cleanup: + raise async with db.scoped_session(self.session_maker) as session: await session.execute( diff --git a/src/basic_memory/services/project_service.py b/src/basic_memory/services/project_service.py index 64e047cde..4714b144b 100644 --- a/src/basic_memory/services/project_service.py +++ b/src/basic_memory/services/project_service.py @@ -4,6 +4,7 @@ import json import os import shutil +from collections.abc import Callable from datetime import datetime from pathlib import Path from typing import TYPE_CHECKING, Dict, Optional, Sequence @@ -16,6 +17,7 @@ from basic_memory import db from basic_memory.models import Project from basic_memory.repository.project_repository import ProjectRepository +from basic_memory.repository.search_repository import SearchRepository from basic_memory.repository.embedding_provider_factory import ( configured_embedding_provider_identity, ) @@ -42,6 +44,8 @@ if TYPE_CHECKING: # pragma: no cover from basic_memory.services.file_service import FileService +type ProjectSearchRepositoryFactory = Callable[[int], SearchRepository] + class ProjectService: """Service for managing Basic Memory projects.""" @@ -53,12 +57,14 @@ def __init__( repository: ProjectRepository, session_maker: async_sessionmaker[AsyncSession], file_service: Optional["FileService"] = None, + search_repository_factory: ProjectSearchRepositoryFactory | None = None, ): """Initialize the project service.""" super().__init__() self.repository = repository self.session_maker = session_maker self.file_service = file_service + self._search_repository_factory = search_repository_factory @property def config_manager(self) -> ConfigManager: @@ -335,6 +341,7 @@ async def remove_project(self, name: str, delete_notes: bool = False) -> None: if not project: raise ValueError(f"Project '{name}' not found") # pragma: no cover + project_id = project.id project_path = project.path # Check if project is default @@ -346,6 +353,16 @@ async def remove_project(self, name: str, delete_notes: bool = False) -> None: if is_default: raise ValueError(f"Cannot remove the default project '{name}'") # pragma: no cover + # Trigger: project deletion can remove the only SQL ownership manifest for + # vectors stored by an extension such as Milvus. + # Why: external storage has no database cascade and cannot reconcile after + # the project manifest disappears. + # Outcome: delete adapter-owned vectors while project ownership is still known. + if self._search_repository_factory is not None: + search_repository = self._search_repository_factory(project_id) + await search_repository.delete_project_vector_rows(strict_adapter_cleanup=True) + + async with db.scoped_session(self.session_maker) as session: # Remove from config if it exists there (may not exist in cloud mode) try: self.config_manager.remove_project(name) @@ -356,7 +373,7 @@ async def remove_project(self, name: str, delete_notes: bool = False) -> None: ) # Remove from database - await self.repository.delete(session, project.id) + await self.repository.delete(session, project_id) logger.info(f"Project '{name}' removed from configuration and database") diff --git a/tests/repository/test_semantic_search_base.py b/tests/repository/test_semantic_search_base.py index 69b966c1a..a3a2ccd95 100644 --- a/tests/repository/test_semantic_search_base.py +++ b/tests/repository/test_semantic_search_base.py @@ -184,6 +184,31 @@ async def fake_scoped_session(_session_maker): assert any(statement.startswith("DELETE FROM search_vector_chunks") for statement in statements) +@pytest.mark.asyncio +async def test_project_vector_cleanup_handles_legacy_manifest_without_status(monkeypatch): + """Legacy vector manifests should be removed without querying absent status columns.""" + repo = _ConcreteRepo() + session = AsyncMock() + connection = AsyncMock() + connection.run_sync.side_effect = [True, False] + session.connection.return_value = connection + session.execute.return_value = SimpleNamespace( + scalars=lambda: SimpleNamespace(all=lambda: [41]) + ) + + @asynccontextmanager + async def fake_scoped_session(_session_maker): + yield session + + monkeypatch.setattr(search_repository_base_module.db, "scoped_session", fake_scoped_session) + + await repo.delete_project_vector_rows() + + statements = [str(call.args[0]) for call in session.execute.await_args_list] + assert not any("embedding_status" in statement for statement in statements) + assert any(statement.startswith("DELETE FROM search_vector_chunks") for statement in statements) + + @pytest.mark.asyncio async def test_project_vector_cleanup_uses_available_adapter(monkeypatch): """An available adapter should receive every manifest-owned entity deletion.""" @@ -251,6 +276,39 @@ async def fake_scoped_session(_session_maker): assert any(statement.startswith("DELETE FROM search_vector_chunks") for statement in statements) +@pytest.mark.asyncio +async def test_strict_project_vector_cleanup_preserves_manifest_after_adapter_failure(monkeypatch): + """Project deletion should retain ownership when its external adapter is unavailable.""" + repo = _ConcreteRepo() + adapter: Any = SimpleNamespace( + initialize=AsyncMock(side_effect=RuntimeError("adapter unavailable")), + delete_entity=AsyncMock(), + ) + repo._semantic_vector_index = adapter + repo._semantic_vector_index_name = "milvus" + session = AsyncMock() + connection = AsyncMock() + connection.run_sync.return_value = True + session.connection.return_value = connection + session.execute.return_value = SimpleNamespace( + scalars=lambda: SimpleNamespace(all=lambda: [41]) + ) + + @asynccontextmanager + async def fake_scoped_session(_session_maker): + yield session + + monkeypatch.setattr(search_repository_base_module.db, "scoped_session", fake_scoped_session) + + with pytest.raises(RuntimeError, match="adapter unavailable"): + await repo.delete_project_vector_rows(strict_adapter_cleanup=True) + + statements = [str(call.args[0]) for call in session.execute.await_args_list] + assert not any( + statement.startswith("DELETE FROM search_vector_chunks") for statement in statements + ) + + @pytest.mark.asyncio async def test_sync_entity_vectors_batch_flushes_at_configured_threshold(monkeypatch): """Batch sync should flush queued jobs at semantic_embedding_sync_batch_size boundaries.""" diff --git a/tests/services/test_project_service_operations.py b/tests/services/test_project_service_operations.py index b6bf8f84b..7aa1219f1 100644 --- a/tests/services/test_project_service_operations.py +++ b/tests/services/test_project_service_operations.py @@ -3,6 +3,8 @@ import os import tempfile from pathlib import Path +from types import SimpleNamespace +from unittest.mock import AsyncMock, Mock import pytest @@ -73,3 +75,50 @@ async def test_add_project_to_config(project_service: ProjectService, config_man # Clean up if test_project_name in project_service.projects: config_manager.remove_project(test_project_name) + + +@pytest.mark.asyncio +async def test_remove_project_cleans_external_vectors_before_database_delete( + project_service: ProjectService, + monkeypatch, +): + """Project removal must clean adapter storage before deleting SQL ownership.""" + project_name = f"external-vector-project-{os.urandom(4).hex()}" + search_repository = SimpleNamespace(delete_project_vector_rows=AsyncMock()) + search_repository_factory = Mock(return_value=search_repository) + service = ProjectService( + repository=project_service.repository, + session_maker=project_service.session_maker, + file_service=project_service.file_service, + search_repository_factory=search_repository_factory, + ) + + with tempfile.TemporaryDirectory() as temp_dir: + async with db.scoped_session(service.session_maker) as session: + project = await service.repository.create( + session, + { + "name": project_name, + "path": temp_dir, + "permalink": project_name, + "is_active": True, + }, + ) + project_id = project.id + + original_delete = service.repository.delete + + async def delete_after_vector_cleanup(session, entity_id: int) -> bool: + search_repository.delete_project_vector_rows.assert_awaited_once_with( + strict_adapter_cleanup=True + ) + return await original_delete(session, entity_id) + + monkeypatch.setattr(service.repository, "delete", delete_after_vector_cleanup) + + await service.remove_project(project_name) + + search_repository_factory.assert_called_once_with(project_id) + search_repository.delete_project_vector_rows.assert_awaited_once_with( + strict_adapter_cleanup=True + ) From 3a9dda503cb92cf98998e21da3f4db268e74d112 Mon Sep 17 00:00:00 2001 From: phernandez Date: Wed, 22 Jul 2026 09:59:12 -0500 Subject: [PATCH 08/24] fix(core): preserve vector manifests without adapters Signed-off-by: phernandez --- .../repository/search_repository_base.py | 13 ++++++++ tests/repository/test_semantic_search_base.py | 32 ++++++++++++++++++- 2 files changed, 44 insertions(+), 1 deletion(-) diff --git a/src/basic_memory/repository/search_repository_base.py b/src/basic_memory/repository/search_repository_base.py index 8d9ba1fba..fbd7dfa3c 100644 --- a/src/basic_memory/repository/search_repository_base.py +++ b/src/basic_memory/repository/search_repository_base.py @@ -31,6 +31,7 @@ from basic_memory.repository.semantic_errors import ( SemanticDependenciesMissingError, SemanticSearchDisabledError, + SemanticVectorIndexExtensionError, ) from basic_memory.repository.semantic_vector_index import ( SemanticVectorIndex, @@ -659,6 +660,18 @@ async def delete_project_vector_rows(self, *, strict_adapter_cleanup: bool = Fal ) await session.commit() + # Trigger: project deletion requires strict cleanup, but semantic search + # was disabled before this repository was composed. + # Why: deleting the manifest without an adapter would orphan externally + # stored vectors and discard the only durable ownership list. + # Outcome: retain the manifest and require the caller to restore the + # configured adapter before retrying project deletion. + if entity_ids and strict_adapter_cleanup and not hasattr(self, "_semantic_vector_index"): + raise SemanticVectorIndexExtensionError( + "Cannot delete project vectors because the configured semantic vector " + "adapter is unavailable. Enable semantic search and retry project deletion." + ) + if hasattr(self, "_semantic_vector_index"): try: await self._semantic_vector_index.initialize() diff --git a/tests/repository/test_semantic_search_base.py b/tests/repository/test_semantic_search_base.py index a3a2ccd95..78dc5dde8 100644 --- a/tests/repository/test_semantic_search_base.py +++ b/tests/repository/test_semantic_search_base.py @@ -17,7 +17,10 @@ _PreparedEntityVectorSync, ) from basic_memory.repository.sqlite_search_repository import SQLiteSearchRepository -from basic_memory.repository.semantic_errors import SemanticSearchDisabledError +from basic_memory.repository.semantic_errors import ( + SemanticSearchDisabledError, + SemanticVectorIndexExtensionError, +) from basic_memory.schemas.search import SearchItemType, SearchRetrievalMode @@ -309,6 +312,33 @@ async def fake_scoped_session(_session_maker): ) +@pytest.mark.asyncio +async def test_strict_project_vector_cleanup_preserves_manifest_without_adapter(monkeypatch): + """Disabled semantic search must retain external-vector ownership for retry.""" + repo = _ConcreteRepo() + session = AsyncMock() + connection = AsyncMock() + connection.run_sync.return_value = True + session.connection.return_value = connection + session.execute.return_value = SimpleNamespace( + scalars=lambda: SimpleNamespace(all=lambda: [41]) + ) + + @asynccontextmanager + async def fake_scoped_session(_session_maker): + yield session + + monkeypatch.setattr(search_repository_base_module.db, "scoped_session", fake_scoped_session) + + with pytest.raises(SemanticVectorIndexExtensionError, match="adapter is unavailable"): + await repo.delete_project_vector_rows(strict_adapter_cleanup=True) + + statements = [str(call.args[0]) for call in session.execute.await_args_list] + assert not any( + statement.startswith("DELETE FROM search_vector_chunks") for statement in statements + ) + + @pytest.mark.asyncio async def test_sync_entity_vectors_batch_flushes_at_configured_threshold(monkeypatch): """Batch sync should flush queued jobs at semantic_embedding_sync_batch_size boundaries.""" From fdc75d7f6aaa9613fd7420145809c0d7e9f84f34 Mon Sep 17 00:00:00 2001 From: phernandez Date: Wed, 22 Jul 2026 10:16:00 -0500 Subject: [PATCH 09/24] fix(core): close vector lifecycle gaps Signed-off-by: phernandez --- src/basic_memory/deps/services.py | 12 ++- src/basic_memory/index/local_dependencies.py | 5 + src/basic_memory/index/local_runtime.py | 1 + .../indexing/directory_delete_runner.py | 25 +++-- .../indexing/project_index_maintenance.py | 23 ++++- .../indexing/project_index_runtime.py | 5 + .../repository/accepted_note_repositories.py | 18 +++- .../accepted_note_search_repository.py | 31 +++++-- .../accepted_note_vector_cleanup.py | 49 +++++++++- .../repository/search_repository.py | 17 +++- .../repository/search_repository_base.py | 28 +++++- src/basic_memory/services/project_service.py | 20 +++- .../test_external_vector_cleanup.py | 91 +++++++++++++++++++ tests/repository/test_semantic_search_base.py | 33 +++++++ .../repository/test_semantic_vector_index.py | 19 ++++ .../test_project_service_embedding_status.py | 42 +++++++++ 16 files changed, 392 insertions(+), 27 deletions(-) create mode 100644 tests/repository/test_external_vector_cleanup.py diff --git a/src/basic_memory/deps/services.py b/src/basic_memory/deps/services.py index 0f7e63edb..eba5a0836 100644 --- a/src/basic_memory/deps/services.py +++ b/src/basic_memory/deps/services.py @@ -171,7 +171,9 @@ async def get_directory_delete_service( return DirectoryDeleteService( session_maker=session_maker, runtime=DirectoryDeleteRuntime( - store=RepositoryDirectoryDeleteAcceptanceStore(), + store=RepositoryDirectoryDeleteAcceptanceStore( + external_vector_cleaner=search_service.repository + ), file_delete_enqueuer=LocalDirectoryFileDeleteEnqueuer(file_service=file_service), relation_cleanup_refresher=LocalDirectoryDeleteRelationCleanupRefresher( session_maker=session_maker, @@ -312,7 +314,13 @@ async def get_note_content_mutation_service( app_config: AppConfigDep, ) -> NoteContentMutationService: """Create the local accepted-note mutation facade for API routes.""" - accepted_note_repositories = AcceptedNoteRepositories() + accepted_note_repositories = AcceptedNoteRepositories( + external_vector_cleaner_factory=lambda project_id: create_search_repository( + session_maker=session_maker, + project_id=project_id, + app_config=app_config, + ) + ) return NoteContentMutationService( session_maker=session_maker, mutation_dependencies=AcceptedNoteMutationDependencies( diff --git a/src/basic_memory/index/local_dependencies.py b/src/basic_memory/index/local_dependencies.py index 045fad596..834fdf246 100644 --- a/src/basic_memory/index/local_dependencies.py +++ b/src/basic_memory/index/local_dependencies.py @@ -63,6 +63,9 @@ ObservationRepository, RelationRepository, ) +from basic_memory.repository.accepted_note_vector_cleanup import ( + ProjectIndexExternalVectorCleaner, +) from basic_memory.repository.search_repository import create_search_repository from basic_memory.runtime.storage import ProjectId, RuntimeFilePath from basic_memory.services import EntityService, FileService @@ -206,6 +209,7 @@ class LocalIndexProjectDependencies: link_resolver: RelationResolutionLinkResolver search_service: LocalIndexSearchService entity_service: LocalIndexEntityService + external_vector_cleaner: ProjectIndexExternalVectorCleaner | None = None class LocalIndexProjectDependencyProvider(Protocol): @@ -658,4 +662,5 @@ async def build_local_index_project_dependencies( link_resolver=link_resolver, search_service=search_service, entity_service=entity_service, + external_vector_cleaner=search_repository, ) diff --git a/src/basic_memory/index/local_runtime.py b/src/basic_memory/index/local_runtime.py index 1582024a9..989f37657 100644 --- a/src/basic_memory/index/local_runtime.py +++ b/src/basic_memory/index/local_runtime.py @@ -265,6 +265,7 @@ async def runtime_for_project(self, project: Project) -> StorageEventIndexRuntim maintenance_store = RepositoryProjectIndexMaintenanceStore( session_maker=dependencies.session_maker, project_id=dependencies.project_id, + external_vector_cleaner=dependencies.external_vector_cleaner, move_content_updater=LocalProjectIndexMoveContentUpdater( entity_service=dependencies.entity_service, file_service=dependencies.file_service, diff --git a/src/basic_memory/indexing/directory_delete_runner.py b/src/basic_memory/indexing/directory_delete_runner.py index ac8d40aa2..bec6c1f52 100644 --- a/src/basic_memory/indexing/directory_delete_runner.py +++ b/src/basic_memory/indexing/directory_delete_runner.py @@ -10,8 +10,11 @@ from sqlalchemy import bindparam, delete, select, text from sqlalchemy.ext.asyncio import AsyncSession -from basic_memory.indexing.project_index_maintenance import delete_project_index_vector_rows from basic_memory.models import Entity, NoteContent, Project, Relation +from basic_memory.repository.accepted_note_vector_cleanup import ( + ProjectIndexExternalVectorCleaner, + delete_project_index_vector_rows, +) from basic_memory.runtime.cleanup import ( RuntimeDeleteStatus, RuntimeDirectoryFileSnapshot, @@ -144,6 +147,8 @@ class DirectoryDeleteRuntime: class RepositoryDirectoryDeleteAcceptanceStore: """Repository-backed directory-delete acceptance store.""" + external_vector_cleaner: ProjectIndexExternalVectorCleaner | None = None + async def load_project_id( self, session: AsyncSession, @@ -235,11 +240,19 @@ async def delete_directory_entities( ).bindparams(bindparam("entity_ids", expanding=True)), {"project_id": project_id, "entity_ids": deleted_entity_ids}, ) - await delete_project_index_vector_rows( - session, - project_id=project_id, - entity_ids=deleted_entity_ids, - ) + if self.external_vector_cleaner is None: + await delete_project_index_vector_rows( + session, + project_id=project_id, + entity_ids=deleted_entity_ids, + ) + else: + await delete_project_index_vector_rows( + session, + project_id=project_id, + entity_ids=deleted_entity_ids, + external_vector_cleaner=self.external_vector_cleaner, + ) await session.execute(delete(Entity).where(Entity.id.in_(deleted_entity_ids))) return relation_cleanup_entity_ids diff --git a/src/basic_memory/indexing/project_index_maintenance.py b/src/basic_memory/indexing/project_index_maintenance.py index 428607dd4..324e9b868 100644 --- a/src/basic_memory/indexing/project_index_maintenance.py +++ b/src/basic_memory/indexing/project_index_maintenance.py @@ -13,6 +13,7 @@ from basic_memory import db from basic_memory.models import Entity, NoteContent, Relation from basic_memory.repository.accepted_note_vector_cleanup import ( + ProjectIndexExternalVectorCleaner, delete_project_index_vector_rows, ) from basic_memory.runtime.storage import ProjectId @@ -337,6 +338,7 @@ async def delete_project_index_entities( *, project_id: ProjectId, entity_ids: Sequence[int], + external_vector_cleaner: ProjectIndexExternalVectorCleaner | None = None, ) -> frozenset[int]: """Delete indexed entities and return surviving relation sources needing repair.""" deleted_entity_ids = tuple(entity_ids) @@ -362,11 +364,19 @@ async def delete_project_index_entities( "relation_row_type": "relation", } await session.execute(DELETE_PROJECT_INDEX_SEARCH_ROWS_SQL, delete_params) - await delete_project_index_vector_rows( - session, - project_id=project_id, - entity_ids=deleted_entity_ids, - ) + if external_vector_cleaner is None: + await delete_project_index_vector_rows( + session, + project_id=project_id, + entity_ids=deleted_entity_ids, + ) + else: + await delete_project_index_vector_rows( + session, + project_id=project_id, + entity_ids=deleted_entity_ids, + external_vector_cleaner=external_vector_cleaner, + ) await session.execute( delete(Entity).where( Entity.project_id == project_id, @@ -542,6 +552,7 @@ class RepositoryProjectIndexMaintenanceStore: session_maker: async_sessionmaker[AsyncSession] project_id: ProjectId + external_vector_cleaner: ProjectIndexExternalVectorCleaner | None = None move_content_updater: ProjectIndexMoveContentUpdater | None = None delete_path_verifier: ProjectIndexDeletePathVerifier = TrustPlannedProjectIndexDeleteVerifier() # Trigger: an entity occupies a move destination at apply time. @@ -616,6 +627,7 @@ async def apply_project_index_move_batch( session, project_id=self.project_id, entity_ids=tuple(replaced_entity_ids), + external_vector_cleaner=self.external_vector_cleaner, ) await self._execute_move_batch_updates( session, @@ -854,6 +866,7 @@ async def apply_project_index_delete_batch( session, project_id=self.project_id, entity_ids=deleted_entity_ids, + external_vector_cleaner=self.external_vector_cleaner, ) return ProjectIndexDeleteBatchResult( diff --git a/src/basic_memory/indexing/project_index_runtime.py b/src/basic_memory/indexing/project_index_runtime.py index 9844a3783..fd7a1806c 100644 --- a/src/basic_memory/indexing/project_index_runtime.py +++ b/src/basic_memory/indexing/project_index_runtime.py @@ -46,6 +46,9 @@ RepositoryProjectIndexMaintenanceStore, StoreProjectIndexMaintenanceRunner, ) +from basic_memory.repository.accepted_note_vector_cleanup import ( + ProjectIndexExternalVectorCleaner, +) from basic_memory.runtime.storage import ProjectId if TYPE_CHECKING: # pragma: no cover @@ -169,6 +172,7 @@ def build_default_project_index_runtime( vector_sync: VectorSyncExecutor, entity_repository: RelationResolutionEntityRepository, entity_indexer: RelationResolutionEntityIndexer, + external_vector_cleaner: ProjectIndexExternalVectorCleaner | None = None, ) -> ProjectIndexRuntime: """Compose the default repository-backed project-index runtime.""" vector_entity_source = RepositoryVectorSyncEntitySource( @@ -178,6 +182,7 @@ def build_default_project_index_runtime( maintenance_store = RepositoryProjectIndexMaintenanceStore( session_maker=session_maker, project_id=project_id, + external_vector_cleaner=external_vector_cleaner, ) return ProjectIndexRuntime( project_id=project_id, diff --git a/src/basic_memory/repository/accepted_note_repositories.py b/src/basic_memory/repository/accepted_note_repositories.py index beb58efb2..07678cabc 100644 --- a/src/basic_memory/repository/accepted_note_repositories.py +++ b/src/basic_memory/repository/accepted_note_repositories.py @@ -1,9 +1,13 @@ """Project-scoped repositories for accepted-note mutations.""" +from collections.abc import Callable from dataclasses import dataclass from basic_memory.repository import NoteContentRepository, ObservationRepository, RelationRepository from basic_memory.repository.accepted_note_search_repository import AcceptedNoteSearchRepository +from basic_memory.repository.accepted_note_vector_cleanup import ( + ProjectIndexExternalVectorCleaner, +) from basic_memory.repository.entity_repository import EntityRepository from basic_memory.runtime.storage import ProjectId @@ -12,6 +16,10 @@ class AcceptedNoteRepositories: """Core repository bundle using the caller-owned transaction.""" + external_vector_cleaner_factory: ( + Callable[[ProjectId], ProjectIndexExternalVectorCleaner] | None + ) = None + def entity_repository(self, project_id: ProjectId) -> EntityRepository: return EntityRepository(project_id=project_id) @@ -22,7 +30,15 @@ def note_content_repository(self, project_id: ProjectId) -> NoteContentRepositor return NoteContentRepository(project_id=project_id) def search_repository(self, project_id: ProjectId) -> AcceptedNoteSearchRepository: - return AcceptedNoteSearchRepository(project_id=project_id) + external_vector_cleaner = ( + self.external_vector_cleaner_factory(project_id) + if self.external_vector_cleaner_factory is not None + else None + ) + return AcceptedNoteSearchRepository( + project_id=project_id, + external_vector_cleaner=external_vector_cleaner, + ) def observation_repository(self, project_id: ProjectId) -> ObservationRepository: return ObservationRepository(project_id=project_id) diff --git a/src/basic_memory/repository/accepted_note_search_repository.py b/src/basic_memory/repository/accepted_note_search_repository.py index a64ff3059..30d06bdd9 100644 --- a/src/basic_memory/repository/accepted_note_search_repository.py +++ b/src/basic_memory/repository/accepted_note_search_repository.py @@ -9,7 +9,10 @@ from sqlalchemy.ext.asyncio import AsyncSession from basic_memory.repository.accepted_note_search_row import AcceptedNoteSearchRow -from basic_memory.repository.accepted_note_vector_cleanup import delete_project_index_vector_rows +from basic_memory.repository.accepted_note_vector_cleanup import ( + ProjectIndexExternalVectorCleaner, + delete_project_index_vector_rows, +) type SearchIndexSqlValue = str | int | datetime | None type SearchIndexSqlParams = dict[str, SearchIndexSqlValue] @@ -106,8 +109,14 @@ def accepted_note_search_insert_params( class AcceptedNoteSearchRepository: """Explicit-session repository for accepted-note hot search refreshes.""" - def __init__(self, *, project_id: int) -> None: + def __init__( + self, + *, + project_id: int, + external_vector_cleaner: ProjectIndexExternalVectorCleaner | None = None, + ) -> None: self.project_id = project_id + self.external_vector_cleaner = external_vector_cleaner async def refresh_entity( self, @@ -147,8 +156,16 @@ async def delete_entity_vectors( entity_id: int, ) -> None: """Delete semantic vector rows for one accepted-note entity.""" - await delete_project_index_vector_rows( - session, - project_id=self.project_id, - entity_ids=(entity_id,), - ) + if self.external_vector_cleaner is None: + await delete_project_index_vector_rows( + session, + project_id=self.project_id, + entity_ids=(entity_id,), + ) + else: + await delete_project_index_vector_rows( + session, + project_id=self.project_id, + entity_ids=(entity_id,), + external_vector_cleaner=self.external_vector_cleaner, + ) diff --git a/src/basic_memory/repository/accepted_note_vector_cleanup.py b/src/basic_memory/repository/accepted_note_vector_cleanup.py index 14f1f8b18..7961d4184 100644 --- a/src/basic_memory/repository/accepted_note_vector_cleanup.py +++ b/src/basic_memory/repository/accepted_note_vector_cleanup.py @@ -1,11 +1,13 @@ """Repository-owned cleanup for accepted-note vector search rows.""" from collections.abc import Sequence +from typing import Protocol from sqlalchemy import bindparam, text from sqlalchemy.ext.asyncio import AsyncSession from basic_memory.repository.project_repository import _load_sqlite_vec_on_session +from basic_memory.repository.semantic_errors import SemanticVectorIndexExtensionError from basic_memory.runtime.storage import ProjectId @@ -49,6 +51,25 @@ ) """).bindparams(bindparam("deleted_entity_ids", expanding=True)) +SELECT_PROJECT_INDEX_EXTERNAL_VECTOR_INDEXES_SQL = text(""" + SELECT DISTINCT vector_index + FROM search_vector_chunks + WHERE project_id = :project_id + AND entity_id IN :deleted_entity_ids + AND vector_index <> 'pgvector' +""").bindparams(bindparam("deleted_entity_ids", expanding=True)) + + +class ProjectIndexExternalVectorCleaner(Protocol): + """Project-scoped capability for deleting vectors from extension storage.""" + + async def delete_external_entity_vectors( + self, + entity_ids: Sequence[int], + *, + vector_index_names: frozenset[str], + ) -> None: ... + def project_index_session_dialect_name(session: AsyncSession) -> str: """Return the SQLAlchemy dialect name for project-index maintenance.""" @@ -73,6 +94,7 @@ async def delete_project_index_vector_rows( *, project_id: ProjectId, entity_ids: Sequence[int], + external_vector_cleaner: ProjectIndexExternalVectorCleaner | None = None, ) -> None: """Delete backend vector rows for project-index entity deletes when tables exist.""" deleted_entity_ids = tuple(entity_ids) @@ -87,8 +109,33 @@ async def delete_project_index_vector_rows( "project_id": project_id, "deleted_entity_ids": deleted_entity_ids, } + dialect_name = project_index_session_dialect_name(session) + + # Trigger: the manifest says some vectors live outside PostgreSQL. + # Why: deleting the manifest first would discard the only durable ownership + # list and leave extension data with no retry path. + # Outcome: require the project-scoped adapter to delete those entities before + # the caller-owned transaction removes built-in vectors and manifest rows. + if dialect_name == "postgresql": + external_result = await session.execute( + SELECT_PROJECT_INDEX_EXTERNAL_VECTOR_INDEXES_SQL, + delete_params, + ) + external_vector_indexes = frozenset( + str(vector_index) for vector_index in external_result.scalars() + ) + if external_vector_indexes: + if external_vector_cleaner is None: + raise SemanticVectorIndexExtensionError( + "Cannot delete externally indexed entity vectors without a " + "project-scoped semantic vector adapter." + ) + await external_vector_cleaner.delete_external_entity_vectors( + deleted_entity_ids, + vector_index_names=external_vector_indexes, + ) + if "search_vector_embeddings" in vector_table_names: - dialect_name = project_index_session_dialect_name(session) if dialect_name == "sqlite": if await _load_sqlite_vec_on_session(session): await session.execute( diff --git a/src/basic_memory/repository/search_repository.py b/src/basic_memory/repository/search_repository.py index 5139c481a..62b640a9b 100644 --- a/src/basic_memory/repository/search_repository.py +++ b/src/basic_memory/repository/search_repository.py @@ -6,6 +6,7 @@ - PostgresSearchRepository: Uses tsvector/tsquery with GIN indexes """ +from collections.abc import Sequence from datetime import datetime from typing import Any, Callable, List, Optional, Protocol @@ -17,7 +18,10 @@ from basic_memory.repository.postgres_search_repository import PostgresSearchRepository from basic_memory.repository.search_index_row import SearchIndexRow from basic_memory.repository.search_repository_base import VectorSyncBatchResult -from basic_memory.repository.semantic_vector_index_factory import create_semantic_vector_index +from basic_memory.repository.semantic_vector_index_factory import ( + create_semantic_vector_index, + resolve_semantic_vector_index_name, +) from basic_memory.repository.sqlite_search_repository import SQLiteSearchRepository from basic_memory.schemas.search import SearchItemType, SearchRetrievalMode @@ -97,6 +101,15 @@ async def delete_entity_vector_rows(self, entity_id: int) -> None: """Delete semantic vector chunks and embeddings for one entity.""" ... + async def delete_external_entity_vectors( + self, + entity_ids: Sequence[int], + *, + vector_index_names: frozenset[str], + ) -> None: + """Delete DB-first entity vectors through an extension adapter.""" + ... + async def delete_project_vector_rows(self, *, strict_adapter_cleanup: bool = False) -> None: """Delete all semantic vector chunks and embeddings for this project.""" ... @@ -151,7 +164,7 @@ def create_search_repository( # Outcome: resolve the cached singleton here once and inject it, so the provider # is the single source of truth across all callers of this factory. embedding_provider = None - vector_index_name = None + vector_index_name = resolve_semantic_vector_index_name(config, database_backend) vector_index = None if config.semantic_search_enabled: embedding_provider = create_embedding_provider(config) diff --git a/src/basic_memory/repository/search_repository_base.py b/src/basic_memory/repository/search_repository_base.py index fbd7dfa3c..7c53854b1 100644 --- a/src/basic_memory/repository/search_repository_base.py +++ b/src/basic_memory/repository/search_repository_base.py @@ -2,7 +2,7 @@ import time from abc import ABC, abstractmethod -from collections.abc import Iterable +from collections.abc import Iterable, Sequence from contextlib import asynccontextmanager from dataclasses import replace from datetime import datetime @@ -610,6 +610,32 @@ async def delete_entity_vector_rows(self, entity_id: int) -> None: ) ) + async def delete_external_entity_vectors( + self, + entity_ids: Sequence[int], + *, + vector_index_names: frozenset[str], + ) -> None: + """Delete DB-first entity vectors through the configured extension adapter.""" + if not vector_index_names: + return + configured_index = self._semantic_vector_index_name + if vector_index_names != frozenset({configured_index}): + raise SemanticVectorIndexExtensionError( + "Cannot delete external vectors owned by " + f"{sorted(vector_index_names)!r} with configured adapter " + f"{configured_index!r}." + ) + if not hasattr(self, "_semantic_vector_index"): + raise SemanticVectorIndexExtensionError( + f"Semantic vector adapter {configured_index!r} is unavailable. " + "Enable semantic search and retry the entity deletion." + ) + + await self._semantic_vector_index.initialize() + for entity_id in entity_ids: + await self._semantic_vector_index.delete_entity(entity_id) + async def delete_project_vector_rows(self, *, strict_adapter_cleanup: bool = False) -> None: """Delete this project's vectors through the configured storage adapter. diff --git a/src/basic_memory/services/project_service.py b/src/basic_memory/services/project_service.py index 4714b144b..fd3850070 100644 --- a/src/basic_memory/services/project_service.py +++ b/src/basic_memory/services/project_service.py @@ -1083,7 +1083,21 @@ async def get_embedding_status(self, project_id: int) -> EmbeddingStatus: table_result = await self.repository.execute_query(session, table_check_sql, {}) vector_tables_exist = (table_result.scalar() or 0) == 1 - if not vector_tables_exist: + manifest_schema_current = vector_tables_exist + if vector_tables_exist and not is_postgres: + columns_result = await self.repository.execute_query( + session, + text("PRAGMA table_info(search_vector_chunks)"), + {}, + ) + manifest_columns = {str(row[1]) for row in columns_result.fetchall()} + manifest_schema_current = { + "embedding_model", + "vector_index", + "embedding_status", + }.issubset(manifest_columns) + + if not manifest_schema_current: # Count distinct entities in search index for the recommendation message si_result = await self.repository.execute_query( session, @@ -1106,7 +1120,9 @@ async def get_embedding_status(self, project_id: int) -> EmbeddingStatus: vector_tables_exist=False, reindex_recommended=True, reindex_reason=( - "Vector manifest not initialized — run: bm reindex --embeddings" + "Vector manifest schema is outdated — run: bm reindex --embeddings" + if vector_tables_exist + else "Vector manifest not initialized — run: bm reindex --embeddings" ), ) diff --git a/tests/repository/test_external_vector_cleanup.py b/tests/repository/test_external_vector_cleanup.py new file mode 100644 index 000000000..62b48f232 --- /dev/null +++ b/tests/repository/test_external_vector_cleanup.py @@ -0,0 +1,91 @@ +"""Tests for DB-first cleanup of externally stored semantic vectors.""" + +from collections.abc import Sequence +from typing import Any, cast + +import pytest +from sqlalchemy.ext.asyncio import AsyncSession + +from basic_memory.repository.accepted_note_vector_cleanup import ( + delete_project_index_vector_rows, +) +from basic_memory.repository.semantic_errors import SemanticVectorIndexExtensionError + + +class _ScalarResult: + def __init__(self, values: Sequence[object] = ()) -> None: + self._values = values + + def scalars(self) -> Sequence[object]: + return self._values + + +class _PostgresSession: + def __init__(self, results: Sequence[_ScalarResult]) -> None: + self._results = iter(results) + self.executed: list[tuple[str, dict[str, object]]] = [] + self._bind = type("Bind", (), {"dialect": type("Dialect", (), {"name": "postgresql"})()})() + + def get_bind(self) -> object: + return self._bind + + async def execute(self, statement: object, params: dict[str, object] | None = None) -> Any: + self.executed.append((str(statement), params or {})) + return next(self._results) + + +class _RecordingExternalCleaner: + def __init__(self) -> None: + self.calls: list[tuple[tuple[int, ...], frozenset[str]]] = [] + + async def delete_external_entity_vectors( + self, + entity_ids: Sequence[int], + *, + vector_index_names: frozenset[str], + ) -> None: + self.calls.append((tuple(entity_ids), vector_index_names)) + + +@pytest.mark.asyncio +async def test_external_vectors_are_deleted_before_their_manifest() -> None: + session = _PostgresSession( + [ + _ScalarResult(["search_vector_chunks"]), + _ScalarResult(["milvus"]), + _ScalarResult(), + ] + ) + cleaner = _RecordingExternalCleaner() + + await delete_project_index_vector_rows( + cast(AsyncSession, session), + project_id=7, + entity_ids=[41, 42], + external_vector_cleaner=cleaner, + ) + + assert cleaner.calls == [((41, 42), frozenset({"milvus"}))] + assert str(session.executed[-1][0]).lstrip().startswith("DELETE FROM search_vector_chunks") + + +@pytest.mark.asyncio +async def test_external_vector_manifest_is_preserved_without_a_cleaner() -> None: + session = _PostgresSession( + [ + _ScalarResult(["search_vector_chunks"]), + _ScalarResult(["milvus"]), + ] + ) + + with pytest.raises(SemanticVectorIndexExtensionError, match="project-scoped"): + await delete_project_index_vector_rows( + cast(AsyncSession, session), + project_id=7, + entity_ids=[41], + ) + + assert not any( + statement.lstrip().startswith("DELETE FROM search_vector_chunks") + for statement, _params in session.executed + ) diff --git a/tests/repository/test_semantic_search_base.py b/tests/repository/test_semantic_search_base.py index 78dc5dde8..61bdbaca3 100644 --- a/tests/repository/test_semantic_search_base.py +++ b/tests/repository/test_semantic_search_base.py @@ -339,6 +339,39 @@ async def fake_scoped_session(_session_maker): ) +@pytest.mark.asyncio +async def test_external_entity_cleanup_uses_matching_project_adapter() -> None: + """DB-first deletes should invoke the configured extension before manifest removal.""" + repo = _ConcreteRepo() + adapter: Any = SimpleNamespace( + initialize=AsyncMock(), + delete_entity=AsyncMock(), + ) + repo._semantic_vector_index = adapter + repo._semantic_vector_index_name = "milvus" + + await repo.delete_external_entity_vectors( + [41, 42], + vector_index_names=frozenset({"milvus"}), + ) + + adapter.initialize.assert_awaited_once() + assert adapter.delete_entity.await_args_list == [((41,), {}), ((42,), {})] + + +@pytest.mark.asyncio +async def test_external_entity_cleanup_rejects_mismatched_adapter() -> None: + """A configured adapter must not delete rows owned by another extension.""" + repo = _ConcreteRepo() + repo._semantic_vector_index_name = "milvus" + + with pytest.raises(SemanticVectorIndexExtensionError, match="owned by"): + await repo.delete_external_entity_vectors( + [41], + vector_index_names=frozenset({"pinecone"}), + ) + + @pytest.mark.asyncio async def test_sync_entity_vectors_batch_flushes_at_configured_threshold(monkeypatch): """Batch sync should flush queued jobs at semantic_embedding_sync_batch_size boundaries.""" diff --git a/tests/repository/test_semantic_vector_index.py b/tests/repository/test_semantic_vector_index.py index 1c30bd125..5a32c26b7 100644 --- a/tests/repository/test_semantic_vector_index.py +++ b/tests/repository/test_semantic_vector_index.py @@ -294,3 +294,22 @@ def test_search_repository_composition_root_injects_selected_adapter(monkeypatch assert isinstance(repository, PostgresSearchRepository) assert repository._semantic_vector_index_name == "milvus" assert repository._semantic_vector_index is index + + +def test_disabled_search_repository_retains_configured_adapter_name(monkeypatch) -> None: + """Cleanup must identify external ownership without loading an embedding model.""" + monkeypatch.setattr( + "basic_memory.repository.search_repository.create_embedding_provider", + lambda _config: pytest.fail("disabled search must not create an embedding provider"), + ) + + repository = create_search_repository( + MagicMock(), + project_id=7, + app_config=_postgres_config(semantic_search_enabled=False), + database_backend=DatabaseBackend.POSTGRES, + ) + + assert isinstance(repository, PostgresSearchRepository) + assert repository._semantic_vector_index_name == "milvus" + assert not hasattr(repository, "_semantic_vector_index") diff --git a/tests/services/test_project_service_embedding_status.py b/tests/services/test_project_service_embedding_status.py index 3b197ae94..42ad36831 100644 --- a/tests/services/test_project_service_embedding_status.py +++ b/tests/services/test_project_service_embedding_status.py @@ -130,6 +130,48 @@ async def test_embedding_status_vector_tables_missing( assert "Vector manifest not initialized" in (status.reindex_reason or "") +@pytest.mark.asyncio +async def test_embedding_status_treats_legacy_sqlite_manifest_as_unavailable( + project_service: ProjectService, + test_graph, + test_project, +): + """Legacy SQLite manifests should recommend rebuild instead of querying new columns.""" + if _is_postgres(): + pytest.skip("The legacy in-place manifest schema only applies to SQLite.") + + await _execute(project_service, text("DROP TABLE search_vector_chunks"), {}) + await _execute( + project_service, + text( + "CREATE TABLE search_vector_chunks (" + "id INTEGER PRIMARY KEY AUTOINCREMENT, " + "entity_id INTEGER NOT NULL, " + "project_id INTEGER NOT NULL, " + "chunk_key TEXT NOT NULL, " + "chunk_text TEXT NOT NULL, " + "source_hash TEXT NOT NULL, " + "entity_fingerprint TEXT NOT NULL, " + "embedding_model TEXT NOT NULL, " + "updated_at TEXT NOT NULL DEFAULT CURRENT_TIMESTAMP)" + ), + {}, + ) + + with patch.object( + type(project_service), + "config_manager", + new_callable=lambda: property( + lambda self: _config_manager_with(semantic_search_enabled=True) + ), + ): + status = await project_service.get_embedding_status(test_project.id) + + assert status.vector_tables_exist is False + assert status.reindex_recommended is True + assert "schema is outdated" in (status.reindex_reason or "") + + @pytest.mark.asyncio async def test_embedding_status_entities_without_chunks( project_service: ProjectService, test_graph, test_project From d3cf2133d172a3867089b99f4c46229356320765 Mon Sep 17 00:00:00 2001 From: phernandez Date: Wed, 22 Jul 2026 10:40:22 -0500 Subject: [PATCH 10/24] fix(core): preserve vector ownership on cleanup Signed-off-by: phernandez --- .../repository/search_repository_base.py | 84 +++++++++++------ tests/repository/test_semantic_search_base.py | 89 ++++++++++++++----- 2 files changed, 126 insertions(+), 47 deletions(-) diff --git a/src/basic_memory/repository/search_repository_base.py b/src/basic_memory/repository/search_repository_base.py index 7c53854b1..b41f12e72 100644 --- a/src/basic_memory/repository/search_repository_base.py +++ b/src/basic_memory/repository/search_repository_base.py @@ -59,6 +59,7 @@ SMALL_NOTE_CONTENT_LIMIT = 2000 OVERSIZED_ENTITY_VECTOR_SHARD_SIZE = semantic_vector_sync.OVERSIZED_ENTITY_VECTOR_SHARD_SIZE _SQLITE_MAX_PREPARE_WINDOW = semantic_vector_sync.SQLITE_MAX_PREPARE_WINDOW +_BUILT_IN_VECTOR_INDEX_NAMES = frozenset({"pgvector", "sqlite-vec"}) # Entity, observation, and relation rows in search_index carry ids from independent # auto-increment sequences, so a bare id is ambiguous across row types. Every map in @@ -644,10 +645,12 @@ async def delete_project_vector_rows(self, *, strict_adapter_cleanup: bool = Fal A full reindex clears the manifest even when semantic search is disabled, so stale ready rows cannot become current if the feature is re-enabled. Adapter cleanup is best-effort because the manifest remains the search - authority and an unavailable extension must not block an FTS rebuild. + authority. Ownership mismatches fail closed because only the manifest + identifies the extension that can safely remove previously written vectors. Project deletion opts into strict cleanup so external data ownership is preserved for a retry instead of being discarded after an adapter failure. """ + configured_index = self._semantic_vector_index_name async with db.scoped_session(self.session_maker) as session: connection = await session.connection() manifest_exists = await connection.run_sync( @@ -656,21 +659,59 @@ async def delete_project_vector_rows(self, *, strict_adapter_cleanup: bool = Fal if not manifest_exists: return - manifest_has_embedding_status = await connection.run_sync( - lambda sync_connection: any( - column["name"] == "embedding_status" + manifest_columns = await connection.run_sync( + lambda sync_connection: { + str(column["name"]) for column in inspect(sync_connection).get_columns("search_vector_chunks") - ) + } ) + manifest_has_embedding_status = "embedding_status" in manifest_columns - result = await session.execute( - text( - "SELECT DISTINCT entity_id FROM search_vector_chunks " - "WHERE project_id = :project_id ORDER BY entity_id" - ), - {"project_id": self.project_id}, + entity_ids_by_vector_index: dict[str, list[int]] = {} + if "vector_index" in manifest_columns: + result = await session.execute( + text( + "SELECT DISTINCT entity_id, vector_index FROM search_vector_chunks " + "WHERE project_id = :project_id ORDER BY vector_index, entity_id" + ), + {"project_id": self.project_id}, + ) + for entity_id, vector_index in result.all(): + entity_ids_by_vector_index.setdefault(str(vector_index), []).append( + int(entity_id) + ) + else: + result = await session.execute( + text( + "SELECT DISTINCT entity_id FROM search_vector_chunks " + "WHERE project_id = :project_id ORDER BY entity_id" + ), + {"project_id": self.project_id}, + ) + legacy_vector_index = ( + "sqlite-vec" if connection.dialect.name == "sqlite" else "pgvector" + ) + entity_ids_by_vector_index[legacy_vector_index] = [ + int(entity_id) for entity_id in result.scalars().all() + ] + + external_vector_indexes = ( + frozenset(entity_ids_by_vector_index) - _BUILT_IN_VECTOR_INDEX_NAMES ) - entity_ids = [int(entity_id) for entity_id in result.scalars().all()] + + # Trigger: manifests belong to an external index other than the available adapter. + # Why: adapter configuration can change after vectors were written, and deleting + # the manifest would discard the only durable routing information for old vectors. + # Outcome: fail before touching any adapter or manifest so the owner can be restored. + if external_vector_indexes and ( + not hasattr(self, "_semantic_vector_index") + or external_vector_indexes != frozenset({configured_index}) + ): + raise SemanticVectorIndexExtensionError( + "Cannot delete project vectors owned by external indexes " + f"{sorted(external_vector_indexes)!r} with configured adapter " + f"{configured_index!r}. Restore the owning adapter and retry." + ) # Trigger: the manifest predates embedding lifecycle state. # Why: legacy SQLite schemas must reach cleanup before lazy schema repair runs. @@ -686,22 +727,11 @@ async def delete_project_vector_rows(self, *, strict_adapter_cleanup: bool = Fal ) await session.commit() - # Trigger: project deletion requires strict cleanup, but semantic search - # was disabled before this repository was composed. - # Why: deleting the manifest without an adapter would orphan externally - # stored vectors and discard the only durable ownership list. - # Outcome: retain the manifest and require the caller to restore the - # configured adapter before retrying project deletion. - if entity_ids and strict_adapter_cleanup and not hasattr(self, "_semantic_vector_index"): - raise SemanticVectorIndexExtensionError( - "Cannot delete project vectors because the configured semantic vector " - "adapter is unavailable. Enable semantic search and retry project deletion." - ) - - if hasattr(self, "_semantic_vector_index"): + adapter_entity_ids = entity_ids_by_vector_index.get(configured_index, []) + if adapter_entity_ids and hasattr(self, "_semantic_vector_index"): try: await self._semantic_vector_index.initialize() - for entity_id in entity_ids: + for entity_id in adapter_entity_ids: await self._semantic_vector_index.delete_entity(entity_id) except Exception as exc: # Trigger: a configured external adapter cannot initialize or delete. @@ -716,7 +746,7 @@ async def delete_project_vector_rows(self, *, strict_adapter_cleanup: bool = Fal vector_index=self._semantic_vector_index_name, error=exc, ) - if strict_adapter_cleanup: + if strict_adapter_cleanup and configured_index not in _BUILT_IN_VECTOR_INDEX_NAMES: raise async with db.scoped_session(self.session_maker) as session: diff --git a/tests/repository/test_semantic_search_base.py b/tests/repository/test_semantic_search_base.py index 61bdbaca3..4dd983b4b 100644 --- a/tests/repository/test_semantic_search_base.py +++ b/tests/repository/test_semantic_search_base.py @@ -169,9 +169,9 @@ async def test_project_vector_cleanup_clears_disabled_manifest(monkeypatch): repo = _ConcreteRepo() session = AsyncMock() connection = AsyncMock() - connection.run_sync.return_value = True + connection.run_sync.side_effect = [True, {"embedding_status", "vector_index"}] session.connection.return_value = connection - entity_result = SimpleNamespace(scalars=lambda: SimpleNamespace(all=lambda: [41, 42])) + entity_result = SimpleNamespace(all=lambda: [(41, "pgvector"), (42, "pgvector")]) session.execute.return_value = entity_result @asynccontextmanager @@ -193,7 +193,8 @@ async def test_project_vector_cleanup_handles_legacy_manifest_without_status(mon repo = _ConcreteRepo() session = AsyncMock() connection = AsyncMock() - connection.run_sync.side_effect = [True, False] + connection.run_sync.side_effect = [True, set()] + connection.dialect.name = "sqlite" session.connection.return_value = connection session.execute.return_value = SimpleNamespace( scalars=lambda: SimpleNamespace(all=lambda: [41]) @@ -224,11 +225,9 @@ async def test_project_vector_cleanup_uses_available_adapter(monkeypatch): repo._semantic_vector_index_name = "milvus" session = AsyncMock() connection = AsyncMock() - connection.run_sync.return_value = True + connection.run_sync.side_effect = [True, {"embedding_status", "vector_index"}] session.connection.return_value = connection - session.execute.return_value = SimpleNamespace( - scalars=lambda: SimpleNamespace(all=lambda: [41, 42]) - ) + session.execute.return_value = SimpleNamespace(all=lambda: [(41, "milvus"), (42, "milvus")]) @asynccontextmanager async def fake_scoped_session(_session_maker): @@ -257,11 +256,9 @@ async def test_project_vector_cleanup_clears_manifest_after_adapter_failure(monk repo._semantic_vector_index_name = "milvus" session = AsyncMock() connection = AsyncMock() - connection.run_sync.return_value = True + connection.run_sync.side_effect = [True, {"embedding_status", "vector_index"}] session.connection.return_value = connection - session.execute.return_value = SimpleNamespace( - scalars=lambda: SimpleNamespace(all=lambda: [41]) - ) + session.execute.return_value = SimpleNamespace(all=lambda: [(41, "milvus")]) @asynccontextmanager async def fake_scoped_session(_session_maker): @@ -291,11 +288,9 @@ async def test_strict_project_vector_cleanup_preserves_manifest_after_adapter_fa repo._semantic_vector_index_name = "milvus" session = AsyncMock() connection = AsyncMock() - connection.run_sync.return_value = True + connection.run_sync.side_effect = [True, {"embedding_status", "vector_index"}] session.connection.return_value = connection - session.execute.return_value = SimpleNamespace( - scalars=lambda: SimpleNamespace(all=lambda: [41]) - ) + session.execute.return_value = SimpleNamespace(all=lambda: [(41, "milvus")]) @asynccontextmanager async def fake_scoped_session(_session_maker): @@ -318,11 +313,9 @@ async def test_strict_project_vector_cleanup_preserves_manifest_without_adapter( repo = _ConcreteRepo() session = AsyncMock() connection = AsyncMock() - connection.run_sync.return_value = True + connection.run_sync.side_effect = [True, {"embedding_status", "vector_index"}] session.connection.return_value = connection - session.execute.return_value = SimpleNamespace( - scalars=lambda: SimpleNamespace(all=lambda: [41]) - ) + session.execute.return_value = SimpleNamespace(all=lambda: [(41, "milvus")]) @asynccontextmanager async def fake_scoped_session(_session_maker): @@ -330,7 +323,7 @@ async def fake_scoped_session(_session_maker): monkeypatch.setattr(search_repository_base_module.db, "scoped_session", fake_scoped_session) - with pytest.raises(SemanticVectorIndexExtensionError, match="adapter is unavailable"): + with pytest.raises(SemanticVectorIndexExtensionError, match="owned by external indexes"): await repo.delete_project_vector_rows(strict_adapter_cleanup=True) statements = [str(call.args[0]) for call in session.execute.await_args_list] @@ -339,6 +332,62 @@ async def fake_scoped_session(_session_maker): ) +@pytest.mark.asyncio +async def test_project_vector_cleanup_preserves_mismatched_external_owner(monkeypatch): + """Cleanup must not route old extension rows through the newly configured adapter.""" + repo = _ConcreteRepo() + adapter: Any = SimpleNamespace( + initialize=AsyncMock(), + delete_entity=AsyncMock(), + ) + repo._semantic_vector_index = adapter + repo._semantic_vector_index_name = "pgvector" + session = AsyncMock() + connection = AsyncMock() + connection.run_sync.side_effect = [True, {"embedding_status", "vector_index"}] + session.connection.return_value = connection + session.execute.return_value = SimpleNamespace(all=lambda: [(41, "milvus")]) + + @asynccontextmanager + async def fake_scoped_session(_session_maker): + yield session + + monkeypatch.setattr(search_repository_base_module.db, "scoped_session", fake_scoped_session) + + with pytest.raises(SemanticVectorIndexExtensionError, match="milvus"): + await repo.delete_project_vector_rows() + + adapter.initialize.assert_not_awaited() + statements = [str(call.args[0]) for call in session.execute.await_args_list] + assert not any(statement.startswith("UPDATE search_vector_chunks") for statement in statements) + assert not any( + statement.startswith("DELETE FROM search_vector_chunks") for statement in statements + ) + + +@pytest.mark.asyncio +async def test_strict_project_vector_cleanup_allows_disabled_builtin_owner(monkeypatch): + """Project deletion should not require the semantic stack for built-in vectors.""" + repo = _ConcreteRepo() + repo._semantic_vector_index_name = "pgvector" + session = AsyncMock() + connection = AsyncMock() + connection.run_sync.side_effect = [True, {"embedding_status", "vector_index"}] + session.connection.return_value = connection + session.execute.return_value = SimpleNamespace(all=lambda: [(41, "pgvector")]) + + @asynccontextmanager + async def fake_scoped_session(_session_maker): + yield session + + monkeypatch.setattr(search_repository_base_module.db, "scoped_session", fake_scoped_session) + + await repo.delete_project_vector_rows(strict_adapter_cleanup=True) + + statements = [str(call.args[0]) for call in session.execute.await_args_list] + assert any(statement.startswith("DELETE FROM search_vector_chunks") for statement in statements) + + @pytest.mark.asyncio async def test_external_entity_cleanup_uses_matching_project_adapter() -> None: """DB-first deletes should invoke the configured extension before manifest removal.""" From 4c5e3a0c0be05a354a40882dfc94dd65ad772d02 Mon Sep 17 00:00:00 2001 From: phernandez Date: Wed, 22 Jul 2026 10:55:28 -0500 Subject: [PATCH 11/24] fix(core): guard entity vector ownership Signed-off-by: phernandez --- .../repository/search_repository_base.py | 61 ++++++++++++++----- .../repository/semantic_vector_sync.py | 5 ++ tests/repository/test_semantic_search_base.py | 18 ++++++ .../test_sqlite_vector_search_repository.py | 12 ++-- 4 files changed, 76 insertions(+), 20 deletions(-) diff --git a/src/basic_memory/repository/search_repository_base.py b/src/basic_memory/repository/search_repository_base.py index b41f12e72..2f001e0a6 100644 --- a/src/basic_memory/repository/search_repository_base.py +++ b/src/basic_memory/repository/search_repository_base.py @@ -360,6 +360,14 @@ async def _delete_entity_chunks( entity_id: int, ) -> None: """Stage an entity deletion by making its manifest rows non-searchable.""" + vector_index_result = await session.execute( + text( + "SELECT DISTINCT vector_index FROM search_vector_chunks " + "WHERE project_id = :project_id AND entity_id = :entity_id" + ), + {"project_id": self.project_id, "entity_id": entity_id}, + ) + self._assert_manifest_vector_ownership(vector_index_result.scalars().all()) await session.execute( text( "UPDATE search_vector_chunks SET embedding_status = 'pending' " @@ -379,6 +387,15 @@ async def _delete_stale_chunks( return params = {f"stale_id_{index}": row_id for index, row_id in enumerate(stale_ids)} placeholders = ", ".join(f":stale_id_{index}" for index in range(len(stale_ids))) + vector_index_result = await session.execute( + text( + "SELECT DISTINCT vector_index FROM search_vector_chunks " + f"WHERE project_id = :project_id AND entity_id = :entity_id " + f"AND id IN ({placeholders})" + ), + {**params, "project_id": self.project_id, "entity_id": entity_id}, + ) + self._assert_manifest_vector_ownership(vector_index_result.scalars().all()) await session.execute( text( "UPDATE search_vector_chunks SET embedding_status = 'pending' " @@ -404,6 +421,15 @@ async def _finalize_prepared_vector_deletions( return if prepared.delete_entity_vectors: + async with db.scoped_session(self.session_maker) as session: + vector_index_result = await session.execute( + text( + "SELECT DISTINCT vector_index FROM search_vector_chunks " + "WHERE project_id = :project_id AND entity_id = :entity_id" + ), + {"project_id": self.project_id, "entity_id": prepared.entity_id}, + ) + self._assert_manifest_vector_ownership(vector_index_result.scalars().all()) await self._semantic_vector_index.delete_entity(prepared.entity_id) async with db.scoped_session(self.session_maker) as session: await session.execute( @@ -422,15 +448,17 @@ async def _finalize_prepared_vector_deletions( async with db.scoped_session(self.session_maker) as session: result = await session.execute( text( - "SELECT entity_id, chunk_key FROM search_vector_chunks " + "SELECT entity_id, chunk_key, vector_index FROM search_vector_chunks " f"WHERE project_id = :project_id AND entity_id = :entity_id " f"AND id IN ({placeholders})" ), {**params, "project_id": self.project_id, "entity_id": prepared.entity_id}, ) + rows = result.mappings().all() + self._assert_manifest_vector_ownership(row["vector_index"] for row in rows) keys = [ VectorKey(entity_id=int(row["entity_id"]), chunk_key=str(row["chunk_key"])) - for row in result.mappings().all() + for row in rows ] await self._semantic_vector_index.delete(keys) @@ -637,6 +665,21 @@ async def delete_external_entity_vectors( for entity_id in entity_ids: await self._semantic_vector_index.delete_entity(entity_id) + def _assert_manifest_vector_ownership(self, vector_index_names: Iterable[object]) -> None: + """Reject cleanup that cannot reach every externally owned vector.""" + recorded_indexes = frozenset(str(name) for name in vector_index_names if str(name)) + external_indexes = recorded_indexes - _BUILT_IN_VECTOR_INDEX_NAMES + configured_index = self._semantic_vector_index_name + if external_indexes and ( + not hasattr(self, "_semantic_vector_index") + or external_indexes != frozenset({configured_index}) + ): + raise SemanticVectorIndexExtensionError( + "Cannot mutate vector manifests owned by external indexes " + f"{sorted(external_indexes)!r} with configured adapter " + f"{configured_index!r}. Restore the owning adapter and retry." + ) + async def delete_project_vector_rows(self, *, strict_adapter_cleanup: bool = False) -> None: """Delete this project's vectors through the configured storage adapter. @@ -695,23 +738,11 @@ async def delete_project_vector_rows(self, *, strict_adapter_cleanup: bool = Fal int(entity_id) for entity_id in result.scalars().all() ] - external_vector_indexes = ( - frozenset(entity_ids_by_vector_index) - _BUILT_IN_VECTOR_INDEX_NAMES - ) - # Trigger: manifests belong to an external index other than the available adapter. # Why: adapter configuration can change after vectors were written, and deleting # the manifest would discard the only durable routing information for old vectors. # Outcome: fail before touching any adapter or manifest so the owner can be restored. - if external_vector_indexes and ( - not hasattr(self, "_semantic_vector_index") - or external_vector_indexes != frozenset({configured_index}) - ): - raise SemanticVectorIndexExtensionError( - "Cannot delete project vectors owned by external indexes " - f"{sorted(external_vector_indexes)!r} with configured adapter " - f"{configured_index!r}. Restore the owning adapter and retry." - ) + self._assert_manifest_vector_ownership(entity_ids_by_vector_index) # Trigger: the manifest predates embedding lifecycle state. # Why: legacy SQLite schemas must reach cleanup before lazy schema repair runs. diff --git a/src/basic_memory/repository/semantic_vector_sync.py b/src/basic_memory/repository/semantic_vector_sync.py index 7c5d3a871..13f927054 100644 --- a/src/basic_memory/repository/semantic_vector_sync.py +++ b/src/basic_memory/repository/semantic_vector_sync.py @@ -985,6 +985,11 @@ async def upsert_scheduled_chunk_records( embedding_model: str, ) -> list[tuple[int, str]]: """Upsert scheduled chunk rows and return embedding jobs.""" + repository._assert_manifest_vector_ownership( + current.vector_index + for record in scheduled_records + if (current := existing_by_key.get(record["chunk_key"])) is not None + ) timestamp_expr = repository._timestamp_now_expr() embedding_jobs: list[tuple[int, str]] = [] for record in scheduled_records: diff --git a/tests/repository/test_semantic_search_base.py b/tests/repository/test_semantic_search_base.py index 4dd983b4b..61f5a9c7e 100644 --- a/tests/repository/test_semantic_search_base.py +++ b/tests/repository/test_semantic_search_base.py @@ -365,6 +365,24 @@ async def fake_scoped_session(_session_maker): ) +@pytest.mark.asyncio +async def test_entity_vector_cleanup_preserves_mismatched_external_owner() -> None: + """Entity cleanup must fail before staging rows owned by another extension.""" + repo = _ConcreteRepo() + adapter: Any = SimpleNamespace() + repo._semantic_vector_index = adapter + repo._semantic_vector_index_name = "pgvector" + session = AsyncMock() + session.execute.return_value = SimpleNamespace( + scalars=lambda: SimpleNamespace(all=lambda: ["milvus"]) + ) + + with pytest.raises(SemanticVectorIndexExtensionError, match="milvus"): + await SearchRepositoryBase._delete_entity_chunks(repo, session, 41) + + assert session.execute.await_count == 1 + + @pytest.mark.asyncio async def test_strict_project_vector_cleanup_allows_disabled_builtin_owner(monkeypatch): """Project deletion should not require the semantic stack for built-in vectors.""" diff --git a/tests/repository/test_sqlite_vector_search_repository.py b/tests/repository/test_sqlite_vector_search_repository.py index e92605e3a..80c60aa2e 100644 --- a/tests/repository/test_sqlite_vector_search_repository.py +++ b/tests/repository/test_sqlite_vector_search_repository.py @@ -18,6 +18,7 @@ from basic_memory.repository.prefixing_provider import PrefixingEmbeddingProvider from basic_memory.repository import search_repository_base as search_repository_base_module from basic_memory.repository.search_index_row import SearchIndexRow +from basic_memory.repository.semantic_errors import SemanticVectorIndexExtensionError from basic_memory.repository.semantic_vector_index import ( VectorIndexScope, VectorKey, @@ -695,8 +696,8 @@ async def test_adapter_matches_hydrate_only_current_ready_manifest_rows(search_r @pytest.mark.asyncio -async def test_vector_index_switch_reembeds_manifest_rows(search_repository): - """Changing the selected adapter identity invalidates otherwise unchanged chunks.""" +async def test_external_vector_index_switch_preserves_manifest_ownership(search_repository): + """Changing an external adapter identity must retain the old cleanup route.""" if not isinstance(search_repository, SQLiteSearchRepository): pytest.skip("Semantic manifest behavior is exercised through local SQLite.") @@ -718,7 +719,8 @@ async def test_vector_index_switch_reembeds_manifest_rows(search_repository): await search_repository.sync_entity_vectors(114) search_repository._semantic_vector_index_name = "recording-b" - await search_repository.sync_entity_vectors(114) + with pytest.raises(SemanticVectorIndexExtensionError, match="recording-a"): + await search_repository.sync_entity_vectors(114) async with db.scoped_session(search_repository.session_maker) as session: state = await session.execute( @@ -728,9 +730,9 @@ async def test_vector_index_switch_reembeds_manifest_rows(search_repository): ), {"project_id": search_repository.project_id, "entity_id": 114}, ) - assert state.all() == [("recording-b", "ready")] + assert state.all() == [("recording-a", "ready")] - assert len(adapter.upsert_calls) == 2 + assert len(adapter.upsert_calls) == 1 @pytest.mark.asyncio From 67f27ddc9e116125e81fbf6669796f4e2410a5ab Mon Sep 17 00:00:00 2001 From: phernandez Date: Wed, 22 Jul 2026 11:05:40 -0500 Subject: [PATCH 12/24] fix(core): guard postgres vector ownership Signed-off-by: phernandez --- .../repository/postgres_search_repository.py | 5 ++ .../test_postgres_search_repository_unit.py | 60 +++++++++++++++++-- 2 files changed, 60 insertions(+), 5 deletions(-) diff --git a/src/basic_memory/repository/postgres_search_repository.py b/src/basic_memory/repository/postgres_search_repository.py index d370301ce..5962ad9cd 100644 --- a/src/basic_memory/repository/postgres_search_repository.py +++ b/src/basic_memory/repository/postgres_search_repository.py @@ -399,6 +399,11 @@ async def _upsert_scheduled_chunk_records( if not scheduled_records: return [] + self._assert_manifest_vector_ownership( + current.vector_index + for record in scheduled_records + if (current := existing_by_key.get(record["chunk_key"])) is not None + ) upsert_params: dict[str, object] = { "project_id": self.project_id, "entity_id": entity_id, diff --git a/tests/repository/test_postgres_search_repository_unit.py b/tests/repository/test_postgres_search_repository_unit.py index cbc12bb5f..8f4e79f76 100644 --- a/tests/repository/test_postgres_search_repository_unit.py +++ b/tests/repository/test_postgres_search_repository_unit.py @@ -14,10 +14,14 @@ from basic_memory.config import BasicMemoryConfig, DatabaseBackend from basic_memory.repository.pgvector_index import PgVectorIndex from basic_memory.repository.postgres_search_repository import PostgresSearchRepository -from basic_memory.repository.search_repository_base import _PreparedEntityVectorSync +from basic_memory.repository.search_repository_base import ( + VectorChunkState, + _PreparedEntityVectorSync, +) from basic_memory.repository.semantic_errors import ( SemanticDependenciesMissingError, SemanticSearchDisabledError, + SemanticVectorIndexExtensionError, ) @@ -242,11 +246,14 @@ class TestDeleteStaleChunks: async def test_delete_stale_chunks_builds_correct_params(self): repo = _make_repo() session = AsyncMock() + ownership_result = MagicMock() + ownership_result.scalars.return_value.all.return_value = [] + session.execute.side_effect = [ownership_result, MagicMock()] stale_ids = [10, 20, 30] await repo._delete_stale_chunks(session, stale_ids, entity_id=5) - session.execute.assert_called_once() - call_args = session.execute.call_args + assert session.execute.await_count == 2 + call_args = session.execute.await_args_list[1] params = call_args[0][1] assert params["stale_id_0"] == 10 assert params["stale_id_1"] == 20 @@ -265,14 +272,57 @@ class TestDeleteEntityChunks: async def test_delete_entity_chunks_executes_sql(self): repo = _make_repo() session = AsyncMock() + ownership_result = MagicMock() + ownership_result.scalars.return_value.all.return_value = [] + session.execute.side_effect = [ownership_result, MagicMock()] await repo._delete_entity_chunks(session, entity_id=42) - session.execute.assert_called_once() - call_args = session.execute.call_args + assert session.execute.await_count == 2 + call_args = session.execute.await_args_list[1] params = call_args[0][1] assert params["project_id"] == repo.project_id assert params["entity_id"] == 42 +@pytest.mark.asyncio +async def test_postgres_upsert_preserves_external_vector_ownership() -> None: + """Postgres must reject an adapter switch before overwriting manifest ownership.""" + repo = _make_repo( + semantic_enabled=True, + embedding_provider=StubEmbeddingProvider(), + ) + repo._semantic_vector_index_name = "recording-b" + session = AsyncMock() + + with pytest.raises(SemanticVectorIndexExtensionError, match="recording-a"): + await repo._upsert_scheduled_chunk_records( + session, + entity_id=42, + scheduled_records=[ + { + "chunk_key": "entity:42:0", + "chunk_text": "changed", + "source_hash": "new-hash", + } + ], + existing_by_key={ + "entity:42:0": VectorChunkState( + id=7, + chunk_key="entity:42:0", + source_hash="old-hash", + entity_fingerprint="old-fingerprint", + embedding_model="stub:4:document", + has_embedding=True, + vector_index="recording-a", + embedding_status="ready", + ) + }, + entity_fingerprint="new-fingerprint", + embedding_model="stub:4:document", + ) + + session.execute.assert_not_awaited() + + class TestBatchPrepareWindow: """Cover the shared batched prepare window used by Postgres.""" From d9aa17f76824013cbe47ac46816483cc0a251f3b Mon Sep 17 00:00:00 2001 From: phernandez Date: Wed, 22 Jul 2026 11:21:24 -0500 Subject: [PATCH 13/24] fix(core): close vector cleanup lifecycle gaps Signed-off-by: phernandez --- src/basic_memory/db.py | 31 ++++++----- .../repository/search_repository_base.py | 36 ++++++++++++- .../repository/sqlite_search_repository.py | 23 +++++++++ tests/db/test_semantic_search_startup.py | 51 +++++++++++++++++++ tests/repository/test_semantic_search_base.py | 23 +++++++-- .../test_sqlite_vector_search_repository.py | 51 +++++++++++++++++++ 6 files changed, 196 insertions(+), 19 deletions(-) create mode 100644 tests/db/test_semantic_search_startup.py diff --git a/src/basic_memory/db.py b/src/basic_memory/db.py index b73da2f93..568a3b524 100644 --- a/src/basic_memory/db.py +++ b/src/basic_memory/db.py @@ -21,9 +21,6 @@ ) from sqlalchemy.pool import AsyncAdaptedQueuePool, NullPool -from basic_memory.repository.postgres_search_repository import PostgresSearchRepository -from basic_memory.repository.sqlite_search_repository import SQLiteSearchRepository - # ----------------------------------------------------------------------------- # Windows event loop policy # ----------------------------------------------------------------------------- @@ -562,17 +559,23 @@ async def run_migrations( else: session_maker = _session_maker - # Initialize the search index schema - # For SQLite: Create FTS5 virtual table - # For Postgres: No-op (tsvector column added by migrations) - # The project_id is not used for init_search_index, so we pass a dummy value - if ( - database_type == DatabaseType.POSTGRES - or app_config.database_backend == DatabaseBackend.POSTGRES - ): - await PostgresSearchRepository(session_maker, 1).init_search_index() - else: - await SQLiteSearchRepository(session_maker, 1).init_search_index() + # Import lazily because backend repositories import this module for session + # management. Startup must still use the composition factory so configured + # semantic-vector extensions are available during index initialization. + from basic_memory.repository.search_repository import create_search_repository + + database_backend = ( + DatabaseBackend.POSTGRES + if database_type == DatabaseType.POSTGRES + else app_config.database_backend + ) + search_repository = create_search_repository( + session_maker=session_maker, + project_id=1, + app_config=app_config, + database_backend=database_backend, + ) + await search_repository.init_search_index() except Exception as e: # pragma: no cover logger.error(f"Error running migrations: {e}") diff --git a/src/basic_memory/repository/search_repository_base.py b/src/basic_memory/repository/search_repository_base.py index 2f001e0a6..5d54da95e 100644 --- a/src/basic_memory/repository/search_repository_base.py +++ b/src/basic_memory/repository/search_repository_base.py @@ -661,10 +661,43 @@ async def delete_external_entity_vectors( "Enable semantic search and retry the entity deletion." ) + deleted_entity_ids = tuple(dict.fromkeys(entity_ids)) + if not deleted_entity_ids: + return + + # Trigger: DB-first deletion runs inside a caller-owned SQL transaction. + # Why: the external adapter cannot participate in that transaction. If its + # delete succeeds and the caller later rolls back, ready manifests would + # incorrectly claim the now-missing vectors are searchable. + # Outcome: commit a non-searchable retry marker in an independent session + # before touching external storage; a retried delete remains idempotent. + params = { + "project_id": self.project_id, + "vector_index": configured_index, + **{ + f"entity_id_{index}": entity_id + for index, entity_id in enumerate(deleted_entity_ids) + }, + } + placeholders = ", ".join(f":entity_id_{index}" for index in range(len(deleted_entity_ids))) + async with db.scoped_session(self.session_maker) as session: + await session.execute( + text( + "UPDATE search_vector_chunks SET embedding_status = 'pending' " + "WHERE project_id = :project_id AND vector_index = :vector_index " + f"AND entity_id IN ({placeholders})" + ), + params, + ) + await session.commit() + await self._semantic_vector_index.initialize() - for entity_id in entity_ids: + for entity_id in deleted_entity_ids: await self._semantic_vector_index.delete_entity(entity_id) + async def _delete_project_builtin_vector_rows(self, session: AsyncSession) -> None: + """Delete backend-owned vector rows before their SQL manifest is removed.""" + def _assert_manifest_vector_ownership(self, vector_index_names: Iterable[object]) -> None: """Reject cleanup that cannot reach every externally owned vector.""" recorded_indexes = frozenset(str(name) for name in vector_index_names if str(name)) @@ -781,6 +814,7 @@ async def delete_project_vector_rows(self, *, strict_adapter_cleanup: bool = Fal raise async with db.scoped_session(self.session_maker) as session: + await self._delete_project_builtin_vector_rows(session) await session.execute( text("DELETE FROM search_vector_chunks WHERE project_id = :project_id"), {"project_id": self.project_id}, diff --git a/src/basic_memory/repository/sqlite_search_repository.py b/src/basic_memory/repository/sqlite_search_repository.py index bb97b293c..91aa8eba5 100644 --- a/src/basic_memory/repository/sqlite_search_repository.py +++ b/src/basic_memory/repository/sqlite_search_repository.py @@ -547,6 +547,29 @@ async def _delete_stale_chunks( ) -> None: await super()._delete_stale_chunks(session, stale_ids, entity_id) + async def _delete_project_builtin_vector_rows(self, session: AsyncSession) -> None: + """Delete sqlite-vec rows atomically with their project manifest.""" + table_result = await session.execute( + text( + "SELECT name, sql FROM sqlite_master WHERE type = 'table' " + "AND name IN ('search_vector_chunks', 'search_vector_embeddings')" + ) + ) + table_sql = {str(name): str(sql or "") for name, sql in table_result.all()} + if not {"search_vector_chunks", "search_vector_embeddings"}.issubset(table_sql): + return + + vector_sql = table_sql["search_vector_embeddings"] + if "using vec0" in vector_sql.lower(): + await self._ensure_sqlite_vec_loaded(session) + await session.execute( + text( + "DELETE FROM search_vector_embeddings WHERE rowid IN (" + "SELECT id FROM search_vector_chunks WHERE project_id = :project_id)" + ), + {"project_id": self.project_id}, + ) + async def drop_vector_tables(self) -> None: """Drop SQLite vector tables on a sqlite-vec-enabled connection.""" async with db.scoped_session(self.session_maker) as session: diff --git a/tests/db/test_semantic_search_startup.py b/tests/db/test_semantic_search_startup.py new file mode 100644 index 000000000..37a81e2c9 --- /dev/null +++ b/tests/db/test_semantic_search_startup.py @@ -0,0 +1,51 @@ +"""Regression tests for semantic search repository startup composition.""" + +from types import SimpleNamespace +from unittest.mock import AsyncMock + +import pytest + +from basic_memory import db +from basic_memory.config import BasicMemoryConfig, DatabaseBackend +from basic_memory.repository import search_repository as search_repository_module + + +@pytest.mark.asyncio +async def test_migrations_initialize_search_through_repository_factory( + monkeypatch, tmp_path +) -> None: + """Startup must load configured vector extensions through the repository factory.""" + config = BasicMemoryConfig( + env="test", + database_backend=DatabaseBackend.POSTGRES, + database_url="postgresql+asyncpg://test:test@localhost/test", + semantic_search_enabled=True, + semantic_vector_index="milvus", + ) + session_maker = object() + repository = SimpleNamespace(init_search_index=AsyncMock()) + factory_calls: list[dict[str, object]] = [] + + def fake_create_search_repository(**kwargs): + factory_calls.append(kwargs) + return repository + + monkeypatch.setattr(db, "_session_maker", session_maker) + monkeypatch.setattr(db.command, "upgrade", lambda *_args: None) + monkeypatch.setattr( + search_repository_module, + "create_search_repository", + fake_create_search_repository, + ) + + await db.run_migrations(config, db.DatabaseType.POSTGRES) + + assert factory_calls == [ + { + "session_maker": session_maker, + "project_id": 1, + "app_config": config, + "database_backend": DatabaseBackend.POSTGRES, + } + ] + repository.init_search_index.assert_awaited_once() diff --git a/tests/repository/test_semantic_search_base.py b/tests/repository/test_semantic_search_base.py index 61f5a9c7e..aa1330cec 100644 --- a/tests/repository/test_semantic_search_base.py +++ b/tests/repository/test_semantic_search_base.py @@ -407,15 +407,25 @@ async def fake_scoped_session(_session_maker): @pytest.mark.asyncio -async def test_external_entity_cleanup_uses_matching_project_adapter() -> None: - """DB-first deletes should invoke the configured extension before manifest removal.""" +async def test_external_entity_cleanup_uses_matching_project_adapter(monkeypatch) -> None: + """DB-first deletes should durably stage manifests before invoking the extension.""" repo = _ConcreteRepo() + events: list[str] = [] adapter: Any = SimpleNamespace( - initialize=AsyncMock(), - delete_entity=AsyncMock(), + initialize=AsyncMock(side_effect=lambda: events.append("initialize")), + delete_entity=AsyncMock(side_effect=lambda _entity_id: events.append("delete")), ) repo._semantic_vector_index = adapter repo._semantic_vector_index_name = "milvus" + session = AsyncMock() + session.execute.side_effect = lambda *_args, **_kwargs: events.append("stage") + session.commit.side_effect = lambda: events.append("commit") + + @asynccontextmanager + async def fake_scoped_session(_session_maker): + yield session + + monkeypatch.setattr(search_repository_base_module.db, "scoped_session", fake_scoped_session) await repo.delete_external_entity_vectors( [41, 42], @@ -424,6 +434,11 @@ async def test_external_entity_cleanup_uses_matching_project_adapter() -> None: adapter.initialize.assert_awaited_once() assert adapter.delete_entity.await_args_list == [((41,), {}), ((42,), {})] + assert events == ["stage", "commit", "initialize", "delete", "delete"] + stage_statement = str(session.execute.await_args.args[0]) + assert stage_statement.startswith( + "UPDATE search_vector_chunks SET embedding_status = 'pending'" + ) @pytest.mark.asyncio diff --git a/tests/repository/test_sqlite_vector_search_repository.py b/tests/repository/test_sqlite_vector_search_repository.py index 80c60aa2e..48da2caee 100644 --- a/tests/repository/test_sqlite_vector_search_repository.py +++ b/tests/repository/test_sqlite_vector_search_repository.py @@ -290,6 +290,57 @@ async def test_sqlite_vec_recreated_storage_invalidates_ready_manifest(search_re assert result.scalar_one() == "pending" +@pytest.mark.asyncio +async def test_disabled_semantic_cleanup_deletes_sqlite_vec_rows(search_repository): + """Project cleanup must not strand sqlite-vec rows when semantic search is disabled.""" + if not isinstance(search_repository, SQLiteSearchRepository): + pytest.skip("sqlite-vec storage cleanup is local SQLite-only.") + + _enable_semantic(search_repository) + await search_repository.init_search_index() + embedding_identity = search_repository._embedding_model_key() + + async with db.scoped_session(search_repository.session_maker) as session: + index = cast(SQLiteVecIndex, search_repository._semantic_vector_index) + 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 (" + "906, 906, :project_id, 'entity:906:0', 'text', 'hash', " + "'fingerprint', :embedding_model, 'sqlite-vec', 'ready')" + ), + { + "project_id": search_repository.project_id, + "embedding_model": embedding_identity, + }, + ) + await session.execute( + text( + "INSERT INTO search_vector_embeddings (rowid, embedding) VALUES (906, :embedding)" + ), + {"embedding": "[1.0, 0.0, 0.0, 0.0]"}, + ) + await session.commit() + + search_repository._semantic_enabled = False + del search_repository._semantic_vector_index + await search_repository.delete_project_vector_rows() + + async with db.scoped_session(search_repository.session_maker) as session: + manifest_count = await session.scalar( + text("SELECT COUNT(*) FROM search_vector_chunks WHERE project_id = :project_id"), + {"project_id": search_repository.project_id}, + ) + vector_count = await session.scalar( + text("SELECT COUNT(*) FROM search_vector_embeddings WHERE rowid = 906") + ) + assert manifest_count == 0 + assert vector_count == 0 + + @pytest.mark.asyncio async def test_sqlite_vec_reconciliation_is_project_scoped(search_repository): """Reconciliation removes orphan/local stale rows without touching another project.""" From 7e7a456723fb9b7d40ee51f569d6c90ed8c45dd2 Mon Sep 17 00:00:00 2001 From: phernandez Date: Wed, 22 Jul 2026 11:34:55 -0500 Subject: [PATCH 14/24] fix(core): bound vector cleanup and hydration Signed-off-by: phernandez --- .../repository/search_repository_base.py | 69 +++++++++++-------- src/basic_memory/services/project_service.py | 12 +++- tests/repository/test_semantic_search_base.py | 37 ++++++++++ .../test_project_service_operations.py | 41 +++++++++++ 4 files changed, 126 insertions(+), 33 deletions(-) diff --git a/src/basic_memory/repository/search_repository_base.py b/src/basic_memory/repository/search_repository_base.py index 5d54da95e..a513e0a5f 100644 --- a/src/basic_memory/repository/search_repository_base.py +++ b/src/basic_memory/repository/search_repository_base.py @@ -53,6 +53,7 @@ # --- Semantic search constants --- VECTOR_FILTER_SCAN_LIMIT = 50000 +VECTOR_HYDRATION_BATCH_SIZE = 250 FUSION_BONUS = 0.3 FTS_GATE_THRESHOLD = 0.0 TOP_CHUNKS_PER_RESULT = 5 @@ -232,37 +233,45 @@ async def _hydrate_vector_matches( if not matches: return [] - params: dict[str, object] = { - "project_id": self.project_id, - "vector_index": self._semantic_vector_index_name, - "embedding_model": self._embedding_model_key(), - } - predicates: list[str] = [] - for index, match in enumerate(matches): - params[f"entity_id_{index}"] = match.key.entity_id - params[f"chunk_key_{index}"] = match.key.chunk_key - predicates.append( - f"(entity_id = :entity_id_{index} AND chunk_key = :chunk_key_{index})" - ) + chunks_by_key: dict[VectorKey, str] = {} + for batch_start in range(0, len(matches), VECTOR_HYDRATION_BATCH_SIZE): + batch = matches[batch_start : batch_start + VECTOR_HYDRATION_BATCH_SIZE] + params: dict[str, object] = { + "project_id": self.project_id, + "vector_index": self._semantic_vector_index_name, + "embedding_model": self._embedding_model_key(), + } + predicates: list[str] = [] + for index, match in enumerate(batch): + params[f"entity_id_{index}"] = match.key.entity_id + params[f"chunk_key_{index}"] = match.key.chunk_key + predicates.append( + f"(entity_id = :entity_id_{index} AND chunk_key = :chunk_key_{index})" + ) - result = await session.execute( - text( - "SELECT entity_id, chunk_key, chunk_text FROM search_vector_chunks " - "WHERE project_id = :project_id " - "AND vector_index = :vector_index " - "AND embedding_model = :embedding_model " - "AND embedding_status = 'ready' " - "AND (" + " OR ".join(predicates) + ")" - ), - params, - ) - chunks_by_key = { - VectorKey( - entity_id=int(row["entity_id"]), - chunk_key=str(row["chunk_key"]), - ): str(row["chunk_text"]) - for row in result.mappings().all() - } + # Constraint: adapters may return thousands of candidates for deep pages. + # PostgreSQL and SQLite both cap bind parameters, so hydrate in fixed-size + # batches while retaining the adapter's original ranking in the final list. + result = await session.execute( + text( + "SELECT entity_id, chunk_key, chunk_text FROM search_vector_chunks " + "WHERE project_id = :project_id " + "AND vector_index = :vector_index " + "AND embedding_model = :embedding_model " + "AND embedding_status = 'ready' " + "AND (" + " OR ".join(predicates) + ")" + ), + params, + ) + chunks_by_key.update( + { + VectorKey( + entity_id=int(row["entity_id"]), + chunk_key=str(row["chunk_key"]), + ): str(row["chunk_text"]) + for row in result.mappings().all() + } + ) return [ { "entity_id": match.key.entity_id, diff --git a/src/basic_memory/services/project_service.py b/src/basic_memory/services/project_service.py index fd3850070..de7631b80 100644 --- a/src/basic_memory/services/project_service.py +++ b/src/basic_memory/services/project_service.py @@ -17,7 +17,7 @@ from basic_memory import db from basic_memory.models import Project from basic_memory.repository.project_repository import ProjectRepository -from basic_memory.repository.search_repository import SearchRepository +from basic_memory.repository.search_repository import SearchRepository, create_search_repository from basic_memory.repository.embedding_provider_factory import ( configured_embedding_provider_identity, ) @@ -358,9 +358,15 @@ async def remove_project(self, name: str, delete_notes: bool = False) -> None: # Why: external storage has no database cascade and cannot reconcile after # the project manifest disappears. # Outcome: delete adapter-owned vectors while project ownership is still known. - if self._search_repository_factory is not None: + if self._search_repository_factory is None: + search_repository = create_search_repository( + session_maker=self.session_maker, + project_id=project_id, + app_config=self.config_manager.config, + ) + else: search_repository = self._search_repository_factory(project_id) - await search_repository.delete_project_vector_rows(strict_adapter_cleanup=True) + await search_repository.delete_project_vector_rows(strict_adapter_cleanup=True) async with db.scoped_session(self.session_maker) as session: # Remove from config if it exists there (may not exist in cloud mode) diff --git a/tests/repository/test_semantic_search_base.py b/tests/repository/test_semantic_search_base.py index aa1330cec..4b89f2e4c 100644 --- a/tests/repository/test_semantic_search_base.py +++ b/tests/repository/test_semantic_search_base.py @@ -21,6 +21,7 @@ SemanticSearchDisabledError, SemanticVectorIndexExtensionError, ) +from basic_memory.repository.semantic_vector_index import VectorKey, VectorMatch from basic_memory.schemas.search import SearchItemType, SearchRetrievalMode @@ -89,6 +90,42 @@ def _distance_to_similarity(self, distance: float) -> float: return 1.0 / (1.0 + max(distance, 0.0)) +@pytest.mark.asyncio +async def test_vector_match_hydration_batches_large_adapter_results() -> None: + """Deep vector pages must not create an unbounded SQL bind-parameter list.""" + repo = _ConcreteRepo() + repo._semantic_vector_index_name = "milvus" + repo._embedding_provider = SimpleNamespace(model_name="stub", dimensions=4) + matches = [ + VectorMatch( + key=VectorKey(entity_id=entity_id, chunk_key=f"entity:{entity_id}:0"), + similarity=0.9, + ) + for entity_id in range(600) + ] + session = AsyncMock() + + def hydrated_batch(_statement, params): + rows = [ + { + "entity_id": value, + "chunk_key": params[f"chunk_key_{index}"], + "chunk_text": f"chunk {value}", + } + for index in range((len(params) - 3) // 2) + if (value := params[f"entity_id_{index}"]) is not None + ] + return SimpleNamespace(mappings=lambda: SimpleNamespace(all=lambda: rows)) + + session.execute.side_effect = hydrated_batch + + hydrated = await repo._hydrate_vector_matches(session, matches) + + assert session.execute.await_count == 3 + assert max(len(call.args[1]) for call in session.execute.await_args_list) == 503 + assert [row["entity_id"] for row in hydrated] == list(range(600)) + + # --- SQLite SemanticSearchDisabledError --- diff --git a/tests/services/test_project_service_operations.py b/tests/services/test_project_service_operations.py index 7aa1219f1..867382959 100644 --- a/tests/services/test_project_service_operations.py +++ b/tests/services/test_project_service_operations.py @@ -9,6 +9,7 @@ import pytest from basic_memory import db +from basic_memory.services import project_service as project_service_module from basic_memory.services.project_service import ProjectService @@ -122,3 +123,43 @@ async def delete_after_vector_cleanup(session, entity_id: int) -> bool: search_repository.delete_project_vector_rows.assert_awaited_once_with( strict_adapter_cleanup=True ) + + +@pytest.mark.asyncio +async def test_remove_project_composes_vector_cleanup_without_injected_factory( + project_service: ProjectService, + monkeypatch, +): + """Legacy service construction must still preserve external vector ownership.""" + project_name = f"fallback-vector-project-{os.urandom(4).hex()}" + search_repository = SimpleNamespace(delete_project_vector_rows=AsyncMock()) + create_search_repository = Mock(return_value=search_repository) + monkeypatch.setattr( + project_service_module, + "create_search_repository", + create_search_repository, + ) + + with tempfile.TemporaryDirectory() as temp_dir: + async with db.scoped_session(project_service.session_maker) as session: + project = await project_service.repository.create( + session, + { + "name": project_name, + "path": temp_dir, + "permalink": project_name, + "is_active": True, + }, + ) + project_id = project.id + + await project_service.remove_project(project_name) + + create_search_repository.assert_called_once() + assert ( + create_search_repository.call_args.kwargs["session_maker"] is project_service.session_maker + ) + assert create_search_repository.call_args.kwargs["project_id"] == project_id + search_repository.delete_project_vector_rows.assert_awaited_once_with( + strict_adapter_cleanup=True + ) From aa4cf3e0a19c359cf8ba2d86d4be0c7d374cb2c0 Mon Sep 17 00:00:00 2001 From: phernandez Date: Sun, 26 Jul 2026 13:45:10 -0500 Subject: [PATCH 15/24] fix(core): address vector lifecycle review findings Signed-off-by: phernandez --- docs/semantic-search.md | 21 ++++++++++++----- .../repository/search_repository.py | 2 +- .../repository/search_repository_base.py | 18 +++++++-------- .../repository/semantic_vector_index.py | 7 +++++- .../semantic_vector_index_factory.py | 10 ++++++++ tests/repository/test_semantic_search_base.py | 11 +++++---- .../repository/test_semantic_vector_index.py | 23 +++++++++++++++++-- 7 files changed, 68 insertions(+), 24 deletions(-) diff --git a/docs/semantic-search.md b/docs/semantic-search.md index 5b9ee19ea..6a1a609d1 100644 --- a/docs/semantic-search.md +++ b/docs/semantic-search.md @@ -461,7 +461,7 @@ bm reindex -p my-project - **Dimension change**: After changing `semantic_embedding_dimensions` - **LiteLLM role change**: After changing `semantic_embedding_document_input_type` or `semantic_embedding_query_input_type` - **Literal prefix change**: After changing `semantic_embedding_document_prefix` or `semantic_embedding_query_prefix` -- **Vector index change**: After changing `semantic_vector_index` (the normal incremental sync also detects the change, while an explicit reindex completes the migration immediately) +- **Vector index change**: After completing the external-adapter cleanup procedure below and changing `semantic_vector_index` The reindex command shows progress with embedded/skipped/error counts: @@ -572,9 +572,13 @@ def create_index( ``` `VectorIndexScope` contains a stable, credential-free database namespace, project ID, embedding -identity, and vector dimensions. Extensions must isolate storage by the complete scope. They own -their client lifecycle, credentials, collection/index creation, vector persistence, and -nearest-neighbour implementation. +identity, and vector dimensions. Extensions must isolate physical storage by +`scope.storage_key`, which contains only the stable database namespace and project ID. +`embedding_identity` and `dimensions` describe the current vector schema for validation and +initialization; adapters must not use those mutable fields to create a second unreachable +project collection when the embedding configuration changes. Extensions own their client +lifecycle, credentials, collection/index creation, vector persistence, and nearest-neighbour +implementation. The returned `SemanticVectorIndex` has five asynchronous operations: @@ -605,5 +609,10 @@ searchable and the next sync safely retries the idempotent operation. Adapter se hydrated only through current, ready manifest rows, so stale or orphaned external matches fail closed. -Switching `semantic_vector_index` invalidates existing manifest rows for incremental re-embedding. -Run `bm reindex --embeddings` after a switch to migrate all eligible content immediately. +Basic Memory deliberately refuses to mutate manifest rows owned by a different external adapter. +Before switching `semantic_vector_index`, keep the old adapter configured and use that extension's +project-scope administrative cleanup to remove the old vectors. Remove the corresponding +`search_vector_chunks` manifest rows only after the external cleanup succeeds. Then switch the +configured adapter and run `bm reindex --embeddings` to populate the new store. If configuration +was switched too early, restore the old adapter first; the ownership check will continue to fail +closed until cleanup is completed. diff --git a/src/basic_memory/repository/search_repository.py b/src/basic_memory/repository/search_repository.py index 90db024d3..e282cc61c 100644 --- a/src/basic_memory/repository/search_repository.py +++ b/src/basic_memory/repository/search_repository.py @@ -111,7 +111,7 @@ async def delete_external_entity_vectors( """Delete DB-first entity vectors through an extension adapter.""" ... - async def delete_project_vector_rows(self, *, strict_adapter_cleanup: bool = False) -> None: + async def delete_project_vector_rows(self, *, strict_adapter_cleanup: bool = True) -> None: """Delete all semantic vector chunks and embeddings for this project.""" ... diff --git a/src/basic_memory/repository/search_repository_base.py b/src/basic_memory/repository/search_repository_base.py index bf93df4fc..1785bb368 100644 --- a/src/basic_memory/repository/search_repository_base.py +++ b/src/basic_memory/repository/search_repository_base.py @@ -736,18 +736,17 @@ def _assert_manifest_vector_ownership(self, vector_index_names: Iterable[object] f"{configured_index!r}. Restore the owning adapter and retry." ) - async def delete_project_vector_rows(self, *, strict_adapter_cleanup: bool = False) -> None: + async def delete_project_vector_rows(self, *, strict_adapter_cleanup: bool = True) -> None: """Delete this project's vectors through the configured storage adapter. Core enumerates ownership from the SQL manifest because the adapter contract intentionally has no project-wide listing or destructive reset. A full reindex clears the manifest even when semantic search is disabled, so stale ready rows cannot become current if the feature is re-enabled. - Adapter cleanup is best-effort because the manifest remains the search - authority. Ownership mismatches fail closed because only the manifest - identifies the extension that can safely remove previously written vectors. - Project deletion opts into strict cleanup so external data ownership is - preserved for a retry instead of being discarded after an adapter failure. + External adapter cleanup is strict by default because the manifest is the + only durable ownership record for remote vectors. Ownership mismatches + fail closed because only the owning extension can safely remove previously + written vectors. """ configured_index = self._semantic_vector_index_name async with db.scoped_session(self.session_maker) as session: @@ -822,10 +821,9 @@ async def delete_project_vector_rows(self, *, strict_adapter_cleanup: bool = Fal await self._semantic_vector_index.delete_entity(entity_id) except Exception as exc: # Trigger: a configured external adapter cannot initialize or delete. - # Why: reindex may discard derived state, but project deletion must - # not discard the only ownership manifest for external data. - # Outcome: strict callers stop for a retry; reindex logs the adapter - # failure and continues with a non-searchable empty manifest. + # Why: reindex and project deletion must not discard the only + # ownership manifest for external data. + # Outcome: strict callers stop for a retry before manifest deletion. logger.warning( "Could not clean semantic vector adapter: " "project_id={project_id} vector_index={vector_index} error={error}", diff --git a/src/basic_memory/repository/semantic_vector_index.py b/src/basic_memory/repository/semantic_vector_index.py index ece213428..3d55b9342 100644 --- a/src/basic_memory/repository/semantic_vector_index.py +++ b/src/basic_memory/repository/semantic_vector_index.py @@ -12,13 +12,18 @@ @dataclass(frozen=True, slots=True) class VectorIndexScope: - """Stable isolation and embedding identity for one project's vectors.""" + """Stable project storage identity plus the current embedding schema.""" namespace: str project_id: int embedding_identity: str dimensions: int + @property + def storage_key(self) -> tuple[str, int]: + """Return the stable isolation key external adapters must use for storage.""" + return (self.namespace, self.project_id) + @dataclass(frozen=True, slots=True) class VectorKey: diff --git a/src/basic_memory/repository/semantic_vector_index_factory.py b/src/basic_memory/repository/semantic_vector_index_factory.py index 3f36fe6bb..3eaff93ba 100644 --- a/src/basic_memory/repository/semantic_vector_index_factory.py +++ b/src/basic_memory/repository/semantic_vector_index_factory.py @@ -3,6 +3,7 @@ from __future__ import annotations import hashlib +import re from importlib.metadata import entry_points from typing import Protocol @@ -52,12 +53,21 @@ def _database_namespace(app_config: BasicMemoryConfig) -> str: """Derive a stable, credential-free namespace from the authoritative database.""" if app_config.database_url: url = make_url(app_config.database_url) + options = url.query.get("options", "") + serialized_options = "|".join(options) if isinstance(options, tuple) else str(options) + search_path_match = re.search( + r"(?:^|\s)-c\s*search_path=([^\s]+)", + serialized_options, + ) + search_path = search_path_match.group(1) if search_path_match else "" locator = "|".join( [ url.get_backend_name(), url.host or "", str(url.port or ""), url.database or "", + url.username or "", + search_path, ] ) else: diff --git a/tests/repository/test_semantic_search_base.py b/tests/repository/test_semantic_search_base.py index 4b89f2e4c..49ddf2049 100644 --- a/tests/repository/test_semantic_search_base.py +++ b/tests/repository/test_semantic_search_base.py @@ -282,8 +282,8 @@ async def fake_scoped_session(_session_maker): @pytest.mark.asyncio -async def test_project_vector_cleanup_clears_manifest_after_adapter_failure(monkeypatch): - """An extension failure should fail closed by removing searchable manifests.""" +async def test_project_vector_cleanup_preserves_manifest_after_adapter_failure(monkeypatch): + """A reindex must retain external ownership when its adapter is unavailable.""" repo = _ConcreteRepo() adapter: Any = SimpleNamespace( initialize=AsyncMock(side_effect=RuntimeError("adapter unavailable")), @@ -305,12 +305,15 @@ async def fake_scoped_session(_session_maker): monkeypatch.setattr(search_repository_base_module.db, "scoped_session", fake_scoped_session) monkeypatch.setattr(search_repository_base_module.logger, "warning", warning) - await repo.delete_project_vector_rows() + with pytest.raises(RuntimeError, match="adapter unavailable"): + await repo.delete_project_vector_rows() adapter.delete_entity.assert_not_awaited() warning.assert_called_once() statements = [str(call.args[0]) for call in session.execute.await_args_list] - assert any(statement.startswith("DELETE FROM search_vector_chunks") for statement in statements) + assert not any( + statement.startswith("DELETE FROM search_vector_chunks") for statement in statements + ) @pytest.mark.asyncio diff --git a/tests/repository/test_semantic_vector_index.py b/tests/repository/test_semantic_vector_index.py index 5a32c26b7..95a141113 100644 --- a/tests/repository/test_semantic_vector_index.py +++ b/tests/repository/test_semantic_vector_index.py @@ -132,20 +132,39 @@ def test_scope_is_stable_credential_free_and_project_isolated() -> None: rotated_password = build_vector_index_scope( _postgres_config( database_url=( - "postgresql+asyncpg://rotated-user:new-secret@db.example.test:5432/memory" - "?sslmode=require" + "postgresql+asyncpg://user:new-secret@db.example.test:5432/memory?sslmode=require" ) ), provider, project_id=7, ) other_project = build_vector_index_scope(_postgres_config(), provider, project_id=8) + other_user = build_vector_index_scope( + _postgres_config( + database_url="postgresql+asyncpg://tenant-user:secret@db.example.test:5432/memory" + ), + provider, + project_id=7, + ) + other_schema = build_vector_index_scope( + _postgres_config( + database_url=( + "postgresql+asyncpg://user:secret@db.example.test:5432/memory" + "?options=-csearch_path%3Dtenant" + ) + ), + provider, + project_id=7, + ) assert first.namespace == rotated_password.namespace assert "secret" not in first.namespace assert first.project_id != other_project.project_id + assert first.namespace != other_user.namespace + assert first.namespace != other_schema.namespace assert first.embedding_identity == "StubEmbeddingProvider:stub-model:3" assert first.dimensions == 3 + assert first.storage_key == rotated_password.storage_key def test_missing_configured_extension_fails_without_fallback(monkeypatch) -> None: From 89062ae3f8a8c1044f975824ecadf93d753a8aa3 Mon Sep 17 00:00:00 2001 From: phernandez Date: Sun, 26 Jul 2026 13:54:44 -0500 Subject: [PATCH 16/24] fix(core): forward external vector cleanup Signed-off-by: phernandez --- src/basic_memory/index/local_project.py | 1 + ...est_local_project_vector_cleaner_wiring.py | 43 +++++++++++++++++++ 2 files changed, 44 insertions(+) create mode 100644 tests/index/test_local_project_vector_cleaner_wiring.py diff --git a/src/basic_memory/index/local_project.py b/src/basic_memory/index/local_project.py index 8826c99db..a781a034c 100644 --- a/src/basic_memory/index/local_project.py +++ b/src/basic_memory/index/local_project.py @@ -592,6 +592,7 @@ def runtime_from_dependencies( maintenance_store = RepositoryProjectIndexMaintenanceStore( session_maker=dependencies.session_maker, project_id=dependencies.project_id, + external_vector_cleaner=dependencies.external_vector_cleaner, move_content_updater=LocalProjectIndexMoveContentUpdater( entity_service=dependencies.entity_service, file_service=dependencies.file_service, diff --git a/tests/index/test_local_project_vector_cleaner_wiring.py b/tests/index/test_local_project_vector_cleaner_wiring.py new file mode 100644 index 000000000..fcc7c9571 --- /dev/null +++ b/tests/index/test_local_project_vector_cleaner_wiring.py @@ -0,0 +1,43 @@ +"""Composition coverage for external vector cleanup in full project indexing.""" + +from unittest.mock import Mock + +from basic_memory.index.local_dependencies import LocalIndexProjectDependencies +from basic_memory.index.local_project import LocalProjectIndexRuntimeFactory +from basic_memory.indexing.project_index_maintenance import ( + RepositoryProjectIndexMaintenanceStore, + StoreProjectIndexMaintenanceRunner, +) + + +class RuntimeFactoryEntityService: + app_config = None + + async def resolve_permalink(self, *args: object, **kwargs: object) -> str: + return "local" + + +def test_full_project_runtime_forwards_external_vector_cleaner() -> None: + cleaner = Mock() + dependencies = LocalIndexProjectDependencies( + file_service=Mock(), + file_indexer=Mock(), + file_batch_indexer=Mock(), + session_maker=Mock(), + project_id=42, + entity_repository=Mock(), + relation_repository=Mock(), + link_resolver=Mock(), + search_service=Mock(), + entity_service=RuntimeFactoryEntityService(), + external_vector_cleaner=cleaner, + ) + + runtime = LocalProjectIndexRuntimeFactory().runtime_from_dependencies(dependencies) + + assert isinstance(runtime.maintenance_runner, StoreProjectIndexMaintenanceRunner) + assert isinstance( + runtime.maintenance_runner.delete_store, + RepositoryProjectIndexMaintenanceStore, + ) + assert runtime.maintenance_runner.delete_store.external_vector_cleaner is cleaner From cbf8e50b3b934697491789c1032d8e1b6ca23cce Mon Sep 17 00:00:00 2001 From: phernandez Date: Sun, 26 Jul 2026 14:07:51 -0500 Subject: [PATCH 17/24] fix(core): harden vector generation isolation Signed-off-by: phernandez --- src/basic_memory/repository/pgvector_index.py | 3 +- .../repository/search_repository_base.py | 60 ++++++--- .../semantic_vector_index_factory.py | 12 +- tests/repository/test_pgvector_index.py | 2 + tests/repository/test_semantic_search_base.py | 119 +++++++++++++++++- .../repository/test_semantic_vector_index.py | 19 +++ .../test_sqlite_vector_search_repository.py | 2 +- 7 files changed, 192 insertions(+), 25 deletions(-) diff --git a/src/basic_memory/repository/pgvector_index.py b/src/basic_memory/repository/pgvector_index.py index 9f0693bb6..44b5a5356 100644 --- a/src/basic_memory/repository/pgvector_index.py +++ b/src/basic_memory/repository/pgvector_index.py @@ -117,7 +117,8 @@ async def _existing_dimensions(self, session: AsyncSession) -> int | None: exists = await session.execute( text( "SELECT 1 FROM information_schema.tables " - "WHERE table_name = 'search_vector_embeddings'" + "WHERE table_schema = ANY (current_schemas(false)) " + "AND table_name = 'search_vector_embeddings'" ) ) if exists.fetchone() is None: diff --git a/src/basic_memory/repository/search_repository_base.py b/src/basic_memory/repository/search_repository_base.py index 1785bb368..eefc5e59f 100644 --- a/src/basic_memory/repository/search_repository_base.py +++ b/src/basic_memory/repository/search_repository_base.py @@ -1,5 +1,6 @@ """Abstract base class for search repository implementations.""" +import hashlib import time from abc import ABC, abstractmethod from collections.abc import Iterable, Sequence @@ -327,46 +328,65 @@ async def _persist_embeddings( return row_ids = [row_id for row_id, _ in jobs] - params = {f"row_id_{index}": row_id for index, row_id in enumerate(row_ids)} - placeholders = ", ".join(f":row_id_{index}" for index in range(len(row_ids))) + lookup_params = {f"row_id_{index}": row_id for index, row_id in enumerate(row_ids)} + lookup_placeholders = ", ".join(f":row_id_{index}" for index in range(len(row_ids))) async with db.scoped_session(self.session_maker) as session: result = await session.execute( text( - "SELECT id, entity_id, chunk_key FROM search_vector_chunks " - f"WHERE project_id = :project_id AND id IN ({placeholders})" + "SELECT id, entity_id, chunk_key, source_hash FROM search_vector_chunks " + f"WHERE project_id = :project_id AND id IN ({lookup_placeholders})" ), - {**params, "project_id": self.project_id}, + {**lookup_params, "project_id": self.project_id}, ) - keys_by_id = { - int(row["id"]): VectorKey( - entity_id=int(row["entity_id"]), - chunk_key=str(row["chunk_key"]), - ) - for row in result.mappings().all() - } + rows_by_id = {int(row["id"]): row for row in result.mappings().all()} - missing_row_ids = [row_id for row_id in row_ids if row_id not in keys_by_id] + missing_row_ids = [row_id for row_id in row_ids if row_id not in rows_by_id] if missing_row_ids: raise RuntimeError(f"Vector manifest rows disappeared before write: {missing_row_ids}") + current_jobs: list[tuple[int, str, str, list[float]]] = [] + for (row_id, chunk_text), embedding in zip(jobs, embeddings, strict=True): + expected_source_hash = hashlib.sha256(chunk_text.encode("utf-8")).hexdigest() + if str(rows_by_id[row_id]["source_hash"]) != expected_source_hash: + continue + current_jobs.append((row_id, chunk_text, expected_source_hash, embedding)) + if not current_jobs: + return + + params: dict[str, object] = {} + generation_predicates: list[str] = [] records = [ - VectorRecord(key=keys_by_id[row_id], values=tuple(embedding)) - for (row_id, _), embedding in zip(jobs, embeddings, strict=True) + VectorRecord( + key=VectorKey( + entity_id=int(rows_by_id[row_id]["entity_id"]), + chunk_key=str(rows_by_id[row_id]["chunk_key"]), + ), + values=tuple(embedding), + ) + for row_id, _chunk_text, _source_hash, embedding in current_jobs ] + for index, (row_id, _chunk_text, source_hash, _embedding) in enumerate(current_jobs): + params[f"row_id_{index}"] = row_id + params[f"source_hash_{index}"] = source_hash + generation_predicates.append( + f"(id = :row_id_{index} AND source_hash = :source_hash_{index})" + ) await self._semantic_vector_index.upsert(records) # Trigger: the adapter write completed successfully. - # Why: only SQL rows marked ready may hydrate search matches. If this commit - # fails, the adapter's idempotent upsert is retried while the row stays pending. - # Outcome: external partial failures fail closed without cross-store transactions. + # Why: only SQL rows marked ready may hydrate search matches, and a newer sync + # may have replaced a chunk while this job was embedding its previous text. + # Outcome: current source generations become ready; stale completions leave the + # newer manifest pending for retry. async with db.scoped_session(self.session_maker) as session: await session.execute( text( "UPDATE search_vector_chunks SET embedding_status = 'ready', " f"updated_at = {self._timestamp_now_expr()} " - f"WHERE project_id = :project_id AND id IN ({placeholders}) " + "WHERE project_id = :project_id " "AND vector_index = :vector_index " - "AND embedding_model = :embedding_model" + "AND embedding_model = :embedding_model " + "AND (" + " OR ".join(generation_predicates) + ")" ), { **params, diff --git a/src/basic_memory/repository/semantic_vector_index_factory.py b/src/basic_memory/repository/semantic_vector_index_factory.py index 3eaff93ba..f84081e68 100644 --- a/src/basic_memory/repository/semantic_vector_index_factory.py +++ b/src/basic_memory/repository/semantic_vector_index_factory.py @@ -49,12 +49,18 @@ def semantic_embedding_identity(provider: EmbeddingProvider) -> str: return f"{type(provider).__name__}:{embedding_provider_identity(provider)}" +def _serialize_query_value(value: str | tuple[str, ...] | None) -> str: + """Serialize SQLAlchemy URL query values without dropping repeated fields.""" + if isinstance(value, tuple): + return "|".join(value) + return value or "" + + def _database_namespace(app_config: BasicMemoryConfig) -> str: """Derive a stable, credential-free namespace from the authoritative database.""" if app_config.database_url: url = make_url(app_config.database_url) - options = url.query.get("options", "") - serialized_options = "|".join(options) if isinstance(options, tuple) else str(options) + serialized_options = _serialize_query_value(url.query.get("options")) search_path_match = re.search( r"(?:^|\s)-c\s*search_path=([^\s]+)", serialized_options, @@ -65,6 +71,8 @@ def _database_namespace(app_config: BasicMemoryConfig) -> str: url.get_backend_name(), url.host or "", str(url.port or ""), + _serialize_query_value(url.query.get("host")), + _serialize_query_value(url.query.get("port")), url.database or "", url.username or "", search_path, diff --git a/tests/repository/test_pgvector_index.py b/tests/repository/test_pgvector_index.py index f52301733..e7abd0582 100644 --- a/tests/repository/test_pgvector_index.py +++ b/tests/repository/test_pgvector_index.py @@ -137,6 +137,8 @@ async def test_initialize_preserves_manifest_when_storage_is_unchanged(monkeypat sql_calls = _sql_calls(session) assert not any("DROP TABLE IF EXISTS search_vector_embeddings" in sql for sql in sql_calls) assert not any("embedding_status = 'pending'" in sql for sql in sql_calls) + table_probe = next(sql for sql in sql_calls if "information_schema.tables" in sql) + assert "table_schema = ANY (current_schemas(false))" in table_probe @pytest.mark.asyncio diff --git a/tests/repository/test_semantic_search_base.py b/tests/repository/test_semantic_search_base.py index 49ddf2049..67821fee7 100644 --- a/tests/repository/test_semantic_search_base.py +++ b/tests/repository/test_semantic_search_base.py @@ -1,6 +1,8 @@ """Tests for semantic search orchestration in SearchRepositoryBase.""" import asyncio +import hashlib +from collections.abc import Sequence from contextlib import asynccontextmanager from datetime import datetime from types import SimpleNamespace @@ -21,7 +23,12 @@ SemanticSearchDisabledError, SemanticVectorIndexExtensionError, ) -from basic_memory.repository.semantic_vector_index import VectorKey, VectorMatch +from basic_memory.repository.semantic_vector_index import ( + VectorIndexScope, + VectorKey, + VectorMatch, + VectorRecord, +) from basic_memory.schemas.search import SearchItemType, SearchRetrievalMode @@ -90,6 +97,40 @@ def _distance_to_similarity(self, distance: float) -> float: return 1.0 / (1.0 + max(distance, 0.0)) +class _RecordingVectorIndex: + """Protocol-complete adapter that records generation-safe upserts.""" + + scope = VectorIndexScope( + namespace="basic-memory-test", + project_id=1, + embedding_identity="stub:4", + dimensions=4, + ) + + def __init__(self) -> None: + self.upserted_records: list[VectorRecord] = [] + + async def initialize(self) -> None: + return None + + async def upsert(self, records: Sequence[VectorRecord]) -> None: + self.upserted_records.extend(records) + + async def delete(self, keys: Sequence[VectorKey]) -> None: + return None + + async def delete_entity(self, entity_id: int) -> None: + return None + + async def search( + self, + query: Sequence[float], + *, + limit: int, + ) -> list[VectorMatch]: + return [] + + @pytest.mark.asyncio async def test_vector_match_hydration_batches_large_adapter_results() -> None: """Deep vector pages must not create an unbounded SQL bind-parameter list.""" @@ -126,6 +167,82 @@ def hydrated_batch(_statement, params): assert [row["entity_id"] for row in hydrated] == list(range(600)) +@pytest.mark.asyncio +async def test_embedding_persistence_skips_stale_source_generation(monkeypatch) -> None: + """An obsolete embedding job must not overwrite the current adapter record.""" + repo = _ConcreteRepo() + repo._semantic_vector_index_name = "milvus" + repo._embedding_provider = SimpleNamespace(model_name="stub", dimensions=4) + adapter = _RecordingVectorIndex() + repo._semantic_vector_index = adapter + current_text = "new chunk text" + session = AsyncMock() + session.execute.return_value = SimpleNamespace( + mappings=lambda: SimpleNamespace( + all=lambda: [ + { + "id": 7, + "entity_id": 41, + "chunk_key": "entity:41:0", + "source_hash": hashlib.sha256(current_text.encode("utf-8")).hexdigest(), + } + ] + ) + ) + + @asynccontextmanager + async def fake_scoped_session(_session_maker): + yield session + + monkeypatch.setattr(search_repository_base_module.db, "scoped_session", fake_scoped_session) + + await repo._persist_embeddings([(7, "old chunk text")], [[1.0, 0.0, 0.0, 0.0]]) + + assert adapter.upserted_records == [] + assert session.execute.await_count == 1 + + +@pytest.mark.asyncio +async def test_embedding_ready_update_requires_source_generation(monkeypatch) -> None: + """Ready state must belong to the exact source text that produced the vector.""" + repo = _ConcreteRepo() + repo._semantic_vector_index_name = "milvus" + repo._embedding_provider = SimpleNamespace(model_name="stub", dimensions=4) + adapter = _RecordingVectorIndex() + repo._semantic_vector_index = adapter + chunk_text = "current chunk text" + source_hash = hashlib.sha256(chunk_text.encode("utf-8")).hexdigest() + session = AsyncMock() + session.execute.side_effect = [ + SimpleNamespace( + mappings=lambda: SimpleNamespace( + all=lambda: [ + { + "id": 7, + "entity_id": 41, + "chunk_key": "entity:41:0", + "source_hash": source_hash, + } + ] + ) + ), + SimpleNamespace(), + ] + + @asynccontextmanager + async def fake_scoped_session(_session_maker): + yield session + + monkeypatch.setattr(search_repository_base_module.db, "scoped_session", fake_scoped_session) + + await repo._persist_embeddings([(7, chunk_text)], [[1.0, 0.0, 0.0, 0.0]]) + + assert len(adapter.upserted_records) == 1 + ready_statement, ready_params = session.execute.await_args_list[1].args + assert "source_hash = :source_hash_0" in str(ready_statement) + assert ready_params["source_hash_0"] == source_hash + + # --- SQLite SemanticSearchDisabledError --- diff --git a/tests/repository/test_semantic_vector_index.py b/tests/repository/test_semantic_vector_index.py index 95a141113..8a2370f5d 100644 --- a/tests/repository/test_semantic_vector_index.py +++ b/tests/repository/test_semantic_vector_index.py @@ -156,12 +156,31 @@ def test_scope_is_stable_credential_free_and_project_isolated() -> None: provider, project_id=7, ) + first_socket = build_vector_index_scope( + _postgres_config( + database_url=( + "postgresql+asyncpg://user:secret@/memory?host=%2Fvar%2Frun%2Fpostgresql&port=5432" + ) + ), + provider, + project_id=7, + ) + other_socket = build_vector_index_scope( + _postgres_config( + database_url=( + "postgresql+asyncpg://user:secret@/memory?host=%2Ftmp%2Fpostgresql&port=5433" + ) + ), + provider, + project_id=7, + ) assert first.namespace == rotated_password.namespace assert "secret" not in first.namespace assert first.project_id != other_project.project_id assert first.namespace != other_user.namespace assert first.namespace != other_schema.namespace + assert first_socket.namespace != other_socket.namespace assert first.embedding_identity == "StubEmbeddingProvider:stub-model:3" assert first.dimensions == 3 assert first.storage_key == rotated_password.storage_key diff --git a/tests/repository/test_sqlite_vector_search_repository.py b/tests/repository/test_sqlite_vector_search_repository.py index 48da2caee..99161348c 100644 --- a/tests/repository/test_sqlite_vector_search_repository.py +++ b/tests/repository/test_sqlite_vector_search_repository.py @@ -565,7 +565,7 @@ async def test_ready_commit_failure_retries_same_stable_adapter_key( "project_id": search_repository.project_id, "chunk_key": "entity:115:0", "chunk_text": "ready commit retry", - "source_hash": "hash", + "source_hash": hashlib.sha256(b"ready commit retry").hexdigest(), "entity_fingerprint": "fingerprint", "embedding_model": search_repository._embedding_model_key(), "vector_index": "recording", From 58d434f850340ea99acc6da2c76e07e0b8518fd8 Mon Sep 17 00:00:00 2001 From: phernandez Date: Sun, 26 Jul 2026 14:27:14 -0500 Subject: [PATCH 18/24] fix(core): serialize vector generation writes Signed-off-by: phernandez --- docs/semantic-search.md | 18 +- src/basic_memory/models/search.py | 5 +- src/basic_memory/repository/pgvector_index.py | 75 +++++++-- .../repository/search_repository_base.py | 155 ++++++++++++------ .../repository/semantic_vector_index.py | 5 +- .../repository/sqlite_vec_index.py | 68 ++++++-- test-int/test_embedding_status_vec0.py | 16 ++ tests/repository/test_pgvector_index.py | 65 +++++++- tests/repository/test_semantic_search_base.py | 101 ++++++++++++ .../repository/test_semantic_vector_index.py | 7 +- .../test_sqlite_vector_search_repository.py | 2 +- 11 files changed, 422 insertions(+), 95 deletions(-) diff --git a/docs/semantic-search.md b/docs/semantic-search.md index 6a1a609d1..022498255 100644 --- a/docs/semantic-search.md +++ b/docs/semantic-search.md @@ -511,7 +511,7 @@ Vector and hybrid modes return individual observations and relations as first-cl - **Vector storage**: [sqlite-vec](https://github.com/asg017/sqlite-vec) virtual table - **Table creation**: At runtime when semantic search is first used — no migration needed -- **Embedding table**: `search_vector_embeddings` using `vec0(embedding float[N])` where N is the configured dimensions +- **Embedding table**: `search_vector_embeddings` using `vec0(embedding float[N], +source_hash text)` where N is the configured dimensions - **Chunk metadata**: `search_vector_chunks` table stores chunk text, keys, and source hashes The sqlite-vec extension is loaded per-connection. Vector tables are created lazily on first use. @@ -583,13 +583,18 @@ implementation. The returned `SemanticVectorIndex` has five asynchronous operations: - `initialize()` validates or creates backend storage. -- `upsert(records)` idempotently writes vectors by `(entity_id, chunk_key)`. +- `upsert(records)` idempotently writes vectors by `(entity_id, chunk_key)` for each record's + `source_hash` generation. - `delete(keys)` removes stable keys; missing keys are successful no-ops. - `delete_entity(entity_id)` removes all vectors for one entity in the scope. - `search(query, limit)` returns stable keys with normalized cosine similarity in `[0, 1]`. The adapter never receives a SQLAlchemy session and never calls the embedding provider. Basic Memory owns chunking and embedding, while the extension owns vector storage and lookup. +Each `VectorRecord` carries the SHA-256 source generation that produced its values. Basic Memory +holds the matching SQL manifest locked across extension adapter I/O and the ready-state transition +(`FOR UPDATE` on Postgres and a conditional write lock on SQLite), so an older overlapping sync +cannot overwrite a newer generation under the same stable adapter key. Adapters may also implement the separate `SemanticVectorIndexReconciler` capability. After a vector reindex, Basic Memory passes it the complete set of current ready keys so the adapter can @@ -604,10 +609,11 @@ store. Each row records the selected `vector_index`, embedding identity, stable `embedding_status` of `pending` or `ready`. Writes and deletes commit `pending` before calling the adapter. A successful adapter operation then -makes the manifest row ready or removes it. If the external operation fails, the pending row is not -searchable and the next sync safely retries the idempotent operation. Adapter search results are -hydrated only through current, ready manifest rows, so stale or orphaned external matches fail -closed. +makes the manifest row ready or removes it. Vector writes are generation checked inside built-in +adapter transactions; extension writes retain the manifest lock across adapter I/O. If the external +operation fails, the pending row is not searchable and the next sync safely retries the idempotent +operation. Adapter search results are hydrated only through current, ready manifest rows, so stale +or orphaned external matches fail closed. Basic Memory deliberately refuses to mutate manifest rows owned by a different external adapter. Before switching `semantic_vector_index`, keep the old adapter configured and use that extension's diff --git a/src/basic_memory/models/search.py b/src/basic_memory/models/search.py index ee50c7a36..1c9fff6d5 100644 --- a/src/basic_memory/models/search.py +++ b/src/basic_memory/models/search.py @@ -152,6 +152,9 @@ def create_sqlite_search_vector_embeddings(dimensions: int) -> DDL: return DDL( f""" CREATE VIRTUAL TABLE IF NOT EXISTS search_vector_embeddings -USING vec0(embedding float[{dimensions}]) +USING vec0( + 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 44b5a5356..f68be1ba5 100644 --- a/src/basic_memory/repository/pgvector_index.py +++ b/src/basic_memory/repository/pgvector_index.py @@ -58,15 +58,21 @@ async def initialize(self) -> None: existing_dimensions = await self._existing_dimensions(session) storage_missing = existing_dimensions is None + source_hash_missing = ( + existing_dimensions is not None + and not await self._has_source_hash_column(session) + ) dimensions_changed = ( existing_dimensions is not None and existing_dimensions != self.scope.dimensions ) - if dimensions_changed: + if dimensions_changed or source_hash_missing: logger.warning( - "Embedding dimension mismatch: table has {existing}, " - "provider expects {expected}. Recreating vector storage.", + "Vector storage schema mismatch: table dimensions={existing}, " + "provider dimensions={expected}, source_hash_missing={source_hash_missing}. " + "Recreating vector storage.", existing=existing_dimensions, expected=self.scope.dimensions, + source_hash_missing=source_hash_missing, ) await session.execute(text("DROP TABLE IF EXISTS search_vector_embeddings")) @@ -78,6 +84,7 @@ async def initialize(self) -> None: project_id INTEGER NOT NULL, embedding vector({self.scope.dimensions}) NOT NULL, embedding_dims INTEGER NOT NULL, + source_hash TEXT NOT NULL, updated_at TIMESTAMPTZ NOT NULL DEFAULT NOW() ) """) @@ -102,7 +109,7 @@ async def initialize(self) -> None: # Why: SQL manifest rows can otherwise remain `ready` after their vectors # disappeared, causing incremental sync to skip them forever. # Outcome: the normal sync pipeline re-embeds every affected chunk. - if storage_missing or dimensions_changed: + if storage_missing or dimensions_changed or source_hash_missing: await session.execute( text( "UPDATE search_vector_chunks SET embedding_status = 'pending' " @@ -134,6 +141,16 @@ async def _existing_dimensions(self, session: AsyncSession) -> int | None: value = result.scalar_one_or_none() return int(value) if value is not None else None + async def _has_source_hash_column(self, session: AsyncSession) -> bool: + result = await session.execute( + text( + "SELECT 1 FROM pg_attribute " + "WHERE attrelid = 'search_vector_embeddings'::regclass " + "AND attname = 'source_hash'" + ) + ) + return result.scalar_one_or_none() is not None + async def _chunk_ids_by_key( self, session: AsyncSession, @@ -171,30 +188,64 @@ async def upsert(self, records: Sequence[VectorRecord]) -> None: await self.initialize() async with db.scoped_session(self._session_maker) as session: - ids_by_key = await self._chunk_ids_by_key(session, [record.key for record in records]) - missing = [record.key for record in records if record.key not in ids_by_key] + keys = [record.key for record in records] + params: dict[str, object] = {"project_id": self.scope.project_id} + predicates: list[str] = [] + for index, key in enumerate(keys): + params[f"entity_id_{index}"] = key.entity_id + params[f"chunk_key_{index}"] = key.chunk_key + predicates.append( + f"(entity_id = :entity_id_{index} AND chunk_key = :chunk_key_{index})" + ) + result = await session.execute( + text( + "SELECT id, entity_id, chunk_key, source_hash " + "FROM search_vector_chunks " + "WHERE project_id = :project_id AND (" + + " OR ".join(predicates) + + ") FOR UPDATE" + ), + params, + ) + manifest_by_key = { + VectorKey( + entity_id=int(row["entity_id"]), + chunk_key=str(row["chunk_key"]), + ): (int(row["id"]), str(row["source_hash"])) + for row in result.mappings().all() + } + missing = [key for key in keys if key not in manifest_by_key] if missing: raise RuntimeError(f"Vector manifest rows are missing for keys: {missing!r}") - params: dict[str, object] = {"project_id": self.scope.project_id} + current_records = [ + record for record in records if manifest_by_key[record.key][1] == record.source_hash + ] + if not current_records: + return + + params = {"project_id": self.scope.project_id} values: list[str] = [] - for index, record in enumerate(records): - params[f"chunk_id_{index}"] = ids_by_key[record.key] + for index, record in enumerate(current_records): + params[f"chunk_id_{index}"] = manifest_by_key[record.key][0] params[f"embedding_{index}"] = self._format_vector(record.values) params[f"dimensions_{index}"] = len(record.values) + params[f"source_hash_{index}"] = record.source_hash values.append( f"(:chunk_id_{index}, :project_id, " - f"CAST(:embedding_{index} AS vector), :dimensions_{index}, NOW())" + f"CAST(:embedding_{index} AS vector), :dimensions_{index}, " + f":source_hash_{index}, NOW())" ) await session.execute( text(f""" INSERT INTO search_vector_embeddings ( - chunk_id, project_id, embedding, embedding_dims, updated_at + chunk_id, project_id, embedding, embedding_dims, source_hash, updated_at ) VALUES {", ".join(values)} ON CONFLICT (chunk_id) DO UPDATE SET project_id = EXCLUDED.project_id, embedding = EXCLUDED.embedding, embedding_dims = EXCLUDED.embedding_dims, + source_hash = EXCLUDED.source_hash, updated_at = NOW() """), params, @@ -245,6 +296,7 @@ async def delete_orphans(self, _live_keys: Sequence[VectorKey]) -> None: "AND chunks.project_id = :project_id " "AND chunks.vector_index = 'pgvector' " "AND chunks.embedding_model = :embedding_identity " + "AND chunks.source_hash = embeddings.source_hash " "AND chunks.embedding_status = 'ready')" ), { @@ -277,6 +329,7 @@ async def search( "AND c.vector_index = 'pgvector' " "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 " "LIMIT :limit" diff --git a/src/basic_memory/repository/search_repository_base.py b/src/basic_memory/repository/search_repository_base.py index eefc5e59f..738628dac 100644 --- a/src/basic_memory/repository/search_repository_base.py +++ b/src/basic_memory/repository/search_repository_base.py @@ -331,72 +331,121 @@ async def _persist_embeddings( lookup_params = {f"row_id_{index}": row_id for index, row_id in enumerate(row_ids)} lookup_placeholders = ", ".join(f":row_id_{index}" for index in range(len(row_ids))) async with db.scoped_session(self.session_maker) as session: - result = await session.execute( - text( + connection = await session.connection() + dialect_name = connection.dialect.name + external_vector_index = ( + self._semantic_vector_index_name not in _BUILT_IN_VECTOR_INDEX_NAMES + ) + lock_external_write = external_vector_index and dialect_name in {"postgresql", "sqlite"} + if external_vector_index and dialect_name == "sqlite": + # SQLite has no SELECT FOR UPDATE. A conditional no-op write takes + # the database write lock before extension I/O, serializing a newer + # manifest generation behind this adapter write and ready commit. + lookup_statement = ( + "UPDATE search_vector_chunks SET source_hash = source_hash " + f"WHERE project_id = :project_id AND id IN ({lookup_placeholders}) " + "RETURNING id, entity_id, chunk_key, source_hash" + ) + else: + lock_clause = ( + " FOR UPDATE" if external_vector_index and dialect_name == "postgresql" else "" + ) + lookup_statement = ( "SELECT id, entity_id, chunk_key, source_hash FROM search_vector_chunks " f"WHERE project_id = :project_id AND id IN ({lookup_placeholders})" - ), + f"{lock_clause}" + ) + result = await session.execute( + text(lookup_statement), {**lookup_params, "project_id": self.project_id}, ) rows_by_id = {int(row["id"]): row for row in result.mappings().all()} - missing_row_ids = [row_id for row_id in row_ids if row_id not in rows_by_id] - if missing_row_ids: - raise RuntimeError(f"Vector manifest rows disappeared before write: {missing_row_ids}") + missing_row_ids = [row_id for row_id in row_ids if row_id not in rows_by_id] + if missing_row_ids: + raise RuntimeError( + f"Vector manifest rows disappeared before write: {missing_row_ids}" + ) - current_jobs: list[tuple[int, str, str, list[float]]] = [] - for (row_id, chunk_text), embedding in zip(jobs, embeddings, strict=True): - expected_source_hash = hashlib.sha256(chunk_text.encode("utf-8")).hexdigest() - if str(rows_by_id[row_id]["source_hash"]) != expected_source_hash: - continue - current_jobs.append((row_id, chunk_text, expected_source_hash, embedding)) - if not current_jobs: - return + current_jobs: list[tuple[int, str, str, list[float]]] = [] + for (row_id, chunk_text), embedding in zip(jobs, embeddings, strict=True): + expected_source_hash = hashlib.sha256(chunk_text.encode("utf-8")).hexdigest() + if str(rows_by_id[row_id]["source_hash"]) != expected_source_hash: + continue + current_jobs.append((row_id, chunk_text, expected_source_hash, embedding)) + if not current_jobs: + return - params: dict[str, object] = {} - generation_predicates: list[str] = [] - records = [ - VectorRecord( - key=VectorKey( - entity_id=int(rows_by_id[row_id]["entity_id"]), - chunk_key=str(rows_by_id[row_id]["chunk_key"]), - ), - values=tuple(embedding), - ) - for row_id, _chunk_text, _source_hash, embedding in current_jobs - ] - for index, (row_id, _chunk_text, source_hash, _embedding) in enumerate(current_jobs): - params[f"row_id_{index}"] = row_id - params[f"source_hash_{index}"] = source_hash - generation_predicates.append( - f"(id = :row_id_{index} AND source_hash = :source_hash_{index})" - ) - await self._semantic_vector_index.upsert(records) + params: dict[str, object] = {} + generation_predicates: list[str] = [] + records = [ + VectorRecord( + key=VectorKey( + entity_id=int(rows_by_id[row_id]["entity_id"]), + chunk_key=str(rows_by_id[row_id]["chunk_key"]), + ), + source_hash=source_hash, + values=tuple(embedding), + ) + for row_id, _chunk_text, source_hash, embedding in current_jobs + ] + for index, (row_id, _chunk_text, source_hash, _embedding) in enumerate(current_jobs): + params[f"row_id_{index}"] = row_id + params[f"source_hash_{index}"] = source_hash + generation_predicates.append( + f"(id = :row_id_{index} AND source_hash = :source_hash_{index})" + ) + + if lock_external_write: + # Constraint: extension adapters use stable logical keys outside + # the authoritative SQL database. Hold its manifest lock across + # adapter I/O so a newer prepare cannot advance this generation + # before the external write and ready transition complete. + await self._semantic_vector_index.upsert(records) + await self._mark_embedding_jobs_ready( + session, + params=params, + generation_predicates=generation_predicates, + ) + await session.commit() + return - # Trigger: the adapter write completed successfully. - # Why: only SQL rows marked ready may hydrate search matches, and a newer sync - # may have replaced a chunk while this job was embedding its previous text. - # Outcome: current source generations become ready; stale completions leave the - # newer manifest pending for retry. + # Built-in adapters share the authoritative database. They verify and lock + # each record's source_hash inside the same transaction as their vector write. + await self._semantic_vector_index.upsert(records) async with db.scoped_session(self.session_maker) as session: - await session.execute( - text( - "UPDATE search_vector_chunks SET embedding_status = 'ready', " - f"updated_at = {self._timestamp_now_expr()} " - "WHERE project_id = :project_id " - "AND vector_index = :vector_index " - "AND embedding_model = :embedding_model " - "AND (" + " OR ".join(generation_predicates) + ")" - ), - { - **params, - "project_id": self.project_id, - "vector_index": self._semantic_vector_index_name, - "embedding_model": self._embedding_model_key(), - }, + await self._mark_embedding_jobs_ready( + session, + params=params, + generation_predicates=generation_predicates, ) await session.commit() + async def _mark_embedding_jobs_ready( + self, + session: AsyncSession, + *, + params: dict[str, object], + generation_predicates: list[str], + ) -> None: + """Publish only the source generations written by the adapter.""" + await session.execute( + text( + "UPDATE search_vector_chunks SET embedding_status = 'ready', " + f"updated_at = {self._timestamp_now_expr()} " + "WHERE project_id = :project_id " + "AND vector_index = :vector_index " + "AND embedding_model = :embedding_model " + "AND (" + " OR ".join(generation_predicates) + ")" + ), + { + **params, + "project_id": self.project_id, + "vector_index": self._semantic_vector_index_name, + "embedding_model": self._embedding_model_key(), + }, + ) + async def _delete_entity_chunks( self, session: AsyncSession, diff --git a/src/basic_memory/repository/semantic_vector_index.py b/src/basic_memory/repository/semantic_vector_index.py index 3d55b9342..982b69e23 100644 --- a/src/basic_memory/repository/semantic_vector_index.py +++ b/src/basic_memory/repository/semantic_vector_index.py @@ -35,9 +35,10 @@ class VectorKey: @dataclass(frozen=True, slots=True) class VectorRecord: - """One vector value to insert or replace idempotently.""" + """One source-generation vector value to insert or replace idempotently.""" key: VectorKey + source_hash: str values: tuple[float, ...] @@ -65,7 +66,7 @@ async def initialize(self) -> None: ... async def upsert(self, records: Sequence[VectorRecord]) -> None: - """Insert or replace vectors by stable key.""" + """Insert or replace vectors only for each record's source generation.""" ... async def delete(self, keys: Sequence[VectorKey]) -> None: diff --git a/src/basic_memory/repository/sqlite_vec_index.py b/src/basic_memory/repository/sqlite_vec_index.py index cacc2a195..1df56f550 100644 --- a/src/basic_memory/repository/sqlite_vec_index.py +++ b/src/basic_memory/repository/sqlite_vec_index.py @@ -102,18 +102,21 @@ async def initialize(self) -> None: storage_missing = not vector_sql expected_dimensions = f"float[{self.scope.dimensions}]" dimensions_changed = bool(vector_sql and expected_dimensions not in vector_sql) - if dimensions_changed: + source_hash_missing = bool(vector_sql and "+source_hash text" not in vector_sql) + if dimensions_changed or source_hash_missing: logger.warning( - "Embedding dimension mismatch (expected {dimensions}); " - "recreating sqlite-vec storage", + "SQLite vector storage schema mismatch " + "(expected dimensions={dimensions}, " + "source_hash_missing={source_hash_missing}); recreating storage", dimensions=self.scope.dimensions, + source_hash_missing=source_hash_missing, ) await session.execute(text("DROP TABLE IF EXISTS search_vector_embeddings")) await session.execute(create_sqlite_search_vector_embeddings(self.scope.dimensions)) # Missing or dimension-rebuilt vec storage has no vectors, so ready # manifests must become pending before incremental sync inspects them. - if storage_missing or dimensions_changed: + if storage_missing or dimensions_changed or source_hash_missing: await session.execute( text( "UPDATE search_vector_chunks SET embedding_status = 'pending' " @@ -159,12 +162,45 @@ async def upsert(self, records: Sequence[VectorRecord]) -> None: await self.initialize() async with db.scoped_session(self._session_maker) as session: await self._ensure_loaded(session) - rowids_by_key = await self._rowids_by_key(session, [record.key for record in records]) - missing = [record.key for record in records if record.key not in rowids_by_key] - if missing: - raise RuntimeError(f"Vector manifest rows are missing for keys: {missing!r}") + params: dict[str, object] = {"project_id": self.scope.project_id} + predicates: list[str] = [] + records_by_key = {record.key: record for record in records} + for index, record in enumerate(records): + params[f"entity_id_{index}"] = record.key.entity_id + params[f"chunk_key_{index}"] = record.key.chunk_key + params[f"source_hash_{index}"] = record.source_hash + predicates.append( + f"(entity_id = :entity_id_{index} " + f"AND chunk_key = :chunk_key_{index} " + f"AND source_hash = :source_hash_{index})" + ) + + # SQLite has no SELECT FOR UPDATE. This conditional no-op write + # acquires the database write lock and verifies the source generation + # before vec0 rows can be replaced in the same transaction. + result = await session.execute( + text( + "UPDATE search_vector_chunks SET source_hash = source_hash " + "WHERE project_id = :project_id AND (" + + " OR ".join(predicates) + + ") RETURNING id, entity_id, chunk_key" + ), + params, + ) + rowids_by_key = { + VectorKey( + entity_id=int(row["entity_id"]), + chunk_key=str(row["chunk_key"]), + ): int(row["id"]) + for row in result.mappings().all() + } + current_records = [ + records_by_key[key] for key in rowids_by_key if key in records_by_key + ] + if not current_records: + return - rowids = [rowids_by_key[record.key] for record in records] + rowids = [rowids_by_key[record.key] for record in current_records] params = {f"rowid_{index}": rowid for index, rowid in enumerate(rowids)} placeholders = ", ".join(f":rowid_{index}" for index in range(len(rowids))) await session.execute( @@ -173,15 +209,16 @@ async def upsert(self, records: Sequence[VectorRecord]) -> None: ) await session.execute( text( - "INSERT INTO search_vector_embeddings (rowid, embedding) " - "VALUES (:rowid, :embedding)" + "INSERT INTO search_vector_embeddings (rowid, embedding, source_hash) " + "VALUES (:rowid, :embedding, :source_hash)" ), [ { "rowid": rowids_by_key[record.key], "embedding": json.dumps(record.values), + "source_hash": record.source_hash, } - for record in records + for record in current_records ], ) await session.commit() @@ -249,6 +286,8 @@ async def delete_orphans(self, _live_keys: Sequence[VectorKey]) -> None: "WHERE project_id = :project_id AND NOT (" "vector_index = 'sqlite-vec' " "AND embedding_model = :embedding_identity " + "AND search_vector_embeddings.source_hash = " + "search_vector_chunks.source_hash " "AND embedding_status = 'ready'))" ), { @@ -273,13 +312,14 @@ async def search( await self._ensure_loaded(session) result = await session.execute( text( - "WITH vector_matches AS (" - " SELECT rowid, distance FROM search_vector_embeddings " + "WITH vector_matches AS MATERIALIZED (" + " SELECT rowid, distance, source_hash FROM search_vector_embeddings " " WHERE 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 " "WHERE c.project_id = :project_id " "AND c.vector_index = 'sqlite-vec' " "AND c.embedding_status = 'ready' " diff --git a/test-int/test_embedding_status_vec0.py b/test-int/test_embedding_status_vec0.py index 510fe5809..5c6b9bfc8 100644 --- a/test-int/test_embedding_status_vec0.py +++ b/test-int/test_embedding_status_vec0.py @@ -134,10 +134,26 @@ async def test_embedding_status_reads_real_vec0_table(engine_factory, test_proje chunk_id = chunk_result.scalar_one() await session.commit() + # An obsolete embedding result must not claim the stable vec0 row after the + # manifest has advanced to a newer source generation. await search_repo._semantic_vector_index.upsert( [ VectorRecord( key=VectorKey(entity_id=entity_id, chunk_key="chunk-1"), + source_hash="stale-hash", + values=tuple(_unit_vector(dimensions)), + ) + ] + ) + async with db.scoped_session(session_maker) as session: + stale_count = await session.execute(text("SELECT COUNT(*) FROM search_vector_embeddings")) + assert stale_count.scalar_one() == 0 + + await search_repo._semantic_vector_index.upsert( + [ + VectorRecord( + key=VectorKey(entity_id=entity_id, chunk_key="chunk-1"), + source_hash="hash", values=tuple(_unit_vector(dimensions)), ) ] diff --git a/tests/repository/test_pgvector_index.py b/tests/repository/test_pgvector_index.py index e7abd0582..5dd6c5529 100644 --- a/tests/repository/test_pgvector_index.py +++ b/tests/repository/test_pgvector_index.py @@ -53,12 +53,14 @@ def __init__( *, table_exists: bool = False, dimensions: int | None = None, + has_source_hash: bool = True, chunk_rows: list[dict[str, object]] | None = None, search_rows: list[dict[str, object]] | None = None, fail_extension: bool = False, ) -> None: self.table_exists = table_exists self.dimensions = dimensions + self.has_source_hash = has_source_hash self.chunk_rows = chunk_rows or [] self.search_rows = search_rows or [] self.fail_extension = fail_extension @@ -76,6 +78,8 @@ async def execute( raise RuntimeError("extension unavailable") if "information_schema.tables" in sql: return FakeResult(fetchone=(1,) if self.table_exists else None) + if "attname = 'source_hash'" in sql: + return FakeResult(scalar=1 if self.has_source_hash else None) if "SELECT atttypmod" in sql: return FakeResult(scalar=self.dimensions) if "SELECT id, entity_id, chunk_key" in sql: @@ -154,6 +158,20 @@ async def test_initialize_rebuilds_dimension_mismatch_and_invalidates_manifest(m assert any("embedding_status = 'pending'" in sql for sql in sql_calls) +@pytest.mark.asyncio +async def test_initialize_rebuilds_storage_without_source_generation(monkeypatch) -> None: + session = FakeSession(table_exists=True, dimensions=4, has_source_hash=False) + _install_session(monkeypatch, session) + index = PgVectorIndex(MagicMock(), _scope()) + + await index.initialize() + + sql_calls = _sql_calls(session) + assert any("DROP TABLE IF EXISTS search_vector_embeddings" in sql for sql in sql_calls) + assert any("source_hash TEXT NOT NULL" in sql for sql in sql_calls) + assert any("embedding_status = 'pending'" in sql for sql in sql_calls) + + @pytest.mark.asyncio async def test_initialize_reports_missing_pgvector_extension(monkeypatch) -> None: session = FakeSession(fail_extension=True) @@ -170,8 +188,18 @@ async def test_upsert_resolves_stable_keys_and_writes_one_batch(monkeypatch) -> key_b = VectorKey(entity_id=11, chunk_key="entity:11:1") session = FakeSession( chunk_rows=[ - {"id": 101, "entity_id": 11, "chunk_key": key_a.chunk_key}, - {"id": 102, "entity_id": 11, "chunk_key": key_b.chunk_key}, + { + "id": 101, + "entity_id": 11, + "chunk_key": key_a.chunk_key, + "source_hash": "hash-a", + }, + { + "id": 102, + "entity_id": 11, + "chunk_key": key_b.chunk_key, + "source_hash": "hash-b", + }, ] ) _install_session(monkeypatch, session) @@ -180,8 +208,8 @@ async def test_upsert_resolves_stable_keys_and_writes_one_batch(monkeypatch) -> await index.upsert( [ - VectorRecord(key=key_a, values=(1.0, 0.0, 0.0, 0.0)), - VectorRecord(key=key_b, values=(0.0, 1.0, 0.0, 0.0)), + VectorRecord(key=key_a, source_hash="hash-a", values=(1.0, 0.0, 0.0, 0.0)), + VectorRecord(key=key_b, source_hash="hash-b", values=(0.0, 1.0, 0.0, 0.0)), ] ) @@ -191,13 +219,40 @@ async def test_upsert_resolves_stable_keys_and_writes_one_batch(monkeypatch) -> "chunk_id_0": 101, "embedding_0": "[1,0,0,0]", "dimensions_0": 4, + "source_hash_0": "hash-a", "chunk_id_1": 102, "embedding_1": "[0,1,0,0]", "dimensions_1": 4, + "source_hash_1": "hash-b", } assert session.commit_count == 1 +@pytest.mark.asyncio +async def test_upsert_skips_stale_source_generation(monkeypatch) -> None: + key = VectorKey(entity_id=11, chunk_key="entity:11:0") + session = FakeSession( + chunk_rows=[ + { + "id": 101, + "entity_id": 11, + "chunk_key": key.chunk_key, + "source_hash": "new-hash", + } + ] + ) + _install_session(monkeypatch, session) + index = PgVectorIndex(MagicMock(), _scope()) + index._initialized = True + + await index.upsert([VectorRecord(key=key, source_hash="old-hash", values=(1.0, 0.0, 0.0, 0.0))]) + + lock_call = next(call for call in session.calls if "SELECT id, entity_id" in call[0]) + assert "FOR UPDATE" in lock_call[0] + assert not any("INSERT INTO" in sql for sql, _params in session.calls) + assert session.commit_count == 0 + + @pytest.mark.asyncio async def test_upsert_rejects_missing_manifest_key(monkeypatch) -> None: key = VectorKey(entity_id=12, chunk_key="entity:12:0") @@ -207,7 +262,7 @@ async def test_upsert_rejects_missing_manifest_key(monkeypatch) -> None: index._initialized = True with pytest.raises(RuntimeError, match="manifest rows are missing"): - await index.upsert([VectorRecord(key=key, values=(1.0, 0.0, 0.0, 0.0))]) + await index.upsert([VectorRecord(key=key, source_hash="hash", values=(1.0, 0.0, 0.0, 0.0))]) @pytest.mark.asyncio diff --git a/tests/repository/test_semantic_search_base.py b/tests/repository/test_semantic_search_base.py index 67821fee7..2d33454e5 100644 --- a/tests/repository/test_semantic_search_base.py +++ b/tests/repository/test_semantic_search_base.py @@ -243,6 +243,107 @@ async def fake_scoped_session(_session_maker): assert ready_params["source_hash_0"] == source_hash +@pytest.mark.asyncio +async def test_external_postgres_upsert_holds_manifest_lock_through_ready_commit( + monkeypatch, +) -> None: + """External writes and ready publication share one manifest-lock transaction.""" + repo = _ConcreteRepo() + repo._semantic_vector_index_name = "milvus" + repo._embedding_provider = SimpleNamespace(model_name="stub", dimensions=4) + adapter = _RecordingVectorIndex() + repo._semantic_vector_index = adapter + chunk_text = "current chunk text" + source_hash = hashlib.sha256(chunk_text.encode("utf-8")).hexdigest() + session = AsyncMock() + session.connection.return_value = SimpleNamespace(dialect=SimpleNamespace(name="postgresql")) + session.execute.side_effect = [ + SimpleNamespace( + mappings=lambda: SimpleNamespace( + all=lambda: [ + { + "id": 7, + "entity_id": 41, + "chunk_key": "entity:41:0", + "source_hash": source_hash, + } + ] + ) + ), + SimpleNamespace(), + ] + context_count = 0 + + @asynccontextmanager + async def fake_scoped_session(_session_maker): + nonlocal context_count + context_count += 1 + yield session + + monkeypatch.setattr(search_repository_base_module.db, "scoped_session", fake_scoped_session) + + await repo._persist_embeddings([(7, chunk_text)], [[1.0, 0.0, 0.0, 0.0]]) + + lock_statement = session.execute.await_args_list[0].args[0] + ready_statement = session.execute.await_args_list[1].args[0] + assert "FOR UPDATE" in str(lock_statement) + assert "embedding_status = 'ready'" in str(ready_statement) + assert context_count == 1 + assert session.commit.await_count == 1 + assert adapter.upserted_records[0].source_hash == source_hash + + +@pytest.mark.asyncio +async def test_external_sqlite_upsert_holds_write_lock_through_ready_commit( + monkeypatch, +) -> None: + """SQLite external writes use one write-lock transaction through publication.""" + repo = _ConcreteRepo() + repo._semantic_vector_index_name = "milvus" + repo._embedding_provider = SimpleNamespace(model_name="stub", dimensions=4) + adapter = _RecordingVectorIndex() + repo._semantic_vector_index = adapter + chunk_text = "current chunk text" + source_hash = hashlib.sha256(chunk_text.encode("utf-8")).hexdigest() + session = AsyncMock() + session.connection.return_value = SimpleNamespace(dialect=SimpleNamespace(name="sqlite")) + session.execute.side_effect = [ + SimpleNamespace( + mappings=lambda: SimpleNamespace( + all=lambda: [ + { + "id": 7, + "entity_id": 41, + "chunk_key": "entity:41:0", + "source_hash": source_hash, + } + ] + ) + ), + SimpleNamespace(), + ] + context_count = 0 + + @asynccontextmanager + async def fake_scoped_session(_session_maker): + nonlocal context_count + context_count += 1 + yield session + + monkeypatch.setattr(search_repository_base_module.db, "scoped_session", fake_scoped_session) + + await repo._persist_embeddings([(7, chunk_text)], [[1.0, 0.0, 0.0, 0.0]]) + + lock_statement = session.execute.await_args_list[0].args[0] + ready_statement = session.execute.await_args_list[1].args[0] + assert "UPDATE search_vector_chunks SET source_hash = source_hash" in str(lock_statement) + assert "RETURNING id, entity_id, chunk_key, source_hash" in str(lock_statement) + assert "embedding_status = 'ready'" in str(ready_statement) + assert context_count == 1 + assert session.commit.await_count == 1 + assert adapter.upserted_records[0].source_hash == source_hash + + # --- SQLite SemanticSearchDisabledError --- diff --git a/tests/repository/test_semantic_vector_index.py b/tests/repository/test_semantic_vector_index.py index 8a2370f5d..3b852df23 100644 --- a/tests/repository/test_semantic_vector_index.py +++ b/tests/repository/test_semantic_vector_index.py @@ -103,14 +103,17 @@ def test_vector_contract_values_and_dimension_validation() -> None: dimensions=3, ) key = VectorKey(entity_id=11, chunk_key="entity:11:0") - record = VectorRecord(key=key, values=(1.0, 0.0, 0.0)) + record = VectorRecord(key=key, source_hash="hash", values=(1.0, 0.0, 0.0)) assert isinstance(StubVectorIndex(scope), SemanticVectorIndex) validate_vector_dimensions(scope, [record]) validate_query_dimensions(scope, [1.0, 0.0, 0.0]) with pytest.raises(ValueError, match="expected 3, got 2"): - validate_vector_dimensions(scope, [VectorRecord(key=key, values=(1.0, 0.0))]) + validate_vector_dimensions( + scope, + [VectorRecord(key=key, source_hash="hash", values=(1.0, 0.0))], + ) with pytest.raises(ValueError, match="expected 3, got 1"): validate_query_dimensions(scope, [1.0]) diff --git a/tests/repository/test_sqlite_vector_search_repository.py b/tests/repository/test_sqlite_vector_search_repository.py index 99161348c..99243585b 100644 --- a/tests/repository/test_sqlite_vector_search_repository.py +++ b/tests/repository/test_sqlite_vector_search_repository.py @@ -582,7 +582,7 @@ async def fail_ready_commit(session_maker): nonlocal context_count context_count += 1 async with original_scoped_session(session_maker) as session: - if context_count == 2: + if context_count == 1: async def fail_commit() -> None: raise RuntimeError("ready commit failed") From 459ee6fc8673ff1c126cb8daa9808c0f2a65fb5a Mon Sep 17 00:00:00 2001 From: phernandez Date: Sun, 26 Jul 2026 14:54:17 -0500 Subject: [PATCH 19/24] fix(core): bind vector deletes to manifest generations Signed-off-by: phernandez --- docs/semantic-search.md | 14 +- src/basic_memory/repository/pgvector_index.py | 33 ++- .../repository/postgres_search_repository.py | 23 +- .../repository/search_repository_base.py | 215 ++++++++++++------ .../repository/semantic_vector_index.py | 12 +- .../repository/semantic_vector_sync.py | 51 ++++- .../repository/sqlite_search_repository.py | 23 +- .../repository/sqlite_vec_index.py | 32 ++- tests/repository/test_hybrid_fusion.py | 15 +- tests/repository/test_pgvector_index.py | 14 +- .../test_postgres_search_repository_unit.py | 28 +-- tests/repository/test_semantic_search_base.py | 31 ++- .../repository/test_semantic_vector_index.py | 3 +- tests/repository/test_semantic_vector_sync.py | 36 ++- .../test_sqlite_vector_search_repository.py | 157 ++++++++++++- tests/repository/test_vector_pagination.py | 15 +- tests/repository/test_vector_threshold.py | 15 +- 17 files changed, 575 insertions(+), 142 deletions(-) diff --git a/docs/semantic-search.md b/docs/semantic-search.md index 022498255..e4a69e6d0 100644 --- a/docs/semantic-search.md +++ b/docs/semantic-search.md @@ -585,16 +585,20 @@ The returned `SemanticVectorIndex` has five asynchronous operations: - `initialize()` validates or creates backend storage. - `upsert(records)` idempotently writes vectors by `(entity_id, chunk_key)` for each record's `source_hash` generation. -- `delete(keys)` removes stable keys; missing keys are successful no-ops. +- `delete(records)` removes stable keys only for each record's `source_hash` generation; stale or + missing records are successful no-ops. - `delete_entity(entity_id)` removes all vectors for one entity in the scope. - `search(query, limit)` returns stable keys with normalized cosine similarity in `[0, 1]`. The adapter never receives a SQLAlchemy session and never calls the embedding provider. Basic Memory owns chunking and embedding, while the extension owns vector storage and lookup. -Each `VectorRecord` carries the SHA-256 source generation that produced its values. Basic Memory -holds the matching SQL manifest locked across extension adapter I/O and the ready-state transition -(`FOR UPDATE` on Postgres and a conditional write lock on SQLite), so an older overlapping sync -cannot overwrite a newer generation under the same stable adapter key. +The built-in pgvector and sqlite-vec adapters additionally remove each pending SQL manifest row in +the same database transaction as its vector so no newer generation can enter between those steps. +Each `VectorRecord` and `VectorDeletion` carries the SHA-256 source generation that produced its +value. Basic Memory holds the matching SQL manifest locked across extension adapter I/O and the +ready-state transition or deletion (`FOR UPDATE` on Postgres and a conditional write lock on +SQLite), so an older overlapping sync cannot overwrite or remove a newer generation under the same +stable adapter key. Adapters may also implement the separate `SemanticVectorIndexReconciler` capability. After a vector reindex, Basic Memory passes it the complete set of current ready keys so the adapter can diff --git a/src/basic_memory/repository/pgvector_index.py b/src/basic_memory/repository/pgvector_index.py index f68be1ba5..ed15bb967 100644 --- a/src/basic_memory/repository/pgvector_index.py +++ b/src/basic_memory/repository/pgvector_index.py @@ -13,6 +13,7 @@ from basic_memory.repository.semantic_errors import SemanticDependenciesMissingError from basic_memory.repository.semantic_vector_index import ( SemanticVectorIndex, + VectorDeletion, VectorIndexScope, VectorKey, VectorMatch, @@ -252,13 +253,33 @@ async def upsert(self, records: Sequence[VectorRecord]) -> None: ) await session.commit() - async def delete(self, keys: Sequence[VectorKey]) -> None: - if not keys: + async def delete(self, records: Sequence[VectorDeletion]) -> None: + if not records: return await self.initialize() async with db.scoped_session(self._session_maker) as session: - ids_by_key = await self._chunk_ids_by_key(session, keys) - chunk_ids = list(ids_by_key.values()) + params: dict[str, object] = {"project_id": self.scope.project_id} + predicates: list[str] = [] + for index, record in enumerate(records): + params[f"entity_id_{index}"] = record.key.entity_id + params[f"chunk_key_{index}"] = record.key.chunk_key + params[f"source_hash_{index}"] = record.source_hash + predicates.append( + f"(entity_id = :entity_id_{index} " + f"AND chunk_key = :chunk_key_{index} " + f"AND source_hash = :source_hash_{index} " + "AND embedding_status = 'pending')" + ) + result = await session.execute( + text( + "SELECT id, entity_id, chunk_key FROM search_vector_chunks " + "WHERE project_id = :project_id AND (" + + " OR ".join(predicates) + + ") FOR UPDATE" + ), + params, + ) + chunk_ids = [int(row["id"]) for row in result.mappings().all()] if chunk_ids: params = {f"chunk_id_{index}": value for index, value in enumerate(chunk_ids)} placeholders = ", ".join(f":chunk_id_{index}" for index in range(len(chunk_ids))) @@ -268,6 +289,10 @@ async def delete(self, keys: Sequence[VectorKey]) -> None: ), params, ) + await session.execute( + text(f"DELETE FROM search_vector_chunks WHERE id IN ({placeholders})"), + params, + ) await session.commit() async def delete_entity(self, entity_id: int) -> None: diff --git a/src/basic_memory/repository/postgres_search_repository.py b/src/basic_memory/repository/postgres_search_repository.py index 99e93538b..231d94459 100644 --- a/src/basic_memory/repository/postgres_search_repository.py +++ b/src/basic_memory/repository/postgres_search_repository.py @@ -3,6 +3,7 @@ import asyncio import json import re +from collections.abc import Sequence from datetime import datetime from typing import List, Optional @@ -27,6 +28,7 @@ from basic_memory.repository.metadata_filters import parse_metadata_filters from basic_memory.repository.semantic_errors import SemanticDependenciesMissingError from basic_memory.repository.semantic_vector_index import SemanticVectorIndex +from basic_memory.repository.semantic_vector_sync import StagedVectorDeletion from basic_memory.repository.semantic_vector_index_factory import ( build_vector_index_scope, resolve_semantic_vector_index_name, @@ -474,16 +476,29 @@ async def _delete_entity_chunks( self, session: AsyncSession, entity_id: int, - ) -> None: - await super()._delete_entity_chunks(session, entity_id) + *, + expected_deletions: Sequence[StagedVectorDeletion] | None = None, + ) -> list[StagedVectorDeletion]: + return await super()._delete_entity_chunks( + session, + entity_id, + expected_deletions=expected_deletions, + ) async def _delete_stale_chunks( self, session: AsyncSession, stale_ids: list[int], entity_id: int, - ) -> None: - await super()._delete_stale_chunks(session, stale_ids, entity_id) + *, + expected_deletions: Sequence[StagedVectorDeletion] | None = None, + ) -> list[StagedVectorDeletion]: + return await super()._delete_stale_chunks( + session, + stale_ids, + entity_id, + expected_deletions=expected_deletions, + ) def _distance_to_similarity(self, distance: float) -> float: """Convert pgvector cosine distance to cosine similarity. diff --git a/src/basic_memory/repository/search_repository_base.py b/src/basic_memory/repository/search_repository_base.py index 738628dac..5294cdad4 100644 --- a/src/basic_memory/repository/search_repository_base.py +++ b/src/basic_memory/repository/search_repository_base.py @@ -42,6 +42,7 @@ from basic_memory.repository.semantic_vector_index import ( SemanticVectorIndex, SemanticVectorIndexReconciler, + VectorDeletion, VectorKey, VectorMatch, VectorRecord, @@ -51,6 +52,7 @@ EntityVectorShardPlan as _EntityVectorShardPlan, PendingEmbeddingJob as _PendingEmbeddingJob, PreparedEntityVectorSync as _PreparedEntityVectorSync, + StagedVectorDeletion as _StagedVectorDeletion, VectorChunkState, VectorSyncBatchResult, ) @@ -450,22 +452,14 @@ async def _delete_entity_chunks( self, session: AsyncSession, entity_id: int, - ) -> None: + *, + expected_deletions: Sequence[_StagedVectorDeletion] | None = None, + ) -> list[_StagedVectorDeletion]: """Stage an entity deletion by making its manifest rows non-searchable.""" - vector_index_result = await session.execute( - text( - "SELECT DISTINCT vector_index FROM search_vector_chunks " - "WHERE project_id = :project_id AND entity_id = :entity_id" - ), - {"project_id": self.project_id, "entity_id": entity_id}, - ) - self._assert_manifest_vector_ownership(vector_index_result.scalars().all()) - await session.execute( - text( - "UPDATE search_vector_chunks SET embedding_status = 'pending' " - "WHERE project_id = :project_id AND entity_id = :entity_id" - ), - {"project_id": self.project_id, "entity_id": entity_id}, + return await self._stage_vector_deletions( + session, + entity_id=entity_id, + expected_deletions=expected_deletions, ) async def _delete_stale_chunks( @@ -473,29 +467,76 @@ async def _delete_stale_chunks( session: AsyncSession, stale_ids: list[int], entity_id: int, - ) -> None: + *, + expected_deletions: Sequence[_StagedVectorDeletion] | None = None, + ) -> list[_StagedVectorDeletion]: """Stage stale chunk deletion by making manifest rows non-searchable.""" if not stale_ids: - return - params = {f"stale_id_{index}": row_id for index, row_id in enumerate(stale_ids)} - placeholders = ", ".join(f":stale_id_{index}" for index in range(len(stale_ids))) - vector_index_result = await session.execute( - text( - "SELECT DISTINCT vector_index FROM search_vector_chunks " - f"WHERE project_id = :project_id AND entity_id = :entity_id " - f"AND id IN ({placeholders})" - ), - {**params, "project_id": self.project_id, "entity_id": entity_id}, + return [] + return await self._stage_vector_deletions( + session, + entity_id=entity_id, + row_ids=stale_ids, + expected_deletions=expected_deletions, ) - self._assert_manifest_vector_ownership(vector_index_result.scalars().all()) - await session.execute( + + async def _stage_vector_deletions( + self, + session: AsyncSession, + *, + entity_id: int, + row_ids: Sequence[int] | None = None, + expected_deletions: Sequence[_StagedVectorDeletion] | None = None, + ) -> list[_StagedVectorDeletion]: + """Durably stage and return the exact manifest generations to delete.""" + params: dict[str, object] = { + "project_id": self.project_id, + "entity_id": entity_id, + } + generation_clause = "" + if expected_deletions is not None: + if not expected_deletions: + return [] + predicates: list[str] = [] + for index, deletion in enumerate(expected_deletions): + params[f"row_id_{index}"] = deletion.row_id + params[f"source_hash_{index}"] = deletion.source_hash + params[f"vector_index_{index}"] = deletion.vector_index + predicates.append( + f"(id = :row_id_{index} " + f"AND source_hash = :source_hash_{index} " + f"AND vector_index = :vector_index_{index})" + ) + generation_clause = " AND (" + " OR ".join(predicates) + ")" + elif row_ids is not None: + if not row_ids: + return [] + placeholders: list[str] = [] + for index, row_id in enumerate(row_ids): + params[f"row_id_{index}"] = row_id + placeholders.append(f":row_id_{index}") + generation_clause = " AND id IN (" + ", ".join(placeholders) + ")" + + result = await session.execute( text( "UPDATE search_vector_chunks SET embedding_status = 'pending' " - f"WHERE project_id = :project_id AND entity_id = :entity_id " - f"AND id IN ({placeholders})" + "WHERE project_id = :project_id AND entity_id = :entity_id" + + generation_clause + + " RETURNING id, chunk_key, source_hash, vector_index" ), - {**params, "project_id": self.project_id, "entity_id": entity_id}, + params, ) + staged = [ + _StagedVectorDeletion( + row_id=int(row["id"]), + chunk_key=str(row["chunk_key"]), + source_hash=str(row["source_hash"]), + vector_index=str(row["vector_index"]), + ) + for row in result.mappings().all() + ] + self._assert_manifest_vector_ownership(deletion.vector_index for deletion in staged) + return staged async def _finalize_prepared_vector_deletions( self, @@ -507,61 +548,96 @@ async def _finalize_prepared_vector_deletions( fails, those rows remain non-searchable and the next sync retries the idempotent delete instead of losing cleanup intent. """ - if not prepared.delete_entity_vectors and not prepared.stale_chunk_ids: + if not prepared.staged_deletions: return if not hasattr(self, "_semantic_vector_index"): return - if prepared.delete_entity_vectors: - async with db.scoped_session(self.session_maker) as session: - vector_index_result = await session.execute( - text( - "SELECT DISTINCT vector_index FROM search_vector_chunks " - "WHERE project_id = :project_id AND entity_id = :entity_id" - ), - {"project_id": self.project_id, "entity_id": prepared.entity_id}, - ) - self._assert_manifest_vector_ownership(vector_index_result.scalars().all()) - await self._semantic_vector_index.delete_entity(prepared.entity_id) + params: dict[str, object] = { + "project_id": self.project_id, + "entity_id": prepared.entity_id, + } + predicates: list[str] = [] + for index, deletion in enumerate(prepared.staged_deletions): + params[f"row_id_{index}"] = deletion.row_id + params[f"source_hash_{index}"] = deletion.source_hash + params[f"vector_index_{index}"] = deletion.vector_index + predicates.append( + f"(id = :row_id_{index} " + f"AND source_hash = :source_hash_{index} " + f"AND vector_index = :vector_index_{index})" + ) + generation_clause = " OR ".join(predicates) + deletions = [ + VectorDeletion( + key=VectorKey( + entity_id=prepared.entity_id, + chunk_key=deletion.chunk_key, + ), + source_hash=deletion.source_hash, + ) + for deletion in prepared.staged_deletions + ] + + external_vector_index = self._semantic_vector_index_name not in _BUILT_IN_VECTOR_INDEX_NAMES + if external_vector_index: async with db.scoped_session(self.session_maker) as session: + connection = await session.connection() + if connection.dialect.name == "sqlite": + lock_statement = ( + "UPDATE search_vector_chunks SET source_hash = source_hash " + "WHERE project_id = :project_id AND entity_id = :entity_id " + "AND embedding_status = 'pending' AND (" + + generation_clause + + ") RETURNING id, chunk_key, source_hash, vector_index" + ) + else: + lock_statement = ( + "SELECT id, chunk_key, source_hash, vector_index " + "FROM search_vector_chunks " + "WHERE project_id = :project_id AND entity_id = :entity_id " + "AND embedding_status = 'pending' AND (" + + generation_clause + + ") FOR UPDATE" + ) + result = await session.execute(text(lock_statement), params) + rows = result.mappings().all() + self._assert_manifest_vector_ownership(row["vector_index"] for row in rows) + current_generations = {(int(row["id"]), str(row["source_hash"])) for row in rows} + current_deletions = [ + deletion + for deletion, staged in zip( + deletions, + prepared.staged_deletions, + strict=True, + ) + if (staged.row_id, staged.source_hash) in current_generations + ] + if not current_deletions: + return + await self._semantic_vector_index.delete(current_deletions) await session.execute( text( "DELETE FROM search_vector_chunks " - "WHERE project_id = :project_id AND entity_id = :entity_id" + "WHERE project_id = :project_id AND entity_id = :entity_id " + "AND embedding_status = 'pending' AND (" + generation_clause + ")" ), - {"project_id": self.project_id, "entity_id": prepared.entity_id}, + params, ) await session.commit() return - row_ids = prepared.stale_chunk_ids - params = {f"row_id_{index}": row_id for index, row_id in enumerate(row_ids)} - placeholders = ", ".join(f":row_id_{index}" for index in range(len(row_ids))) - async with db.scoped_session(self.session_maker) as session: - result = await session.execute( - text( - "SELECT entity_id, chunk_key, vector_index FROM search_vector_chunks " - f"WHERE project_id = :project_id AND entity_id = :entity_id " - f"AND id IN ({placeholders})" - ), - {**params, "project_id": self.project_id, "entity_id": prepared.entity_id}, - ) - rows = result.mappings().all() - self._assert_manifest_vector_ownership(row["vector_index"] for row in rows) - keys = [ - VectorKey(entity_id=int(row["entity_id"]), chunk_key=str(row["chunk_key"])) - for row in rows - ] - - await self._semantic_vector_index.delete(keys) + await self._semantic_vector_index.delete(deletions) + if self._semantic_vector_index_name in _BUILT_IN_VECTOR_INDEX_NAMES: + return async with db.scoped_session(self.session_maker) as session: await session.execute( text( "DELETE FROM search_vector_chunks " - f"WHERE project_id = :project_id AND entity_id = :entity_id " - f"AND id IN ({placeholders})" + "WHERE project_id = :project_id AND entity_id = :entity_id " + "AND embedding_status = 'pending' AND (" + generation_clause + ")" ), - {**params, "project_id": self.project_id, "entity_id": prepared.entity_id}, + params, ) await session.commit() @@ -719,7 +795,7 @@ async def delete_entity_vector_rows(self, entity_id: int) -> None: await self._ensure_vector_tables() async with db.scoped_session(self.session_maker) as session: - await self._delete_entity_chunks(session, entity_id) + staged_deletions = await self._delete_entity_chunks(session, entity_id) await session.commit() await self._finalize_prepared_vector_deletions( _PreparedEntityVectorSync( @@ -728,6 +804,7 @@ async def delete_entity_vector_rows(self, entity_id: int) -> None: source_rows_count=0, embedding_jobs=[], delete_entity_vectors=True, + staged_deletions=staged_deletions, ) ) diff --git a/src/basic_memory/repository/semantic_vector_index.py b/src/basic_memory/repository/semantic_vector_index.py index 982b69e23..a68880344 100644 --- a/src/basic_memory/repository/semantic_vector_index.py +++ b/src/basic_memory/repository/semantic_vector_index.py @@ -42,6 +42,14 @@ class VectorRecord: values: tuple[float, ...] +@dataclass(frozen=True, slots=True) +class VectorDeletion: + """One source-generation vector value to remove idempotently.""" + + key: VectorKey + source_hash: str + + @dataclass(frozen=True, slots=True) class VectorMatch: """One nearest-neighbour match with normalized cosine similarity.""" @@ -69,8 +77,8 @@ async def upsert(self, records: Sequence[VectorRecord]) -> None: """Insert or replace vectors only for each record's source generation.""" ... - async def delete(self, keys: Sequence[VectorKey]) -> None: - """Delete vectors by stable key; missing keys are successful no-ops.""" + async def delete(self, records: Sequence[VectorDeletion]) -> None: + """Delete vectors only for each record's source generation.""" ... async def delete_entity(self, entity_id: int) -> None: diff --git a/src/basic_memory/repository/semantic_vector_sync.py b/src/basic_memory/repository/semantic_vector_sync.py index 13f927054..15e5fced1 100644 --- a/src/basic_memory/repository/semantic_vector_sync.py +++ b/src/basic_memory/repository/semantic_vector_sync.py @@ -65,6 +65,17 @@ class PreparedEntityVectorSync: queue_start: float | None = None delete_entity_vectors: bool = False stale_chunk_ids: list[int] = field(default_factory=list) + staged_deletions: list[StagedVectorDeletion] = field(default_factory=list) + + +@dataclass(frozen=True, slots=True) +class StagedVectorDeletion: + """Manifest generation durably staged for adapter deletion.""" + + row_id: int + chunk_key: str + source_hash: str + vector_index: str @dataclass(frozen=True, slots=True) @@ -88,6 +99,7 @@ class DeleteEntityVectorPreparePlan: sync_start: float prepare_start: float source_rows_count: int + expected_deletions: list[StagedVectorDeletion] @dataclass(frozen=True, slots=True) @@ -100,6 +112,7 @@ class UpsertEntityVectorPreparePlan: source_rows_count: int existing_by_key: dict[str, VectorChunkState] stale_ids: list[int] + stale_deletions: list[StagedVectorDeletion] metadata_update_ids: list[int] scheduled_records: list[VectorChunkRecord] entity_fingerprint: str @@ -800,6 +813,15 @@ def delete_entity_chunks() -> DeleteEntityVectorPreparePlan: sync_start=sync_start, prepare_start=prepare_start, source_rows_count=source_rows_count, + expected_deletions=[ + StagedVectorDeletion( + row_id=row.id, + chunk_key=row.chunk_key, + source_hash=row.source_hash, + vector_index=row.vector_index, + ) + for row in existing_rows + ], ) if not source_rows: @@ -818,6 +840,16 @@ def delete_entity_chunks() -> DeleteEntityVectorPreparePlan: stale_ids = [ row.id for chunk_key, row in existing_by_key.items() if chunk_key not in incoming_chunk_keys ] + stale_deletions = [ + StagedVectorDeletion( + row_id=row.id, + chunk_key=row.chunk_key, + source_hash=row.source_hash, + vector_index=row.vector_index, + ) + for chunk_key, row in existing_by_key.items() + if chunk_key not in incoming_chunk_keys + ] orphan_ids = {row.id for row in existing_rows if not row.has_embedding} # Trigger: all persisted chunk metadata already matches this entity's @@ -897,6 +929,7 @@ def delete_entity_chunks() -> DeleteEntityVectorPreparePlan: source_rows_count=source_rows_count, existing_by_key=existing_by_key, stale_ids=stale_ids, + stale_deletions=stale_deletions, metadata_update_ids=metadata_update_ids, scheduled_records=scheduled_records, entity_fingerprint=current_entity_fingerprint, @@ -914,7 +947,11 @@ async def apply_entity_vector_prepare_plan( ) -> PreparedEntityVectorSync: """Apply one planned entity mutation inside the caller-owned transaction.""" if isinstance(plan, DeleteEntityVectorPreparePlan): - await repository._delete_entity_chunks(session, plan.entity_id) + staged_deletions = await repository._delete_entity_chunks( + session, + plan.entity_id, + expected_deletions=plan.expected_deletions, + ) return PreparedEntityVectorSync( entity_id=plan.entity_id, sync_start=plan.sync_start, @@ -922,11 +959,18 @@ async def apply_entity_vector_prepare_plan( embedding_jobs=[], prepare_seconds=time.perf_counter() - plan.prepare_start, delete_entity_vectors=True, + staged_deletions=staged_deletions, ) timestamp_expr = repository._timestamp_now_expr() + staged_deletions: list[StagedVectorDeletion] = [] if plan.stale_ids: - await repository._delete_stale_chunks(session, plan.stale_ids, plan.entity_id) + staged_deletions = await repository._delete_stale_chunks( + session, + plan.stale_ids, + plan.entity_id, + expected_deletions=plan.stale_deletions, + ) for row_id in plan.metadata_update_ids: await session.execute( text( @@ -970,7 +1014,8 @@ async def apply_entity_vector_prepare_plan( remaining_jobs_after_shard=plan.shard_plan.remaining_jobs_after_shard, prepare_seconds=prepare_seconds, queue_start=time.perf_counter(), - stale_chunk_ids=plan.stale_ids, + stale_chunk_ids=[deletion.row_id for deletion in staged_deletions], + staged_deletions=staged_deletions, ) diff --git a/src/basic_memory/repository/sqlite_search_repository.py b/src/basic_memory/repository/sqlite_search_repository.py index f31b129cc..8965e5637 100644 --- a/src/basic_memory/repository/sqlite_search_repository.py +++ b/src/basic_memory/repository/sqlite_search_repository.py @@ -2,6 +2,7 @@ import asyncio import re +from collections.abc import Sequence from contextlib import asynccontextmanager from datetime import datetime from typing import List, Optional @@ -30,6 +31,7 @@ from basic_memory.repository.metadata_filters import parse_metadata_filters, build_sqlite_json_path from basic_memory.repository.semantic_errors import SemanticDependenciesMissingError from basic_memory.repository.semantic_vector_index import SemanticVectorIndex +from basic_memory.repository.semantic_vector_sync import StagedVectorDeletion from basic_memory.repository.semantic_vector_index_factory import build_vector_index_scope from basic_memory.repository.sqlite_vec_index import SQLiteVecIndex from basic_memory.schemas.search import SearchItemType, SearchRetrievalMode @@ -545,16 +547,29 @@ async def _delete_entity_chunks( self, session: AsyncSession, entity_id: int, - ) -> None: - await super()._delete_entity_chunks(session, entity_id) + *, + expected_deletions: Sequence[StagedVectorDeletion] | None = None, + ) -> list[StagedVectorDeletion]: + return await super()._delete_entity_chunks( + session, + entity_id, + expected_deletions=expected_deletions, + ) async def _delete_stale_chunks( self, session: AsyncSession, stale_ids: list[int], entity_id: int, - ) -> None: - await super()._delete_stale_chunks(session, stale_ids, entity_id) + *, + expected_deletions: Sequence[StagedVectorDeletion] | None = None, + ) -> list[StagedVectorDeletion]: + return await super()._delete_stale_chunks( + session, + stale_ids, + entity_id, + expected_deletions=expected_deletions, + ) async def _delete_project_builtin_vector_rows(self, session: AsyncSession) -> None: """Delete sqlite-vec rows atomically with their project manifest.""" diff --git a/src/basic_memory/repository/sqlite_vec_index.py b/src/basic_memory/repository/sqlite_vec_index.py index 1df56f550..c62e66f9d 100644 --- a/src/basic_memory/repository/sqlite_vec_index.py +++ b/src/basic_memory/repository/sqlite_vec_index.py @@ -16,6 +16,7 @@ from basic_memory.repository.semantic_errors import SemanticDependenciesMissingError from basic_memory.repository.semantic_vector_index import ( SemanticVectorIndex, + VectorDeletion, VectorIndexScope, VectorKey, VectorMatch, @@ -223,13 +224,34 @@ async def upsert(self, records: Sequence[VectorRecord]) -> None: ) await session.commit() - async def delete(self, keys: Sequence[VectorKey]) -> None: - if not keys: + async def delete(self, records: Sequence[VectorDeletion]) -> None: + if not records: return await self.initialize() async with db.scoped_session(self._session_maker) as session: await self._ensure_loaded(session) - rowids = list((await self._rowids_by_key(session, keys)).values()) + params: dict[str, object] = {"project_id": self.scope.project_id} + predicates: list[str] = [] + for index, record in enumerate(records): + params[f"entity_id_{index}"] = record.key.entity_id + params[f"chunk_key_{index}"] = record.key.chunk_key + params[f"source_hash_{index}"] = record.source_hash + predicates.append( + f"(entity_id = :entity_id_{index} " + f"AND chunk_key = :chunk_key_{index} " + f"AND source_hash = :source_hash_{index} " + "AND embedding_status = 'pending')" + ) + result = await session.execute( + text( + "UPDATE search_vector_chunks SET source_hash = source_hash " + "WHERE project_id = :project_id AND (" + + " OR ".join(predicates) + + ") RETURNING id" + ), + params, + ) + rowids = [int(row_id) for row_id in result.scalars().all()] if rowids: params = {f"rowid_{index}": rowid for index, rowid in enumerate(rowids)} placeholders = ", ".join(f":rowid_{index}" for index in range(len(rowids))) @@ -237,6 +259,10 @@ async def delete(self, keys: Sequence[VectorKey]) -> None: text(f"DELETE FROM search_vector_embeddings WHERE rowid IN ({placeholders})"), params, ) + await session.execute( + text(f"DELETE FROM search_vector_chunks WHERE id IN ({placeholders})"), + params, + ) await session.commit() async def delete_entity(self, entity_id: int) -> None: diff --git a/tests/repository/test_hybrid_fusion.py b/tests/repository/test_hybrid_fusion.py index e61d02e51..16a6a47ab 100644 --- a/tests/repository/test_hybrid_fusion.py +++ b/tests/repository/test_hybrid_fusion.py @@ -90,11 +90,18 @@ async def _run_vector_query(self, session, query_embedding, candidate_limit): async def _write_embeddings(self, session, jobs, embeddings): pass # pragma: no cover - async def _delete_entity_chunks(self, session, entity_id): - pass # pragma: no cover + async def _delete_entity_chunks(self, session, entity_id, *, expected_deletions=None): + return [] # pragma: no cover - async def _delete_stale_chunks(self, session, stale_ids, entity_id): - pass # pragma: no cover + async def _delete_stale_chunks( + self, + session, + stale_ids, + entity_id, + *, + expected_deletions=None, + ): + return [] # pragma: no cover async def _update_timestamp_sql(self): return "CURRENT_TIMESTAMP" # pragma: no cover diff --git a/tests/repository/test_pgvector_index.py b/tests/repository/test_pgvector_index.py index 5dd6c5529..e1fa8bfe8 100644 --- a/tests/repository/test_pgvector_index.py +++ b/tests/repository/test_pgvector_index.py @@ -12,6 +12,7 @@ from basic_memory.repository.pgvector_index import PgVectorIndex from basic_memory.repository.semantic_errors import SemanticDependenciesMissingError from basic_memory.repository.semantic_vector_index import ( + VectorDeletion, VectorIndexScope, VectorKey, VectorRecord, @@ -274,15 +275,20 @@ async def test_delete_stable_keys_and_entity(monkeypatch) -> None: index._initialized = True await index.delete([]) - await index.delete([key]) + await index.delete([VectorDeletion(key=key, source_hash="hash")]) await index.delete_entity(13) await index.delete_orphans([key]) + delete_lock = next(call for call in session.calls if "SELECT id, entity_id" in call[0]) + assert "source_hash = :source_hash_0" in delete_lock[0] + assert "embedding_status = 'pending'" in delete_lock[0] + assert "FOR UPDATE" in delete_lock[0] delete_calls = [call for call in session.calls if "DELETE FROM" in call[0]] - assert len(delete_calls) == 3 + assert len(delete_calls) == 4 assert delete_calls[0][1] == {"chunk_id_0": 103} - assert delete_calls[1][1] == {"project_id": 7, "entity_id": 13} - assert delete_calls[2][1] == { + assert delete_calls[1][1] == {"chunk_id_0": 103} + assert delete_calls[2][1] == {"project_id": 7, "entity_id": 13} + assert delete_calls[3][1] == { "project_id": 7, "embedding_identity": "stub:4", } diff --git a/tests/repository/test_postgres_search_repository_unit.py b/tests/repository/test_postgres_search_repository_unit.py index 8f4e79f76..9c9fdfaa3 100644 --- a/tests/repository/test_postgres_search_repository_unit.py +++ b/tests/repository/test_postgres_search_repository_unit.py @@ -246,18 +246,19 @@ class TestDeleteStaleChunks: async def test_delete_stale_chunks_builds_correct_params(self): repo = _make_repo() session = AsyncMock() - ownership_result = MagicMock() - ownership_result.scalars.return_value.all.return_value = [] - session.execute.side_effect = [ownership_result, MagicMock()] + stage_result = MagicMock() + stage_result.mappings.return_value.all.return_value = [] + session.execute.return_value = stage_result stale_ids = [10, 20, 30] await repo._delete_stale_chunks(session, stale_ids, entity_id=5) - assert session.execute.await_count == 2 - call_args = session.execute.await_args_list[1] + session.execute.assert_awaited_once() + call_args = session.execute.await_args + assert "RETURNING id, chunk_key, source_hash, vector_index" in str(call_args.args[0]) params = call_args[0][1] - assert params["stale_id_0"] == 10 - assert params["stale_id_1"] == 20 - assert params["stale_id_2"] == 30 + assert params["row_id_0"] == 10 + assert params["row_id_1"] == 20 + assert params["row_id_2"] == 30 assert params["project_id"] == repo.project_id assert params["entity_id"] == 5 @@ -272,12 +273,13 @@ class TestDeleteEntityChunks: async def test_delete_entity_chunks_executes_sql(self): repo = _make_repo() session = AsyncMock() - ownership_result = MagicMock() - ownership_result.scalars.return_value.all.return_value = [] - session.execute.side_effect = [ownership_result, MagicMock()] + stage_result = MagicMock() + stage_result.mappings.return_value.all.return_value = [] + session.execute.return_value = stage_result await repo._delete_entity_chunks(session, entity_id=42) - assert session.execute.await_count == 2 - call_args = session.execute.await_args_list[1] + session.execute.assert_awaited_once() + call_args = session.execute.await_args + assert "RETURNING id, chunk_key, source_hash, vector_index" in str(call_args.args[0]) params = call_args[0][1] assert params["project_id"] == repo.project_id assert params["entity_id"] == 42 diff --git a/tests/repository/test_semantic_search_base.py b/tests/repository/test_semantic_search_base.py index 2d33454e5..65942c777 100644 --- a/tests/repository/test_semantic_search_base.py +++ b/tests/repository/test_semantic_search_base.py @@ -24,6 +24,7 @@ SemanticVectorIndexExtensionError, ) from basic_memory.repository.semantic_vector_index import ( + VectorDeletion, VectorIndexScope, VectorKey, VectorMatch, @@ -84,11 +85,18 @@ async def _run_vector_query(self, session, query_embedding, candidate_limit): async def _write_embeddings(self, session, jobs, embeddings): pass - async def _delete_entity_chunks(self, session, entity_id): - pass + async def _delete_entity_chunks(self, session, entity_id, *, expected_deletions=None): + return [] - async def _delete_stale_chunks(self, session, stale_ids, entity_id): - pass + async def _delete_stale_chunks( + self, + session, + stale_ids, + entity_id, + *, + expected_deletions=None, + ): + return [] async def _update_timestamp_sql(self): return "CURRENT_TIMESTAMP" @@ -116,7 +124,7 @@ async def initialize(self) -> None: async def upsert(self, records: Sequence[VectorRecord]) -> None: self.upserted_records.extend(records) - async def delete(self, keys: Sequence[VectorKey]) -> None: + async def delete(self, records: Sequence[VectorDeletion]) -> None: return None async def delete_entity(self, entity_id: int) -> None: @@ -632,7 +640,16 @@ async def test_entity_vector_cleanup_preserves_mismatched_external_owner() -> No repo._semantic_vector_index_name = "pgvector" session = AsyncMock() session.execute.return_value = SimpleNamespace( - scalars=lambda: SimpleNamespace(all=lambda: ["milvus"]) + mappings=lambda: SimpleNamespace( + all=lambda: [ + { + "id": 7, + "chunk_key": "entity:41:0", + "source_hash": "hash", + "vector_index": "milvus", + } + ] + ) ) with pytest.raises(SemanticVectorIndexExtensionError, match="milvus"): @@ -1040,7 +1057,7 @@ async def _yielding_write_scope(): monkeypatch.setattr(repo, "_fetch_prepare_window_existing_rows", _stub_fetch_existing_rows) monkeypatch.setattr(repo, "_prepare_entity_write_scope", _yielding_write_scope) monkeypatch.setattr(repo, "_prepare_vector_session", AsyncMock()) - monkeypatch.setattr(repo, "_delete_entity_chunks", AsyncMock()) + monkeypatch.setattr(repo, "_delete_entity_chunks", AsyncMock(return_value=[])) monkeypatch.setattr( search_repository_base_module.time, "perf_counter", diff --git a/tests/repository/test_semantic_vector_index.py b/tests/repository/test_semantic_vector_index.py index 3b852df23..9ff46ade7 100644 --- a/tests/repository/test_semantic_vector_index.py +++ b/tests/repository/test_semantic_vector_index.py @@ -17,6 +17,7 @@ from basic_memory.repository.semantic_vector_index import ( SEMANTIC_VECTOR_INDEX_ENTRY_POINT_GROUP, SemanticVectorIndex, + VectorDeletion, VectorIndexScope, VectorKey, VectorMatch, @@ -59,7 +60,7 @@ async def initialize(self) -> None: async def upsert(self, records: Sequence[VectorRecord]) -> None: return None - async def delete(self, keys: Sequence[VectorKey]) -> None: + async def delete(self, records: Sequence[VectorDeletion]) -> None: return None async def delete_entity(self, entity_id: int) -> None: diff --git a/tests/repository/test_semantic_vector_sync.py b/tests/repository/test_semantic_vector_sync.py index 986361aaa..27ba63df6 100644 --- a/tests/repository/test_semantic_vector_sync.py +++ b/tests/repository/test_semantic_vector_sync.py @@ -64,11 +64,18 @@ async def _run_vector_query(self, session, query_embedding, candidate_limit): async def _write_embeddings(self, session, jobs, embeddings): pass - async def _delete_entity_chunks(self, session, entity_id): - pass + async def _delete_entity_chunks(self, session, entity_id, *, expected_deletions=None): + return [] - async def _delete_stale_chunks(self, session, stale_ids, entity_id): - pass + async def _delete_stale_chunks( + self, + session, + stale_ids, + entity_id, + *, + expected_deletions=None, + ): + return [] def _distance_to_similarity(self, distance: float) -> float: return 1.0 / (1.0 + max(distance, 0.0)) @@ -363,6 +370,7 @@ def _stub_plan(repository, *, entity_id, source_rows, existing_rows): sync_start=0.0, prepare_start=0.0, source_rows_count=0, + expected_deletions=[], ) monkeypatch.setattr(semantic_vector_sync.db, "scoped_session", scoped_session) @@ -430,7 +438,7 @@ async def write_scope(): monkeypatch.setattr(semantic_vector_sync.db, "scoped_session", scoped_session) repository = _TestRepository() - delete_entity_chunks = AsyncMock() + delete_entity_chunks = AsyncMock(return_value=[]) monkeypatch.setattr(repository, "_prepare_entity_write_scope", write_scope) monkeypatch.setattr(repository, "_prepare_vector_session", AsyncMock()) monkeypatch.setattr(repository, "_delete_entity_chunks", delete_entity_chunks) @@ -444,7 +452,7 @@ async def write_scope(): ) assert empty_chunks.embedding_jobs == [] - delete_entity_chunks.assert_awaited_once_with(session, 1) + delete_entity_chunks.assert_awaited_once_with(session, 1, expected_deletions=[]) record = { "chunk_key": "new", @@ -472,7 +480,7 @@ async def write_scope(): Mock(return_value="CURRENT_TIMESTAMP"), ) monkeypatch.setattr(repository, "_log_vector_shard_plan", Mock()) - delete_stale_chunks = AsyncMock() + delete_stale_chunks = AsyncMock(return_value=[]) monkeypatch.setattr(repository, "_delete_stale_chunks", delete_stale_chunks) monkeypatch.setattr( repository, @@ -487,7 +495,19 @@ async def write_scope(): existing_rows=[stale_row], ) - delete_stale_chunks.assert_awaited_once_with(session, [7], 1) + delete_stale_chunks.assert_awaited_once_with( + session, + [7], + 1, + expected_deletions=[ + semantic_vector_sync.StagedVectorDeletion( + row_id=7, + chunk_key="old", + source_hash="old-hash", + vector_index="", + ) + ], + ) @pytest.mark.asyncio diff --git a/tests/repository/test_sqlite_vector_search_repository.py b/tests/repository/test_sqlite_vector_search_repository.py index 99243585b..7875d22a6 100644 --- a/tests/repository/test_sqlite_vector_search_repository.py +++ b/tests/repository/test_sqlite_vector_search_repository.py @@ -20,11 +20,16 @@ from basic_memory.repository.search_index_row import SearchIndexRow from basic_memory.repository.semantic_errors import SemanticVectorIndexExtensionError from basic_memory.repository.semantic_vector_index import ( + VectorDeletion, VectorIndexScope, VectorKey, VectorMatch, VectorRecord, ) +from basic_memory.repository.semantic_vector_sync import ( + PreparedEntityVectorSync, + StagedVectorDeletion, +) from basic_memory.repository.sqlite_search_repository import SQLiteSearchRepository from basic_memory.repository import sqlite_vec_index as sqlite_vec_index_module from basic_memory.repository.sqlite_vec_index import SQLITE_VEC_MAX_K, SQLiteVecIndex @@ -91,9 +96,12 @@ async def upsert(self, records: Sequence[VectorRecord]) -> None: raise RuntimeError("adapter write failed") self.records.update({record.key: record.values for record in records}) - async def delete(self, keys: Sequence[VectorKey]) -> None: - for key in keys: - self.records.pop(key, None) + async def delete(self, records: Sequence[VectorDeletion]) -> None: + self.deleted_entities.extend(sorted({record.key.entity_id for record in records})) + if self.fail_delete_entity: + raise RuntimeError("adapter delete failed") + for record in records: + self.records.pop(record.key, None) async def delete_entity(self, entity_id: int) -> None: self.deleted_entities.append(entity_id) @@ -411,6 +419,69 @@ async def test_sqlite_vec_reconciliation_is_project_scoped(search_repository): assert remaining.scalars().all() == [902, 903] +@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.""" + if not isinstance(search_repository, SQLiteSearchRepository): + pytest.skip("sqlite-vec deletion behavior is local SQLite-only.") + + _enable_semantic(search_repository) + await search_repository.init_search_index() + index = cast(SQLiteVecIndex, search_repository._semantic_vector_index) + key = VectorKey(entity_id=907, chunk_key="entity:907:0") + + 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 (" + "907, 907, :project_id, :chunk_key, 'text', 'hash', " + "'fingerprint', :embedding_model, 'sqlite-vec', 'ready')" + ), + { + "project_id": search_repository.project_id, + "chunk_key": key.chunk_key, + "embedding_model": search_repository._embedding_model_key(), + }, + ) + await session.execute( + text( + "INSERT INTO search_vector_embeddings (rowid, embedding, source_hash) " + "VALUES (907, :embedding, 'hash')" + ), + {"embedding": "[1.0, 0.0, 0.0, 0.0]"}, + ) + await session.commit() + + deletion = VectorDeletion(key=key, source_hash="hash") + await index.delete([deletion]) + async with db.scoped_session(search_repository.session_maker) as session: + assert ( + await session.scalar( + text("SELECT COUNT(*) FROM search_vector_embeddings WHERE rowid = 907") + ) + == 1 + ) + await session.execute( + text("UPDATE search_vector_chunks SET embedding_status = 'pending' WHERE id = 907") + ) + await session.commit() + + await index.delete([deletion]) + async with db.scoped_session(search_repository.session_maker) as session: + vector_count = await session.scalar( + text("SELECT COUNT(*) FROM search_vector_embeddings WHERE rowid = 907") + ) + manifest_count = await session.scalar( + text("SELECT COUNT(*) FROM search_vector_chunks WHERE id = 907") + ) + assert vector_count == 0 + assert manifest_count == 0 + + @pytest.mark.asyncio async def test_sqlite_chunk_upsert_and_delete_lifecycle(search_repository): """sync_entity_vectors updates changed chunks and clears vectors when source rows disappear.""" @@ -682,6 +753,86 @@ async def test_adapter_delete_failure_stays_pending_until_retry(search_repositor assert adapter.deleted_entities == [112, 112] +@pytest.mark.asyncio +@pytest.mark.parametrize("delete_entity_vectors", [False, True]) +async def test_staged_delete_preserves_newer_source_generation( + search_repository, + delete_entity_vectors, +): + """A stale finalizer must not delete a newer stable-key vector or manifest.""" + if not isinstance(search_repository, SQLiteSearchRepository): + pytest.skip("Semantic manifest concurrency is exercised through local SQLite.") + + _enable_semantic(search_repository) + adapter = RecordingVectorIndex() + search_repository._semantic_vector_index = adapter + search_repository._semantic_vector_index_name = "recording" + await search_repository.init_search_index() + await search_repository.index_item( + _entity_row( + project_id=search_repository.project_id, + row_id=116, + entity_id=116, + title="Generation Delete Race", + permalink="specs/generation-delete-race", + content_stems="old semantic generation", + ) + ) + await search_repository.sync_entity_vectors(116) + + async with db.scoped_session(search_repository.session_maker) as session: + manifest = await session.execute( + text( + "SELECT id, chunk_key, source_hash, vector_index " + "FROM search_vector_chunks " + "WHERE project_id = :project_id AND entity_id = :entity_id" + ), + {"project_id": search_repository.project_id, "entity_id": 116}, + ) + old_row = manifest.mappings().one() + await session.execute( + text( + "UPDATE search_vector_chunks " + "SET source_hash = 'new-source-hash', embedding_status = 'ready' " + "WHERE id = :row_id" + ), + {"row_id": old_row["id"]}, + ) + await session.commit() + + key = VectorKey(entity_id=116, chunk_key=str(old_row["chunk_key"])) + adapter.records[key] = (0.0, 0.0, 0.0, 1.0) + await search_repository._finalize_prepared_vector_deletions( + PreparedEntityVectorSync( + entity_id=116, + sync_start=0.0, + source_rows_count=0, + embedding_jobs=[], + delete_entity_vectors=delete_entity_vectors, + stale_chunk_ids=[] if delete_entity_vectors else [int(old_row["id"])], + staged_deletions=[ + StagedVectorDeletion( + row_id=int(old_row["id"]), + chunk_key=str(old_row["chunk_key"]), + source_hash=str(old_row["source_hash"]), + vector_index=str(old_row["vector_index"]), + ) + ], + ) + ) + + async with db.scoped_session(search_repository.session_maker) as session: + current = await session.execute( + text( + "SELECT source_hash, embedding_status FROM search_vector_chunks WHERE id = :row_id" + ), + {"row_id": old_row["id"]}, + ) + assert current.one() == ("new-source-hash", "ready") + assert adapter.records[key] == (0.0, 0.0, 0.0, 1.0) + assert adapter.deleted_entities == [] + + @pytest.mark.asyncio async def test_adapter_matches_hydrate_only_current_ready_manifest_rows(search_repository): """Stale external matches fail closed unless SQL says the current row is ready.""" diff --git a/tests/repository/test_vector_pagination.py b/tests/repository/test_vector_pagination.py index f08a98ba0..90686c47e 100644 --- a/tests/repository/test_vector_pagination.py +++ b/tests/repository/test_vector_pagination.py @@ -75,11 +75,18 @@ async def _run_vector_query(self, session, query_embedding, candidate_limit): async def _write_embeddings(self, session, jobs, embeddings): pass # pragma: no cover - async def _delete_entity_chunks(self, session, entity_id): - pass # pragma: no cover + async def _delete_entity_chunks(self, session, entity_id, *, expected_deletions=None): + return [] # pragma: no cover - async def _delete_stale_chunks(self, session, stale_ids, entity_id): - pass # pragma: no cover + async def _delete_stale_chunks( + self, + session, + stale_ids, + entity_id, + *, + expected_deletions=None, + ): + return [] # pragma: no cover async def _update_timestamp_sql(self): return "CURRENT_TIMESTAMP" # pragma: no cover diff --git a/tests/repository/test_vector_threshold.py b/tests/repository/test_vector_threshold.py index 915985182..ed8d17329 100644 --- a/tests/repository/test_vector_threshold.py +++ b/tests/repository/test_vector_threshold.py @@ -79,11 +79,18 @@ async def _run_vector_query(self, session, query_embedding, candidate_limit): async def _write_embeddings(self, session, jobs, embeddings): pass # pragma: no cover - async def _delete_entity_chunks(self, session, entity_id): - pass # pragma: no cover + async def _delete_entity_chunks(self, session, entity_id, *, expected_deletions=None): + return [] # pragma: no cover - async def _delete_stale_chunks(self, session, stale_ids, entity_id): - pass # pragma: no cover + async def _delete_stale_chunks( + self, + session, + stale_ids, + entity_id, + *, + expected_deletions=None, + ): + return [] # pragma: no cover async def _update_timestamp_sql(self): return "CURRENT_TIMESTAMP" # pragma: no cover From b6367daec0e1c608b46d93dcbe1dd6e3db7f28cc Mon Sep 17 00:00:00 2001 From: phernandez Date: Sun, 26 Jul 2026 15:11:08 -0500 Subject: [PATCH 20/24] fix(core): serialize project vector cleanup Signed-off-by: phernandez --- .../repository/search_repository_base.py | 82 +++++++++++++++++-- tests/repository/test_semantic_search_base.py | 76 ++++++++++++++--- 2 files changed, 142 insertions(+), 16 deletions(-) diff --git a/src/basic_memory/repository/search_repository_base.py b/src/basic_memory/repository/search_repository_base.py index 5294cdad4..b851b831f 100644 --- a/src/basic_memory/repository/search_repository_base.py +++ b/src/basic_memory/repository/search_repository_base.py @@ -339,6 +339,12 @@ async def _persist_embeddings( self._semantic_vector_index_name not in _BUILT_IN_VECTOR_INDEX_NAMES ) lock_external_write = external_vector_index and dialect_name in {"postgresql", "sqlite"} + if external_vector_index and dialect_name == "postgresql": + # Trigger: a project-wide cleanup can run beside a watcher flush. + # Why: row locks do not block a newly prepared chunk from being inserted. + # Outcome: every PostgreSQL external write shares the project-row lock + # held through cleanup, adapter deletion, and manifest removal. + await self._lock_external_vector_project(session, dialect_name=dialect_name) if external_vector_index and dialect_name == "sqlite": # SQLite has no SELECT FOR UPDATE. A conditional no-op write takes # the database write lock before extension I/O, serializing a newer @@ -867,6 +873,31 @@ async def delete_external_entity_vectors( async def _delete_project_builtin_vector_rows(self, session: AsyncSession) -> None: """Delete backend-owned vector rows before their SQL manifest is removed.""" + async def _lock_external_vector_project( + self, + session: AsyncSession, + *, + dialect_name: str, + ) -> None: + """Serialize project-wide cleanup with external adapter writes.""" + if dialect_name == "postgresql": + await session.execute( + text("SELECT id FROM project WHERE id = :project_id FOR UPDATE"), + {"project_id": self.project_id}, + ) + return + if dialect_name == "sqlite": + # SQLite has no row-level lock. This no-op write acquires its database + # write lock before ownership is read and keeps it through adapter I/O. + await session.execute( + text("UPDATE project SET id = id WHERE id = :project_id"), + {"project_id": self.project_id}, + ) + return + raise SemanticVectorIndexExtensionError( + f"External vector cleanup does not support SQL dialect {dialect_name!r}." + ) + def _assert_manifest_vector_ownership(self, vector_index_names: Iterable[object]) -> None: """Reject cleanup that cannot reach every externally owned vector.""" recorded_indexes = frozenset(str(name) for name in vector_index_names if str(name)) @@ -897,6 +928,7 @@ async def delete_project_vector_rows(self, *, strict_adapter_cleanup: bool = Tru configured_index = self._semantic_vector_index_name async with db.scoped_session(self.session_maker) as session: connection = await session.connection() + dialect_name = connection.dialect.name manifest_exists = await connection.run_sync( lambda sync_connection: inspect(sync_connection).has_table("search_vector_chunks") ) @@ -911,6 +943,20 @@ async def delete_project_vector_rows(self, *, strict_adapter_cleanup: bool = Tru ) manifest_has_embedding_status = "embedding_status" in manifest_columns + external_adapter_available = ( + configured_index not in _BUILT_IN_VECTOR_INDEX_NAMES + and hasattr(self, "_semantic_vector_index") + ) + if external_adapter_available: + # Constraint: PostgreSQL row locks do not cover future manifest inserts, + # while SQLite releases its write lock at commit. Hold one project-wide + # transaction through ownership discovery, adapter I/O, and manifest + # deletion so a concurrent flush cannot publish an unowned vector. + await self._lock_external_vector_project( + session, + dialect_name=dialect_name, + ) + entity_ids_by_vector_index: dict[str, list[int]] = {} if "vector_index" in manifest_columns: result = await session.execute( @@ -957,7 +1003,37 @@ async def delete_project_vector_rows(self, *, strict_adapter_cleanup: bool = Tru ), {"project_id": self.project_id}, ) + adapter_entity_ids = entity_ids_by_vector_index.get(configured_index, []) + if external_adapter_available: + if adapter_entity_ids: + try: + await self._semantic_vector_index.initialize() + for entity_id in adapter_entity_ids: + await self._semantic_vector_index.delete_entity(entity_id) + except Exception as exc: + # Trigger: a configured external adapter cannot initialize or delete. + # Why: reindex and project deletion must not discard the only + # ownership manifest for external data. + # Outcome: strict callers stop for a retry before manifest deletion. + logger.warning( + "Could not clean semantic vector adapter: " + "project_id={project_id} vector_index={vector_index} error={error}", + project_id=self.project_id, + vector_index=self._semantic_vector_index_name, + error=exc, + ) + if strict_adapter_cleanup: + raise + + await self._delete_project_builtin_vector_rows(session) + await session.execute( + text("DELETE FROM search_vector_chunks WHERE project_id = :project_id"), + {"project_id": self.project_id}, + ) await session.commit() + return + + await session.commit() adapter_entity_ids = entity_ids_by_vector_index.get(configured_index, []) if adapter_entity_ids and hasattr(self, "_semantic_vector_index"): @@ -966,10 +1042,6 @@ async def delete_project_vector_rows(self, *, strict_adapter_cleanup: bool = Tru for entity_id in adapter_entity_ids: await self._semantic_vector_index.delete_entity(entity_id) except Exception as exc: - # Trigger: a configured external adapter cannot initialize or delete. - # Why: reindex and project deletion must not discard the only - # ownership manifest for external data. - # Outcome: strict callers stop for a retry before manifest deletion. logger.warning( "Could not clean semantic vector adapter: " "project_id={project_id} vector_index={vector_index} error={error}", @@ -977,8 +1049,6 @@ async def delete_project_vector_rows(self, *, strict_adapter_cleanup: bool = Tru vector_index=self._semantic_vector_index_name, error=exc, ) - if strict_adapter_cleanup and configured_index not in _BUILT_IN_VECTOR_INDEX_NAMES: - raise async with db.scoped_session(self.session_maker) as session: await self._delete_project_builtin_vector_rows(session) diff --git a/tests/repository/test_semantic_search_base.py b/tests/repository/test_semantic_search_base.py index 65942c777..950724865 100644 --- a/tests/repository/test_semantic_search_base.py +++ b/tests/repository/test_semantic_search_base.py @@ -266,6 +266,7 @@ async def test_external_postgres_upsert_holds_manifest_lock_through_ready_commit session = AsyncMock() session.connection.return_value = SimpleNamespace(dialect=SimpleNamespace(name="postgresql")) session.execute.side_effect = [ + SimpleNamespace(), SimpleNamespace( mappings=lambda: SimpleNamespace( all=lambda: [ @@ -292,9 +293,12 @@ async def fake_scoped_session(_session_maker): await repo._persist_embeddings([(7, chunk_text)], [[1.0, 0.0, 0.0, 0.0]]) - lock_statement = session.execute.await_args_list[0].args[0] - ready_statement = session.execute.await_args_list[1].args[0] - assert "FOR UPDATE" in str(lock_statement) + project_lock_statement = session.execute.await_args_list[0].args[0] + manifest_lock_statement = session.execute.await_args_list[1].args[0] + ready_statement = session.execute.await_args_list[2].args[0] + assert str(project_lock_statement).startswith("SELECT id FROM project") + assert "FOR UPDATE" in str(project_lock_statement) + assert "FOR UPDATE" in str(manifest_lock_statement) assert "embedding_status = 'ready'" in str(ready_statement) assert context_count == 1 assert session.commit.await_count == 1 @@ -477,20 +481,51 @@ async def fake_scoped_session(_session_maker): @pytest.mark.asyncio -async def test_project_vector_cleanup_uses_available_adapter(monkeypatch): - """An available adapter should receive every manifest-owned entity deletion.""" +@pytest.mark.parametrize( + ("dialect_name", "project_lock_prefix"), + [ + ("postgresql", "SELECT id FROM project"), + ("sqlite", "UPDATE project SET id = id"), + ], +) +async def test_project_vector_cleanup_uses_available_adapter( + monkeypatch, + dialect_name, + project_lock_prefix, +): + """External cleanup should keep its project lock through manifest removal.""" repo = _ConcreteRepo() + events: list[str] = [] adapter: Any = SimpleNamespace( - initialize=AsyncMock(), - delete_entity=AsyncMock(), + initialize=AsyncMock(side_effect=lambda: events.append("initialize")), + delete_entity=AsyncMock(side_effect=lambda _entity_id: events.append("delete")), ) repo._semantic_vector_index = adapter repo._semantic_vector_index_name = "milvus" session = AsyncMock() connection = AsyncMock() + connection.dialect.name = dialect_name connection.run_sync.side_effect = [True, {"embedding_status", "vector_index"}] session.connection.return_value = connection - session.execute.return_value = SimpleNamespace(all=lambda: [(41, "milvus"), (42, "milvus")]) + + def execute(statement, _params): + sql = str(statement) + if sql.startswith(project_lock_prefix): + events.append("project_lock") + return SimpleNamespace() + if sql.startswith("SELECT DISTINCT entity_id, vector_index"): + events.append("manifest_read") + return SimpleNamespace(all=lambda: [(41, "milvus"), (42, "milvus")]) + if sql.startswith("UPDATE search_vector_chunks"): + events.append("stage") + return SimpleNamespace() + if sql.startswith("DELETE FROM search_vector_chunks"): + events.append("manifest_delete") + return SimpleNamespace() + raise AssertionError(f"Unexpected SQL: {sql}") + + session.execute.side_effect = execute + session.commit.side_effect = lambda: events.append("commit") @asynccontextmanager async def fake_scoped_session(_session_maker): @@ -505,6 +540,17 @@ async def fake_scoped_session(_session_maker): ((41,), {}), ((42,), {}), ] + assert events == [ + "project_lock", + "manifest_read", + "stage", + "initialize", + "delete", + "delete", + "manifest_delete", + "commit", + ] + assert session.commit.await_count == 1 @pytest.mark.asyncio @@ -519,9 +565,14 @@ async def test_project_vector_cleanup_preserves_manifest_after_adapter_failure(m repo._semantic_vector_index_name = "milvus" session = AsyncMock() connection = AsyncMock() + connection.dialect.name = "postgresql" connection.run_sync.side_effect = [True, {"embedding_status", "vector_index"}] session.connection.return_value = connection - session.execute.return_value = SimpleNamespace(all=lambda: [(41, "milvus")]) + session.execute.side_effect = [ + SimpleNamespace(), + SimpleNamespace(all=lambda: [(41, "milvus")]), + SimpleNamespace(), + ] @asynccontextmanager async def fake_scoped_session(_session_maker): @@ -554,9 +605,14 @@ async def test_strict_project_vector_cleanup_preserves_manifest_after_adapter_fa repo._semantic_vector_index_name = "milvus" session = AsyncMock() connection = AsyncMock() + connection.dialect.name = "postgresql" connection.run_sync.side_effect = [True, {"embedding_status", "vector_index"}] session.connection.return_value = connection - session.execute.return_value = SimpleNamespace(all=lambda: [(41, "milvus")]) + session.execute.side_effect = [ + SimpleNamespace(), + SimpleNamespace(all=lambda: [(41, "milvus")]), + SimpleNamespace(), + ] @asynccontextmanager async def fake_scoped_session(_session_maker): From 609cddf1d8e92883f21a3ac0910c4de9d0f9f4eb Mon Sep 17 00:00:00 2001 From: phernandez Date: Sun, 26 Jul 2026 15:43:06 -0500 Subject: [PATCH 21/24] fix(core): close external vector deletion races Signed-off-by: phernandez --- .../indexing/project_delete_runner.py | 43 ++- .../accepted_note_vector_cleanup.py | 56 ++- .../repository/search_repository.py | 10 +- .../repository/search_repository_base.py | 364 +++++++++++------- .../repository/semantic_vector_sync.py | 66 +++- src/basic_memory/services/project_service.py | 41 +- tests/indexing/test_project_delete_runner.py | 68 +++- .../test_external_vector_cleanup.py | 9 +- tests/repository/test_semantic_search_base.py | 123 +++++- tests/repository/test_semantic_vector_sync.py | 23 +- .../test_project_service_operations.py | 56 ++- 11 files changed, 649 insertions(+), 210 deletions(-) diff --git a/src/basic_memory/indexing/project_delete_runner.py b/src/basic_memory/indexing/project_delete_runner.py index 06f0a7ce0..af49a4c57 100644 --- a/src/basic_memory/indexing/project_delete_runner.py +++ b/src/basic_memory/indexing/project_delete_runner.py @@ -2,7 +2,7 @@ from __future__ import annotations -from collections.abc import Sequence +from collections.abc import Callable, Sequence from dataclasses import dataclass, field from enum import StrEnum from typing import Protocol, Self @@ -14,6 +14,10 @@ from basic_memory import db from basic_memory.models import Entity, NoteContent, Project from basic_memory.repository import ProjectRepository +from basic_memory.repository.accepted_note_vector_cleanup import ( + project_external_vector_index_names, +) +from basic_memory.repository.semantic_errors import SemanticVectorIndexExtensionError from basic_memory.runtime.cleanup import ( RuntimeDeleteStatus, RuntimeFileDeleteResult, @@ -90,6 +94,20 @@ class ProjectDeleteRepository(Protocol): async def delete(self, session: AsyncSession, entity_id: int) -> bool: ... +class ProjectVectorCleaner(Protocol): + """Project-scoped vector cleanup that can join the hard-delete transaction.""" + + async def delete_project_vector_rows( + self, + *, + strict_adapter_cleanup: bool = True, + session: AsyncSession | None = None, + ) -> None: ... + + +type ProjectVectorCleanerFactory = Callable[[int], ProjectVectorCleaner] + + async def load_project_file_snapshots( session: AsyncSession, *, @@ -192,6 +210,7 @@ class RepositoryProjectHardDeleter: session_maker: async_sessionmaker[AsyncSession] project_repository: ProjectDeleteRepository = field(default_factory=ProjectRepository) + project_vector_cleaner_factory: ProjectVectorCleanerFactory | None = None async def hard_delete_project( self, @@ -214,6 +233,28 @@ async def hard_delete_project( ) return ProjectHardDeleteOutcome.reactivated + if self.project_vector_cleaner_factory is not None: + vector_cleaner = self.project_vector_cleaner_factory(request.project_id) + await vector_cleaner.delete_project_vector_rows( + strict_adapter_cleanup=True, + session=session, + ) + else: + external_vector_indexes = await project_external_vector_index_names( + session, + project_id=request.project_id, + ) + if external_vector_indexes: + # Trigger: a host constructed the portable hard deleter without + # the extension adapter that owns this project's vectors. + # Why: deleting the project would cascade the only manifest and + # make remote vector cleanup impossible to retry. + # Outcome: fail closed until the host injects project cleanup. + raise SemanticVectorIndexExtensionError( + "Cannot hard-delete a project with externally owned vectors " + f"{sorted(external_vector_indexes)!r} without a project vector cleaner." + ) + deleted = await self.project_repository.delete(session, request.project_id) return ProjectHardDeleteOutcome.deleted if deleted else ProjectHardDeleteOutcome.missing diff --git a/src/basic_memory/repository/accepted_note_vector_cleanup.py b/src/basic_memory/repository/accepted_note_vector_cleanup.py index 7961d4184..b49d444f3 100644 --- a/src/basic_memory/repository/accepted_note_vector_cleanup.py +++ b/src/basic_memory/repository/accepted_note_vector_cleanup.py @@ -56,18 +56,26 @@ FROM search_vector_chunks WHERE project_id = :project_id AND entity_id IN :deleted_entity_ids - AND vector_index <> 'pgvector' + AND vector_index <> '' + AND vector_index NOT IN ('pgvector', 'sqlite-vec') """).bindparams(bindparam("deleted_entity_ids", expanding=True)) +SELECT_PROJECT_EXTERNAL_VECTOR_INDEXES_SQL = text(""" + SELECT DISTINCT vector_index + FROM search_vector_chunks + WHERE project_id = :project_id + AND vector_index <> '' + AND vector_index NOT IN ('pgvector', 'sqlite-vec') +""") + class ProjectIndexExternalVectorCleaner(Protocol): """Project-scoped capability for deleting vectors from extension storage.""" async def delete_external_entity_vectors( self, + session: AsyncSession, entity_ids: Sequence[int], - *, - vector_index_names: frozenset[str], ) -> None: ... @@ -89,6 +97,23 @@ async def project_index_vector_table_names(session: AsyncSession) -> frozenset[s return frozenset(str(table_name) for table_name in result.scalars()) +async def project_external_vector_index_names( + session: AsyncSession, + *, + project_id: ProjectId, +) -> frozenset[str]: + """Return externally owned vector indexes recorded for one project.""" + vector_table_names = await project_index_vector_table_names(session) + if "search_vector_chunks" not in vector_table_names: + return frozenset() + + result = await session.execute( + SELECT_PROJECT_EXTERNAL_VECTOR_INDEXES_SQL, + {"project_id": project_id}, + ) + return frozenset(str(vector_index) for vector_index in result.scalars()) + + async def delete_project_index_vector_rows( session: AsyncSession, *, @@ -117,22 +142,27 @@ async def delete_project_index_vector_rows( # Outcome: require the project-scoped adapter to delete those entities before # the caller-owned transaction removes built-in vectors and manifest rows. if dialect_name == "postgresql": - external_result = await session.execute( - SELECT_PROJECT_INDEX_EXTERNAL_VECTOR_INDEXES_SQL, - delete_params, - ) - external_vector_indexes = frozenset( - str(vector_index) for vector_index in external_result.scalars() - ) - if external_vector_indexes: - if external_vector_cleaner is None: + if external_vector_cleaner is None: + external_result = await session.execute( + SELECT_PROJECT_INDEX_EXTERNAL_VECTOR_INDEXES_SQL, + delete_params, + ) + external_vector_indexes = frozenset( + str(vector_index) for vector_index in external_result.scalars() + ) + if external_vector_indexes: raise SemanticVectorIndexExtensionError( "Cannot delete externally indexed entity vectors without a " "project-scoped semantic vector adapter." ) + else: + # The cleaner acquires the same project lock used by every external + # manifest prepare/write before it discovers ownership. Keeping that + # lock in this caller-owned transaction closes the insert-between-read- + # and-delete race through the manifest and entity deletes below. await external_vector_cleaner.delete_external_entity_vectors( + session, deleted_entity_ids, - vector_index_names=external_vector_indexes, ) if "search_vector_embeddings" in vector_table_names: diff --git a/src/basic_memory/repository/search_repository.py b/src/basic_memory/repository/search_repository.py index e282cc61c..29d69aaf1 100644 --- a/src/basic_memory/repository/search_repository.py +++ b/src/basic_memory/repository/search_repository.py @@ -104,14 +104,18 @@ async def delete_entity_vector_rows(self, entity_id: int) -> None: async def delete_external_entity_vectors( self, + session: AsyncSession, entity_ids: Sequence[int], - *, - vector_index_names: frozenset[str], ) -> None: """Delete DB-first entity vectors through an extension adapter.""" ... - async def delete_project_vector_rows(self, *, strict_adapter_cleanup: bool = True) -> None: + async def delete_project_vector_rows( + self, + *, + strict_adapter_cleanup: bool = True, + session: AsyncSession | None = None, + ) -> None: """Delete all semantic vector chunks and embeddings for this project.""" ... diff --git a/src/basic_memory/repository/search_repository_base.py b/src/basic_memory/repository/search_repository_base.py index b851b831f..3909ebc4c 100644 --- a/src/basic_memory/repository/search_repository_base.py +++ b/src/basic_memory/repository/search_repository_base.py @@ -235,11 +235,38 @@ async def _run_vector_query( candidate_limit: int, ) -> list[dict]: """Query the configured adapter and hydrate only live, ready manifest rows.""" - matches = await self._semantic_vector_index.search( - query_embedding, - limit=candidate_limit, - ) - return await self._hydrate_vector_matches(session, matches) + if candidate_limit <= 0: + return [] + + external_vector_index = self._semantic_vector_index_name not in _BUILT_IN_VECTOR_INDEX_NAMES + if not external_vector_index: + matches = await self._semantic_vector_index.search( + query_embedding, + limit=candidate_limit, + ) + return await self._hydrate_vector_matches(session, matches) + + scan_limit = min(candidate_limit, VECTOR_FILTER_SCAN_LIMIT) + while True: + matches = await self._semantic_vector_index.search( + query_embedding, + limit=scan_limit, + ) + hydrated = await self._hydrate_vector_matches(session, matches) + if ( + len(hydrated) >= candidate_limit + or len(matches) < scan_limit + or scan_limit >= VECTOR_FILTER_SCAN_LIMIT + ): + return hydrated[:candidate_limit] + + # Trigger: stale, pending, or wrong-model adapter hits consumed the + # requested top-k before manifest hydration. + # Why: returning early lets stale extension data crowd every live + # result out of an otherwise valid semantic search. + # Outcome: retry from the same ranked prefix with bounded geometric + # overfetch until enough live rows survive or the adapter is exhausted. + scan_limit = min(scan_limit * 2, VECTOR_FILTER_SCAN_LIMIT) async def _hydrate_vector_matches( self, @@ -339,12 +366,8 @@ async def _persist_embeddings( self._semantic_vector_index_name not in _BUILT_IN_VECTOR_INDEX_NAMES ) lock_external_write = external_vector_index and dialect_name in {"postgresql", "sqlite"} - if external_vector_index and dialect_name == "postgresql": - # Trigger: a project-wide cleanup can run beside a watcher flush. - # Why: row locks do not block a newly prepared chunk from being inserted. - # Outcome: every PostgreSQL external write shares the project-row lock - # held through cleanup, adapter deletion, and manifest removal. - await self._lock_external_vector_project(session, dialect_name=dialect_name) + if external_vector_index: + await self._lock_external_vector_write(session) if external_vector_index and dialect_name == "sqlite": # SQLite has no SELECT FOR UPDATE. A conditional no-op write takes # the database write lock before extension I/O, serializing a newer @@ -816,37 +839,69 @@ async def delete_entity_vector_rows(self, entity_id: int) -> None: async def delete_external_entity_vectors( self, + session: AsyncSession, + entity_ids: Sequence[int], + ) -> None: + """Delete externally stored entity vectors under the caller's project lock.""" + deleted_entity_ids = tuple(dict.fromkeys(entity_ids)) + await self._lock_external_vector_write(session) + if not deleted_entity_ids: + return + + params = { + "project_id": self.project_id, + **{ + f"entity_id_{index}": entity_id + for index, entity_id in enumerate(deleted_entity_ids) + }, + } + placeholders = ", ".join(f":entity_id_{index}" for index in range(len(deleted_entity_ids))) + ownership_result = await session.execute( + text( + "SELECT DISTINCT vector_index FROM search_vector_chunks " + "WHERE project_id = :project_id " + f"AND entity_id IN ({placeholders})" + ), + params, + ) + recorded_indexes = frozenset( + str(vector_index) for vector_index in ownership_result.scalars().all() + ) + await self._delete_external_entity_vectors_locked( + session, + deleted_entity_ids, + recorded_indexes=recorded_indexes, + ) + + async def _delete_external_entity_vectors_locked( + self, + session: AsyncSession, entity_ids: Sequence[int], *, - vector_index_names: frozenset[str], + recorded_indexes: frozenset[str], ) -> None: - """Delete DB-first entity vectors through the configured extension adapter.""" - if not vector_index_names: + """Delete external vectors after the caller has acquired the project lock.""" + self._assert_manifest_vector_ownership(recorded_indexes) + external_indexes = recorded_indexes - _BUILT_IN_VECTOR_INDEX_NAMES + if not external_indexes: return + + deleted_entity_ids = tuple(dict.fromkeys(entity_ids)) configured_index = self._semantic_vector_index_name - if vector_index_names != frozenset({configured_index}): - raise SemanticVectorIndexExtensionError( - "Cannot delete external vectors owned by " - f"{sorted(vector_index_names)!r} with configured adapter " - f"{configured_index!r}." - ) if not hasattr(self, "_semantic_vector_index"): raise SemanticVectorIndexExtensionError( f"Semantic vector adapter {configured_index!r} is unavailable. " "Enable semantic search and retry the entity deletion." ) - deleted_entity_ids = tuple(dict.fromkeys(entity_ids)) - if not deleted_entity_ids: - return - # Trigger: DB-first deletion runs inside a caller-owned SQL transaction. # Why: the external adapter cannot participate in that transaction. If its # delete succeeds and the caller later rolls back, ready manifests would # incorrectly claim the now-missing vectors are searchable. # Outcome: commit a non-searchable retry marker in an independent session - # before touching external storage; a retried delete remains idempotent. - params = { + # before touching external storage while the caller retains the project + # lock; a retried delete remains idempotent and no new generation can race. + stage_params = { "project_id": self.project_id, "vector_index": configured_index, **{ @@ -855,16 +910,30 @@ async def delete_external_entity_vectors( }, } placeholders = ", ".join(f":entity_id_{index}" for index in range(len(deleted_entity_ids))) - async with db.scoped_session(self.session_maker) as session: + connection = await session.connection() + if connection.dialect.name == "postgresql": + async with db.scoped_session(self.session_maker) as marker_session: + await marker_session.execute( + text( + "UPDATE search_vector_chunks SET embedding_status = 'pending' " + "WHERE project_id = :project_id AND vector_index = :vector_index " + f"AND entity_id IN ({placeholders})" + ), + stage_params, + ) + await marker_session.commit() + else: + # SQLite permits only one writer, so a second marker transaction would + # deadlock behind the project write lock. Keep the marker in the caller + # transaction; extension cleanup still remains serialized. await session.execute( text( "UPDATE search_vector_chunks SET embedding_status = 'pending' " "WHERE project_id = :project_id AND vector_index = :vector_index " f"AND entity_id IN ({placeholders})" ), - params, + stage_params, ) - await session.commit() await self._semantic_vector_index.initialize() for entity_id in deleted_entity_ids: @@ -898,6 +967,23 @@ async def _lock_external_vector_project( f"External vector cleanup does not support SQL dialect {dialect_name!r}." ) + async def _lock_external_vector_write(self, session: AsyncSession) -> None: + """Share one project lock across external manifest mutations and cleanup.""" + if not self._uses_external_vector_index(): + return + + connection = await session.connection() + await self._lock_external_vector_project( + session, + dialect_name=connection.dialect.name, + ) + + def _uses_external_vector_index(self) -> bool: + """Return whether this repository writes vectors outside the SQL backend.""" + return self._semantic_vector_index_name not in _BUILT_IN_VECTOR_INDEX_NAMES and hasattr( + self, "_semantic_vector_index" + ) + def _assert_manifest_vector_ownership(self, vector_index_names: Iterable[object]) -> None: """Reject cleanup that cannot reach every externally owned vector.""" recorded_indexes = frozenset(str(name) for name in vector_index_names if str(name)) @@ -913,7 +999,12 @@ def _assert_manifest_vector_ownership(self, vector_index_names: Iterable[object] f"{configured_index!r}. Restore the owning adapter and retry." ) - async def delete_project_vector_rows(self, *, strict_adapter_cleanup: bool = True) -> None: + async def delete_project_vector_rows( + self, + *, + strict_adapter_cleanup: bool = True, + session: AsyncSession | None = None, + ) -> None: """Delete this project's vectors through the configured storage adapter. Core enumerates ownership from the SQL manifest because the adapter @@ -925,123 +1016,111 @@ async def delete_project_vector_rows(self, *, strict_adapter_cleanup: bool = Tru fail closed because only the owning extension can safely remove previously written vectors. """ - configured_index = self._semantic_vector_index_name - async with db.scoped_session(self.session_maker) as session: - connection = await session.connection() - dialect_name = connection.dialect.name - manifest_exists = await connection.run_sync( - lambda sync_connection: inspect(sync_connection).has_table("search_vector_chunks") + if session is not None: + await self._delete_project_vector_rows_in_session( + session, + strict_adapter_cleanup=strict_adapter_cleanup, ) - if not manifest_exists: - return + return - manifest_columns = await connection.run_sync( - lambda sync_connection: { - str(column["name"]) - for column in inspect(sync_connection).get_columns("search_vector_chunks") - } + async with db.scoped_session(self.session_maker) as owned_session: + changed = await self._delete_project_vector_rows_in_session( + owned_session, + strict_adapter_cleanup=strict_adapter_cleanup, ) - manifest_has_embedding_status = "embedding_status" in manifest_columns + if changed: + await owned_session.commit() - external_adapter_available = ( - configured_index not in _BUILT_IN_VECTOR_INDEX_NAMES - and hasattr(self, "_semantic_vector_index") + async def _delete_project_vector_rows_in_session( + self, + session: AsyncSession, + *, + strict_adapter_cleanup: bool, + ) -> bool: + """Delete project vectors while retaining the caller's transaction boundary.""" + connection = await session.connection() + manifest_exists = await connection.run_sync( + lambda sync_connection: inspect(sync_connection).has_table("search_vector_chunks") + ) + if not manifest_exists: + return False + + manifest_columns = await connection.run_sync( + lambda sync_connection: { + str(column["name"]) + for column in inspect(sync_connection).get_columns("search_vector_chunks") + } + ) + manifest_has_embedding_status = "embedding_status" in manifest_columns + configured_index = self._semantic_vector_index_name + external_adapter_available = ( + configured_index not in _BUILT_IN_VECTOR_INDEX_NAMES + and hasattr(self, "_semantic_vector_index") + ) + if external_adapter_available: + # Constraint: the project row is the only lock that also covers future + # manifest inserts. The caller retains it through adapter I/O, manifest + # removal, and—during hard deletion—the project-row delete itself. + await self._lock_external_vector_write(session) + + entity_ids_by_vector_index: dict[str, list[int]] = {} + if "vector_index" in manifest_columns: + result = await session.execute( + text( + "SELECT DISTINCT entity_id, vector_index FROM search_vector_chunks " + "WHERE project_id = :project_id ORDER BY vector_index, entity_id" + ), + {"project_id": self.project_id}, ) - if external_adapter_available: - # Constraint: PostgreSQL row locks do not cover future manifest inserts, - # while SQLite releases its write lock at commit. Hold one project-wide - # transaction through ownership discovery, adapter I/O, and manifest - # deletion so a concurrent flush cannot publish an unowned vector. - await self._lock_external_vector_project( - session, - dialect_name=dialect_name, - ) + for entity_id, vector_index in result.all(): + entity_ids_by_vector_index.setdefault(str(vector_index), []).append(int(entity_id)) + else: + result = await session.execute( + text( + "SELECT DISTINCT entity_id FROM search_vector_chunks " + "WHERE project_id = :project_id ORDER BY entity_id" + ), + {"project_id": self.project_id}, + ) + legacy_vector_index = ( + "sqlite-vec" if connection.dialect.name == "sqlite" else "pgvector" + ) + entity_ids_by_vector_index[legacy_vector_index] = [ + int(entity_id) for entity_id in result.scalars().all() + ] - entity_ids_by_vector_index: dict[str, list[int]] = {} - if "vector_index" in manifest_columns: - result = await session.execute( - text( - "SELECT DISTINCT entity_id, vector_index FROM search_vector_chunks " - "WHERE project_id = :project_id ORDER BY vector_index, entity_id" - ), - {"project_id": self.project_id}, - ) - for entity_id, vector_index in result.all(): - entity_ids_by_vector_index.setdefault(str(vector_index), []).append( - int(entity_id) - ) - else: - result = await session.execute( - text( - "SELECT DISTINCT entity_id FROM search_vector_chunks " - "WHERE project_id = :project_id ORDER BY entity_id" - ), - {"project_id": self.project_id}, - ) - legacy_vector_index = ( - "sqlite-vec" if connection.dialect.name == "sqlite" else "pgvector" - ) - entity_ids_by_vector_index[legacy_vector_index] = [ - int(entity_id) for entity_id in result.scalars().all() - ] + # Trigger: manifests belong to an external index other than the available adapter. + # Why: adapter configuration can change after vectors were written, and deleting + # the manifest would discard the only durable routing information for old vectors. + # Outcome: fail before touching any adapter or manifest so the owner can be restored. + self._assert_manifest_vector_ownership(entity_ids_by_vector_index) - # Trigger: manifests belong to an external index other than the available adapter. - # Why: adapter configuration can change after vectors were written, and deleting - # the manifest would discard the only durable routing information for old vectors. - # Outcome: fail before touching any adapter or manifest so the owner can be restored. - self._assert_manifest_vector_ownership(entity_ids_by_vector_index) - - # Trigger: the manifest predates embedding lifecycle state. - # Why: legacy SQLite schemas must reach cleanup before lazy schema repair runs. - # Outcome: skip staging only for that obsolete schema; the manifest is still - # deleted below, and later semantic initialization recreates the current schema. - if manifest_has_embedding_status: - await session.execute( - text( - "UPDATE search_vector_chunks SET embedding_status = 'pending' " - "WHERE project_id = :project_id" - ), - {"project_id": self.project_id}, - ) - adapter_entity_ids = entity_ids_by_vector_index.get(configured_index, []) + builtin_indexes = frozenset(entity_ids_by_vector_index) & _BUILT_IN_VECTOR_INDEX_NAMES + if manifest_has_embedding_status and (not external_adapter_available or builtin_indexes): + builtin_filter = "" if external_adapter_available: - if adapter_entity_ids: - try: - await self._semantic_vector_index.initialize() - for entity_id in adapter_entity_ids: - await self._semantic_vector_index.delete_entity(entity_id) - except Exception as exc: - # Trigger: a configured external adapter cannot initialize or delete. - # Why: reindex and project deletion must not discard the only - # ownership manifest for external data. - # Outcome: strict callers stop for a retry before manifest deletion. - logger.warning( - "Could not clean semantic vector adapter: " - "project_id={project_id} vector_index={vector_index} error={error}", - project_id=self.project_id, - vector_index=self._semantic_vector_index_name, - error=exc, - ) - if strict_adapter_cleanup: - raise - - await self._delete_project_builtin_vector_rows(session) - await session.execute( - text("DELETE FROM search_vector_chunks WHERE project_id = :project_id"), - {"project_id": self.project_id}, - ) - await session.commit() - return - - await session.commit() + builtin_filter = " AND vector_index IN ('pgvector', 'sqlite-vec')" + await session.execute( + text( + "UPDATE search_vector_chunks SET embedding_status = 'pending' " + f"WHERE project_id = :project_id{builtin_filter}" + ), + {"project_id": self.project_id}, + ) adapter_entity_ids = entity_ids_by_vector_index.get(configured_index, []) - if adapter_entity_ids and hasattr(self, "_semantic_vector_index"): + if external_adapter_available: try: - await self._semantic_vector_index.initialize() - for entity_id in adapter_entity_ids: - await self._semantic_vector_index.delete_entity(entity_id) + await self._delete_external_entity_vectors_locked( + session, + adapter_entity_ids, + recorded_indexes=frozenset(entity_ids_by_vector_index), + ) except Exception as exc: + # Trigger: a configured external adapter cannot initialize or delete. + # Why: reindex and project deletion must not discard the only + # ownership manifest for external data. + # Outcome: strict callers stop for a retry before manifest deletion. logger.warning( "Could not clean semantic vector adapter: " "project_id={project_id} vector_index={vector_index} error={error}", @@ -1049,14 +1128,15 @@ async def delete_project_vector_rows(self, *, strict_adapter_cleanup: bool = Tru vector_index=self._semantic_vector_index_name, error=exc, ) + if strict_adapter_cleanup: + raise - async with db.scoped_session(self.session_maker) as session: - await self._delete_project_builtin_vector_rows(session) - await session.execute( - text("DELETE FROM search_vector_chunks WHERE project_id = :project_id"), - {"project_id": self.project_id}, - ) - await session.commit() + await self._delete_project_builtin_vector_rows(session) + await session.execute( + text("DELETE FROM search_vector_chunks WHERE project_id = :project_id"), + {"project_id": self.project_id}, + ) + return True async def delete_stale_vector_rows(self) -> None: """Delete vectors whose source entity no longer exists. diff --git a/src/basic_memory/repository/semantic_vector_sync.py b/src/basic_memory/repository/semantic_vector_sync.py index 15e5fced1..cba86999a 100644 --- a/src/basic_memory/repository/semantic_vector_sync.py +++ b/src/basic_memory/repository/semantic_vector_sync.py @@ -732,7 +732,45 @@ async def prepare_entity_vector_jobs_window( async with repository._prepare_entity_write_scope(): async with db.scoped_session(repository.session_maker) as session: await repository._prepare_vector_session(session) - for index, plan in mutation_plans: + await repository._lock_external_vector_write(session) + plans_to_apply = mutation_plans + if repository._uses_external_vector_index(): + # Trigger: entity deletion can finish after the shared read + # snapshot but before this prepare transaction gets the + # project lock. + # Why: applying that stale plan would recreate a manifest + # for an entity the delete just removed. + # Outcome: external writes re-read and re-plan under the + # shared lock before making any manifest mutation. + mutation_entity_ids = [plan.entity_id for _index, plan in mutation_plans] + locked_source_rows = await repository._fetch_prepare_window_source_rows( + session, + mutation_entity_ids, + ) + locked_existing_rows = await repository._fetch_prepare_window_existing_rows( + session, + mutation_entity_ids, + ) + plans_to_apply = [] + for index, original_plan in mutation_plans: + replanned = plan_entity_vector_jobs_prefetched( + repository, + entity_id=original_plan.entity_id, + source_rows=locked_source_rows.get( + original_plan.entity_id, + [], + ), + existing_rows=locked_existing_rows.get( + original_plan.entity_id, + [], + ), + ) + if isinstance(replanned, PreparedEntityVectorSync): + prepared_by_index[index] = replanned + else: + plans_to_apply.append((index, replanned)) + + for index, plan in plans_to_apply: prepared_by_index[index] = await apply_entity_vector_prepare_plan( repository, session, @@ -788,7 +826,31 @@ async def prepare_entity_vector_jobs_prefetched( async with repository._prepare_entity_write_scope(): async with db.scoped_session(repository.session_maker) as session: await repository._prepare_vector_session(session) - prepared = await apply_entity_vector_prepare_plan(repository, session, planned) + await repository._lock_external_vector_write(session) + locked_plan = planned + if repository._uses_external_vector_index(): + locked_source_rows = await repository._fetch_prepare_window_source_rows( + session, + [entity_id], + ) + locked_existing_rows = await repository._fetch_prepare_window_existing_rows( + session, + [entity_id], + ) + locked_plan = plan_entity_vector_jobs_prefetched( + repository, + entity_id=entity_id, + source_rows=locked_source_rows.get(entity_id, []), + existing_rows=locked_existing_rows.get(entity_id, []), + ) + if isinstance(locked_plan, PreparedEntityVectorSync): + return locked_plan + + prepared = await apply_entity_vector_prepare_plan( + repository, + session, + locked_plan, + ) await session.commit() await repository._finalize_prepared_vector_deletions(prepared) return prepared diff --git a/src/basic_memory/services/project_service.py b/src/basic_memory/services/project_service.py index de7631b80..e3a8875ac 100644 --- a/src/basic_memory/services/project_service.py +++ b/src/basic_memory/services/project_service.py @@ -93,6 +93,16 @@ def projects(self) -> Dict[str, str]: """ return self.config_manager.projects + def _project_search_repository(self, project_id: int) -> SearchRepository: + """Build the project-scoped repository that owns semantic vector cleanup.""" + if self._search_repository_factory is not None: + return self._search_repository_factory(project_id) + return create_search_repository( + session_maker=self.session_maker, + project_id=project_id, + app_config=self.config_manager.config, + ) + @property def default_project(self) -> Optional[str]: """Get the name of the default project. @@ -353,22 +363,20 @@ async def remove_project(self, name: str, delete_notes: bool = False) -> None: if is_default: raise ValueError(f"Cannot remove the default project '{name}'") # pragma: no cover - # Trigger: project deletion can remove the only SQL ownership manifest for - # vectors stored by an extension such as Milvus. - # Why: external storage has no database cascade and cannot reconcile after - # the project manifest disappears. - # Outcome: delete adapter-owned vectors while project ownership is still known. - if self._search_repository_factory is None: - search_repository = create_search_repository( - session_maker=self.session_maker, - project_id=project_id, - app_config=self.config_manager.config, - ) - else: - search_repository = self._search_repository_factory(project_id) - await search_repository.delete_project_vector_rows(strict_adapter_cleanup=True) + search_repository = self._project_search_repository(project_id) async with db.scoped_session(self.session_maker) as session: + # Trigger: project deletion can remove the only SQL ownership manifest + # for vectors stored by an extension such as Milvus. + # Why: the project lock must survive adapter cleanup through both the + # manifest and project-row deletes; committing cleanup earlier reopens + # a window where a watcher can publish a new unowned vector. + # Outcome: vector cleanup and project deletion share one transaction. + await search_repository.delete_project_vector_rows( + strict_adapter_cleanup=True, + session=session, + ) + # Remove from config if it exists there (may not exist in cloud mode) try: self.config_manager.remove_project(name) @@ -535,6 +543,11 @@ async def synchronize_projects(self) -> None: # pragma: no cover logger.info( f"Removing project '{name}' from database (deleted from config, source of truth)" ) + search_repository = self._project_search_repository(project.id) + await search_repository.delete_project_vector_rows( + strict_adapter_cleanup=True, + session=session, + ) await self.repository.delete(session, project.id) # Ensure database default project state is consistent diff --git a/tests/indexing/test_project_delete_runner.py b/tests/indexing/test_project_delete_runner.py index b13d40f44..3eca140fe 100644 --- a/tests/indexing/test_project_delete_runner.py +++ b/tests/indexing/test_project_delete_runner.py @@ -2,12 +2,14 @@ from collections.abc import AsyncGenerator from datetime import UTC, datetime +from unittest.mock import AsyncMock import pytest import pytest_asyncio from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker, create_async_engine from sqlalchemy.pool import StaticPool +import basic_memory.indexing.project_delete_runner as project_delete_runner_module from basic_memory.indexing.project_delete_runner import ( ProjectDeletePreflightResult, ProjectHardDeleteOutcome, @@ -17,6 +19,7 @@ ) from basic_memory.models import Base as BasicMemoryBase from basic_memory.models import Entity, NoteContent, Project +from basic_memory.repository.semantic_errors import SemanticVectorIndexExtensionError from basic_memory.runtime.cleanup import ( RuntimeDeleteStatus, RuntimeFileDeleteResult, @@ -87,15 +90,33 @@ async def hard_delete_project( class FakeProjectDeleteRepository: - def __init__(self, *, deleted: bool) -> None: + def __init__(self, *, deleted: bool, events: list[str] | None = None) -> None: self.deleted = deleted self.entity_ids: list[int] = [] + self.events = events async def delete(self, session: AsyncSession, entity_id: int) -> bool: + if self.events is not None: + self.events.append("project_delete") self.entity_ids.append(entity_id) return self.deleted +class FakeProjectVectorCleaner: + def __init__(self, events: list[str]) -> None: + self.events = events + self.calls: list[tuple[bool, AsyncSession | None]] = [] + + async def delete_project_vector_rows( + self, + *, + strict_adapter_cleanup: bool = True, + session: AsyncSession | None = None, + ) -> None: + self.events.append("vector_cleanup") + self.calls.append((strict_adapter_cleanup, session)) + + def project_delete_request( *, project_id: int = 101, @@ -355,6 +376,51 @@ async def test_repository_project_hard_deleter_uses_injected_project_repository( assert repository.entity_ids == [project.id] +@pytest.mark.asyncio +async def test_repository_project_hard_deleter_cleans_vectors_in_delete_transaction( + project_delete_session_maker: async_sessionmaker[AsyncSession], +) -> None: + project = await create_project_with_note(project_delete_session_maker, is_active=False) + events: list[str] = [] + repository = FakeProjectDeleteRepository(deleted=True, events=events) + cleaner = FakeProjectVectorCleaner(events) + + outcome = await RepositoryProjectHardDeleter( + session_maker=project_delete_session_maker, + project_repository=repository, + project_vector_cleaner_factory=lambda _project_id: cleaner, + ).hard_delete_project(project_delete_request(project_id=project.id)) + + assert outcome is ProjectHardDeleteOutcome.deleted + assert events == ["vector_cleanup", "project_delete"] + assert len(cleaner.calls) == 1 + strict_cleanup, cleanup_session = cleaner.calls[0] + assert strict_cleanup is True + assert cleanup_session is not None + + +@pytest.mark.asyncio +async def test_repository_project_hard_deleter_rejects_unowned_external_cleanup( + project_delete_session_maker: async_sessionmaker[AsyncSession], + monkeypatch: pytest.MonkeyPatch, +) -> None: + project = await create_project_with_note(project_delete_session_maker, is_active=False) + monkeypatch.setattr( + project_delete_runner_module, + "project_external_vector_index_names", + AsyncMock(return_value=frozenset({"milvus"})), + ) + + with pytest.raises(SemanticVectorIndexExtensionError, match="project vector cleaner"): + await RepositoryProjectHardDeleter( + session_maker=project_delete_session_maker + ).hard_delete_project(project_delete_request(project_id=project.id)) + + async with project_delete_session_maker() as session: + surviving_project = await session.get(Project, project.id) + assert surviving_project is not None + + @pytest.mark.asyncio async def test_repository_project_hard_deleter_defaults_to_core_project_repository( project_delete_session_maker: async_sessionmaker[AsyncSession], diff --git a/tests/repository/test_external_vector_cleanup.py b/tests/repository/test_external_vector_cleanup.py index 62b48f232..54874e867 100644 --- a/tests/repository/test_external_vector_cleanup.py +++ b/tests/repository/test_external_vector_cleanup.py @@ -36,15 +36,14 @@ async def execute(self, statement: object, params: dict[str, object] | None = No class _RecordingExternalCleaner: def __init__(self) -> None: - self.calls: list[tuple[tuple[int, ...], frozenset[str]]] = [] + self.calls: list[tuple[AsyncSession, tuple[int, ...]]] = [] async def delete_external_entity_vectors( self, + session: AsyncSession, entity_ids: Sequence[int], - *, - vector_index_names: frozenset[str], ) -> None: - self.calls.append((tuple(entity_ids), vector_index_names)) + self.calls.append((session, tuple(entity_ids))) @pytest.mark.asyncio @@ -65,7 +64,7 @@ async def test_external_vectors_are_deleted_before_their_manifest() -> None: external_vector_cleaner=cleaner, ) - assert cleaner.calls == [((41, 42), frozenset({"milvus"}))] + assert cleaner.calls == [(cast(AsyncSession, session), (41, 42))] assert str(session.executed[-1][0]).lstrip().startswith("DELETE FROM search_vector_chunks") diff --git a/tests/repository/test_semantic_search_base.py b/tests/repository/test_semantic_search_base.py index 950724865..78753f38f 100644 --- a/tests/repository/test_semantic_search_base.py +++ b/tests/repository/test_semantic_search_base.py @@ -175,11 +175,43 @@ def hydrated_batch(_statement, params): assert [row["entity_id"] for row in hydrated] == list(range(600)) +@pytest.mark.asyncio +async def test_external_vector_query_overfetches_past_stale_adapter_hits( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Stale top-k extension hits must not crowd live manifest rows out.""" + repo = _ConcreteRepo() + repo._semantic_vector_index_name = "milvus" + + def matches(count: int) -> list[VectorMatch]: + return [ + VectorMatch( + key=VectorKey(entity_id=entity_id, chunk_key=f"entity:{entity_id}:0"), + similarity=0.9, + ) + for entity_id in range(count) + ] + + adapter: Any = SimpleNamespace(search=AsyncMock(side_effect=[matches(2), matches(4)])) + repo._semantic_vector_index = adapter + live_rows = [ + {"entity_id": 2, "chunk_key": "entity:2:0", "best_similarity": 0.9}, + {"entity_id": 3, "chunk_key": "entity:3:0", "best_similarity": 0.8}, + ] + hydrate = AsyncMock(side_effect=[[], live_rows]) + monkeypatch.setattr(repo, "_hydrate_vector_matches", hydrate) + + result = await SearchRepositoryBase._run_vector_query(repo, AsyncMock(), [0.1], 2) + + assert result == live_rows + assert [call.kwargs["limit"] for call in adapter.search.await_args_list] == [2, 4] + + @pytest.mark.asyncio async def test_embedding_persistence_skips_stale_source_generation(monkeypatch) -> None: """An obsolete embedding job must not overwrite the current adapter record.""" repo = _ConcreteRepo() - repo._semantic_vector_index_name = "milvus" + repo._semantic_vector_index_name = "pgvector" repo._embedding_provider = SimpleNamespace(model_name="stub", dimensions=4) adapter = _RecordingVectorIndex() repo._semantic_vector_index = adapter @@ -214,7 +246,7 @@ async def fake_scoped_session(_session_maker): async def test_embedding_ready_update_requires_source_generation(monkeypatch) -> None: """Ready state must belong to the exact source text that produced the vector.""" repo = _ConcreteRepo() - repo._semantic_vector_index_name = "milvus" + repo._semantic_vector_index_name = "pgvector" repo._embedding_provider = SimpleNamespace(model_name="stub", dimensions=4) adapter = _RecordingVectorIndex() repo._semantic_vector_index = adapter @@ -320,6 +352,7 @@ async def test_external_sqlite_upsert_holds_write_lock_through_ready_commit( session = AsyncMock() session.connection.return_value = SimpleNamespace(dialect=SimpleNamespace(name="sqlite")) session.execute.side_effect = [ + SimpleNamespace(), SimpleNamespace( mappings=lambda: SimpleNamespace( all=lambda: [ @@ -346,10 +379,14 @@ async def fake_scoped_session(_session_maker): await repo._persist_embeddings([(7, chunk_text)], [[1.0, 0.0, 0.0, 0.0]]) - lock_statement = session.execute.await_args_list[0].args[0] - ready_statement = session.execute.await_args_list[1].args[0] - assert "UPDATE search_vector_chunks SET source_hash = source_hash" in str(lock_statement) - assert "RETURNING id, entity_id, chunk_key, source_hash" in str(lock_statement) + project_lock_statement = session.execute.await_args_list[0].args[0] + manifest_lock_statement = session.execute.await_args_list[1].args[0] + ready_statement = session.execute.await_args_list[2].args[0] + assert str(project_lock_statement).startswith("UPDATE project SET id = id") + assert "UPDATE search_vector_chunks SET source_hash = source_hash" in str( + manifest_lock_statement + ) + assert "RETURNING id, entity_id, chunk_key, source_hash" in str(manifest_lock_statement) assert "embedding_status = 'ready'" in str(ready_statement) assert context_count == 1 assert session.commit.await_count == 1 @@ -507,6 +544,9 @@ async def test_project_vector_cleanup_uses_available_adapter( connection.dialect.name = dialect_name connection.run_sync.side_effect = [True, {"embedding_status", "vector_index"}] session.connection.return_value = connection + marker_session = AsyncMock() + marker_session.execute.side_effect = lambda *_args, **_kwargs: events.append("durable_stage") + marker_session.commit.side_effect = lambda: events.append("marker_commit") def execute(statement, _params): sql = str(statement) @@ -526,10 +566,11 @@ def execute(statement, _params): session.execute.side_effect = execute session.commit.side_effect = lambda: events.append("commit") + scoped_sessions = iter([session, marker_session] if dialect_name == "postgresql" else [session]) @asynccontextmanager async def fake_scoped_session(_session_maker): - yield session + yield next(scoped_sessions) monkeypatch.setattr(search_repository_base_module.db, "scoped_session", fake_scoped_session) @@ -540,16 +581,20 @@ async def fake_scoped_session(_session_maker): ((41,), {}), ((42,), {}), ] - assert events == [ + expected_events = [ "project_lock", "manifest_read", - "stage", "initialize", "delete", "delete", "manifest_delete", "commit", ] + if dialect_name == "postgresql": + expected_events[2:2] = ["durable_stage", "marker_commit"] + else: + expected_events.insert(2, "stage") + assert events == expected_events assert session.commit.await_count == 1 @@ -571,12 +616,13 @@ async def test_project_vector_cleanup_preserves_manifest_after_adapter_failure(m session.execute.side_effect = [ SimpleNamespace(), SimpleNamespace(all=lambda: [(41, "milvus")]), - SimpleNamespace(), ] + marker_session = AsyncMock() + scoped_sessions = iter([session, marker_session]) @asynccontextmanager async def fake_scoped_session(_session_maker): - yield session + yield next(scoped_sessions) warning = Mock() monkeypatch.setattr(search_repository_base_module.db, "scoped_session", fake_scoped_session) @@ -586,6 +632,7 @@ async def fake_scoped_session(_session_maker): await repo.delete_project_vector_rows() adapter.delete_entity.assert_not_awaited() + marker_session.commit.assert_awaited_once() warning.assert_called_once() statements = [str(call.args[0]) for call in session.execute.await_args_list] assert not any( @@ -611,18 +658,20 @@ async def test_strict_project_vector_cleanup_preserves_manifest_after_adapter_fa session.execute.side_effect = [ SimpleNamespace(), SimpleNamespace(all=lambda: [(41, "milvus")]), - SimpleNamespace(), ] + marker_session = AsyncMock() + scoped_sessions = iter([session, marker_session]) @asynccontextmanager async def fake_scoped_session(_session_maker): - yield session + yield next(scoped_sessions) monkeypatch.setattr(search_repository_base_module.db, "scoped_session", fake_scoped_session) with pytest.raises(RuntimeError, match="adapter unavailable"): await repo.delete_project_vector_rows(strict_adapter_cleanup=True) + marker_session.commit.assert_awaited_once() statements = [str(call.args[0]) for call in session.execute.await_args_list] assert not any( statement.startswith("DELETE FROM search_vector_chunks") for statement in statements @@ -748,25 +797,49 @@ async def test_external_entity_cleanup_uses_matching_project_adapter(monkeypatch ) repo._semantic_vector_index = adapter repo._semantic_vector_index_name = "milvus" - session = AsyncMock() - session.execute.side_effect = lambda *_args, **_kwargs: events.append("stage") - session.commit.side_effect = lambda: events.append("commit") + caller_session = AsyncMock() + caller_session.connection.return_value = SimpleNamespace( + dialect=SimpleNamespace(name="postgresql") + ) + marker_session = AsyncMock() + + def caller_execute(statement, _params): + sql = str(statement) + if sql.startswith("SELECT id FROM project"): + events.append("project_lock") + return SimpleNamespace() + if sql.startswith("SELECT DISTINCT vector_index"): + events.append("ownership_read") + return SimpleNamespace(scalars=lambda: SimpleNamespace(all=lambda: ["milvus"])) + raise AssertionError(f"Unexpected caller SQL: {sql}") + + caller_session.execute.side_effect = caller_execute + marker_session.execute.side_effect = lambda *_args, **_kwargs: events.append("stage") + marker_session.commit.side_effect = lambda: events.append("marker_commit") @asynccontextmanager async def fake_scoped_session(_session_maker): - yield session + yield marker_session monkeypatch.setattr(search_repository_base_module.db, "scoped_session", fake_scoped_session) await repo.delete_external_entity_vectors( + caller_session, [41, 42], - vector_index_names=frozenset({"milvus"}), ) adapter.initialize.assert_awaited_once() assert adapter.delete_entity.await_args_list == [((41,), {}), ((42,), {})] - assert events == ["stage", "commit", "initialize", "delete", "delete"] - stage_statement = str(session.execute.await_args.args[0]) + assert events == [ + "project_lock", + "ownership_read", + "stage", + "marker_commit", + "initialize", + "delete", + "delete", + ] + stage_statement = str(marker_session.execute.await_args.args[0]) assert stage_statement.startswith( "UPDATE search_vector_chunks SET embedding_status = 'pending'" ) @@ -777,11 +850,19 @@ async def test_external_entity_cleanup_rejects_mismatched_adapter() -> None: """A configured adapter must not delete rows owned by another extension.""" repo = _ConcreteRepo() repo._semantic_vector_index_name = "milvus" + adapter: Any = SimpleNamespace() + repo._semantic_vector_index = adapter + session = AsyncMock() + session.connection.return_value = SimpleNamespace(dialect=SimpleNamespace(name="postgresql")) + session.execute.side_effect = [ + SimpleNamespace(), + SimpleNamespace(scalars=lambda: SimpleNamespace(all=lambda: ["pinecone"])), + ] with pytest.raises(SemanticVectorIndexExtensionError, match="owned by"): await repo.delete_external_entity_vectors( + session, [41], - vector_index_names=frozenset({"pinecone"}), ) diff --git a/tests/repository/test_semantic_vector_sync.py b/tests/repository/test_semantic_vector_sync.py index 27ba63df6..75026d174 100644 --- a/tests/repository/test_semantic_vector_sync.py +++ b/tests/repository/test_semantic_vector_sync.py @@ -375,10 +375,17 @@ def _stub_plan(repository, *, entity_id, source_rows, existing_rows): monkeypatch.setattr(semantic_vector_sync.db, "scoped_session", scoped_session) monkeypatch.setattr(repository, "_prepare_vector_session", AsyncMock()) - monkeypatch.setattr(repository, "_fetch_prepare_window_source_rows", AsyncMock(return_value={})) + lock_external_vector_write = AsyncMock() monkeypatch.setattr( - repository, "_fetch_prepare_window_existing_rows", AsyncMock(return_value={}) - ) + repository, + "_lock_external_vector_write", + lock_external_vector_write, + ) + fetch_source_rows = AsyncMock(return_value={}) + fetch_existing_rows = AsyncMock(return_value={}) + monkeypatch.setattr(repository, "_fetch_prepare_window_source_rows", fetch_source_rows) + monkeypatch.setattr(repository, "_fetch_prepare_window_existing_rows", fetch_existing_rows) + monkeypatch.setattr(repository, "_uses_external_vector_index", Mock(return_value=True)) monkeypatch.setattr(repository, "_prepare_entity_write_scope", write_scope) monkeypatch.setattr(semantic_vector_sync, "plan_entity_vector_jobs_prefetched", _stub_plan) monkeypatch.setattr( @@ -398,6 +405,9 @@ def _stub_plan(repository, *, entity_id, source_rows, existing_rows): assert str(prepared[1]) == "write failed" assert isinstance(prepared[3], ValueError) assert str(prepared[3]) == "planning failed" + lock_external_vector_write.assert_awaited_once() + assert fetch_source_rows.await_count == 2 + assert fetch_existing_rows.await_count == 2 @pytest.mark.asyncio @@ -441,6 +451,12 @@ async def write_scope(): delete_entity_chunks = AsyncMock(return_value=[]) monkeypatch.setattr(repository, "_prepare_entity_write_scope", write_scope) monkeypatch.setattr(repository, "_prepare_vector_session", AsyncMock()) + lock_external_vector_write = AsyncMock() + monkeypatch.setattr( + repository, + "_lock_external_vector_write", + lock_external_vector_write, + ) monkeypatch.setattr(repository, "_delete_entity_chunks", delete_entity_chunks) monkeypatch.setattr(repository, "_build_chunk_records", Mock(return_value=[])) @@ -508,6 +524,7 @@ async def write_scope(): ) ], ) + assert lock_external_vector_write.await_count == 2 @pytest.mark.asyncio diff --git a/tests/services/test_project_service_operations.py b/tests/services/test_project_service_operations.py index 867382959..f36c06cb5 100644 --- a/tests/services/test_project_service_operations.py +++ b/tests/services/test_project_service_operations.py @@ -111,7 +111,8 @@ async def test_remove_project_cleans_external_vectors_before_database_delete( async def delete_after_vector_cleanup(session, entity_id: int) -> bool: search_repository.delete_project_vector_rows.assert_awaited_once_with( - strict_adapter_cleanup=True + strict_adapter_cleanup=True, + session=session, ) return await original_delete(session, entity_id) @@ -120,9 +121,51 @@ async def delete_after_vector_cleanup(session, entity_id: int) -> bool: await service.remove_project(project_name) search_repository_factory.assert_called_once_with(project_id) - search_repository.delete_project_vector_rows.assert_awaited_once_with( - strict_adapter_cleanup=True + search_repository.delete_project_vector_rows.assert_awaited_once() + assert ( + search_repository.delete_project_vector_rows.await_args.kwargs["strict_adapter_cleanup"] + is True ) + assert search_repository.delete_project_vector_rows.await_args.kwargs["session"] is not None + + +@pytest.mark.asyncio +async def test_project_reconciliation_cleans_vectors_before_database_delete( + project_service: ProjectService, +) -> None: + """Config reconciliation must preserve extension ownership like explicit removal.""" + project_name = f"reconcile-vector-project-{os.urandom(4).hex()}" + search_repository = SimpleNamespace(delete_project_vector_rows=AsyncMock()) + search_repository_factory = Mock(return_value=search_repository) + service = ProjectService( + repository=project_service.repository, + session_maker=project_service.session_maker, + file_service=project_service.file_service, + search_repository_factory=search_repository_factory, + ) + + async with db.scoped_session(service.session_maker) as session: + project = await service.repository.create( + session, + { + "name": project_name, + "path": f"/tmp/{project_name}", + "permalink": project_name, + "is_active": True, + }, + ) + project_id = project.id + + await service.synchronize_projects() + + search_repository_factory.assert_any_call(project_id) + search_repository.delete_project_vector_rows.assert_awaited() + assert any( + call.kwargs["strict_adapter_cleanup"] is True and call.kwargs["session"] is not None + for call in search_repository.delete_project_vector_rows.await_args_list + ) + async with db.scoped_session(service.session_maker) as session: + assert await service.repository.get_by_name(session, project_name) is None @pytest.mark.asyncio @@ -160,6 +203,9 @@ async def test_remove_project_composes_vector_cleanup_without_injected_factory( create_search_repository.call_args.kwargs["session_maker"] is project_service.session_maker ) assert create_search_repository.call_args.kwargs["project_id"] == project_id - search_repository.delete_project_vector_rows.assert_awaited_once_with( - strict_adapter_cleanup=True + search_repository.delete_project_vector_rows.assert_awaited_once() + assert ( + search_repository.delete_project_vector_rows.await_args.kwargs["strict_adapter_cleanup"] + is True ) + assert search_repository.delete_project_vector_rows.await_args.kwargs["session"] is not None From 93715b2d916e2e0475dad8531601516be7e7937e Mon Sep 17 00:00:00 2001 From: phernandez Date: Sun, 26 Jul 2026 15:58:39 -0500 Subject: [PATCH 22/24] test(core): load sqlite-vec for status assertion Signed-off-by: phernandez --- test-int/test_embedding_status_vec0.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/test-int/test_embedding_status_vec0.py b/test-int/test_embedding_status_vec0.py index 5c6b9bfc8..d9fb3cf0c 100644 --- a/test-int/test_embedding_status_vec0.py +++ b/test-int/test_embedding_status_vec0.py @@ -146,6 +146,9 @@ async def test_embedding_status_reads_real_vec0_table(engine_factory, test_proje ] ) async with db.scoped_session(session_maker) as session: + # sqlite-vec is loaded per connection. Windows may hand this assertion a + # different pooled connection than the adapter used for the upsert. + await search_repo._ensure_sqlite_vec_loaded(session) stale_count = await session.execute(text("SELECT COUNT(*) FROM search_vector_embeddings")) assert stale_count.scalar_one() == 0 From 37074e361cb489ce32fe73d3c8c67eb071be4497 Mon Sep 17 00:00:00 2001 From: phernandez Date: Sun, 26 Jul 2026 16:16:09 -0500 Subject: [PATCH 23/24] fix(core): serialize external vector reconciliation Signed-off-by: phernandez --- .../repository/search_repository_base.py | 14 +++++ tests/repository/test_semantic_search_base.py | 58 +++++++++++++++++++ 2 files changed, 72 insertions(+) diff --git a/src/basic_memory/repository/search_repository_base.py b/src/basic_memory/repository/search_repository_base.py index 3909ebc4c..9882ef04c 100644 --- a/src/basic_memory/repository/search_repository_base.py +++ b/src/basic_memory/repository/search_repository_base.py @@ -1173,7 +1173,16 @@ async def reconcile_vector_index(self) -> None: if not isinstance(self._semantic_vector_index, SemanticVectorIndexReconciler): return + external_vector_index = self._uses_external_vector_index() async with db.scoped_session(self.session_maker) as session: + if external_vector_index: + # Trigger: reconciliation snapshots the live manifest before asking + # an external adapter to delete everything else. + # Why: a concurrent watcher could otherwise publish a new vector + # after the snapshot and have reconciliation delete that live key. + # Outcome: share the project lock with external writes through both + # the manifest read and orphan deletion. + await self._lock_external_vector_write(session) result = await session.execute( text( "SELECT entity_id, chunk_key FROM search_vector_chunks " @@ -1197,6 +1206,11 @@ async def reconcile_vector_index(self) -> None: for row in result.mappings().all() ] + if external_vector_index: + await self._semantic_vector_index.delete_orphans(live_keys) + await session.commit() + return + await self._semantic_vector_index.delete_orphans(live_keys) # ------------------------------------------------------------------ diff --git a/tests/repository/test_semantic_search_base.py b/tests/repository/test_semantic_search_base.py index 78753f38f..5d96f8246 100644 --- a/tests/repository/test_semantic_search_base.py +++ b/tests/repository/test_semantic_search_base.py @@ -393,6 +393,64 @@ async def fake_scoped_session(_session_maker): assert adapter.upserted_records[0].source_hash == source_hash +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("dialect_name", "project_lock_prefix"), + [ + ("postgresql", "SELECT id FROM project"), + ("sqlite", "UPDATE project SET id = id"), + ], +) +async def test_external_reconciliation_holds_project_lock_through_orphan_cleanup( + monkeypatch, + dialect_name, + project_lock_prefix, +) -> None: + """External reconciliation must serialize its snapshot and orphan deletion.""" + repo = _ConcreteRepo() + repo._semantic_enabled = True + repo._semantic_vector_index_name = "milvus" + repo._embedding_provider = SimpleNamespace(model_name="stub", dimensions=4) + events: list[str] = [] + adapter: Any = SimpleNamespace( + scope=_RecordingVectorIndex.scope, + delete_orphans=AsyncMock(side_effect=lambda _live_keys: events.append("delete_orphans")), + ) + repo._semantic_vector_index = adapter + session = AsyncMock() + session.connection.return_value = SimpleNamespace(dialect=SimpleNamespace(name=dialect_name)) + + def execute(statement, _params): + sql = str(statement) + if sql.startswith(project_lock_prefix): + events.append("project_lock") + return SimpleNamespace() + if sql.startswith("SELECT entity_id, chunk_key"): + events.append("manifest_read") + return SimpleNamespace( + mappings=lambda: SimpleNamespace( + all=lambda: [{"entity_id": 41, "chunk_key": "entity:41:0"}] + ) + ) + raise AssertionError(f"Unexpected SQL: {sql}") + + session.execute.side_effect = execute + session.commit.side_effect = lambda: events.append("commit") + + @asynccontextmanager + async def fake_scoped_session(_session_maker): + yield session + + monkeypatch.setattr(search_repository_base_module.db, "scoped_session", fake_scoped_session) + + await repo.reconcile_vector_index() + + assert events == ["project_lock", "manifest_read", "delete_orphans", "commit"] + adapter.delete_orphans.assert_awaited_once_with( + [VectorKey(entity_id=41, chunk_key="entity:41:0")] + ) + + # --- SQLite SemanticSearchDisabledError --- From 582b94aa14847c8c3de91de0009691c1f1b742e8 Mon Sep 17 00:00:00 2001 From: phernandez Date: Sun, 26 Jul 2026 16:36:41 -0500 Subject: [PATCH 24/24] fix(core): verify physical vectors in project status Signed-off-by: phernandez --- src/basic_memory/services/project_service.py | 159 +++++++++--- test-int/test_embedding_status_vec0.py | 10 +- .../test_project_service_embedding_status.py | 235 +++++++++++++++--- 3 files changed, 326 insertions(+), 78 deletions(-) diff --git a/src/basic_memory/services/project_service.py b/src/basic_memory/services/project_service.py index e3a8875ac..fd32f4c83 100644 --- a/src/basic_memory/services/project_service.py +++ b/src/basic_memory/services/project_service.py @@ -12,7 +12,9 @@ from loguru import logger from sqlalchemy import text +from sqlalchemy.exc import OperationalError as SAOperationalError from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker +from sqlalchemy.sql.base import Executable from basic_memory import db from basic_memory.models import Project @@ -1083,27 +1085,35 @@ async def get_embedding_status(self, project_id: int) -> EmbeddingStatus: is_postgres = config.database_backend == DatabaseBackend.POSTGRES vector_index = resolve_semantic_vector_index_name(config, config.database_backend) embedding_identity = configured_embedding_provider_identity(config) + uses_builtin_vector_storage = vector_index in {"pgvector", "sqlite-vec"} - # --- Check vector manifest existence --- - # The SQL manifest is authoritative even when vector values live outside - # the database, so status never probes backend-specific storage tables. + # --- Check vector manifest and built-in storage existence --- + # External adapters expose no portable storage-inspection contract, so + # their status remains manifest-only. The built-ins are SQL-backed and + # must prove that the physical vector table still exists. if is_postgres: table_check_sql = text( - "SELECT COUNT(*) FROM information_schema.tables " - "WHERE table_name = 'search_vector_chunks'" + "SELECT table_name FROM information_schema.tables " + "WHERE table_name IN ('search_vector_chunks', 'search_vector_embeddings')" ) else: table_check_sql = text( - "SELECT COUNT(*) FROM sqlite_master " - "WHERE type = 'table' AND name = 'search_vector_chunks'" + "SELECT name FROM sqlite_master " + "WHERE type = 'table' " + "AND name IN ('search_vector_chunks', 'search_vector_embeddings')" ) async with db.scoped_session(self.session_maker) as session: table_result = await self.repository.execute_query(session, table_check_sql, {}) - vector_tables_exist = (table_result.scalar() or 0) == 1 + existing_vector_tables = {str(name) for name in table_result.scalars().all()} + manifest_exists = "search_vector_chunks" in existing_vector_tables + storage_exists = "search_vector_embeddings" in existing_vector_tables + vector_tables_exist = manifest_exists and ( + storage_exists or not uses_builtin_vector_storage + ) - manifest_schema_current = vector_tables_exist - if vector_tables_exist and not is_postgres: + manifest_schema_current = manifest_exists + if manifest_exists and not is_postgres: columns_result = await self.repository.execute_query( session, text("PRAGMA table_info(search_vector_chunks)"), @@ -1116,7 +1126,7 @@ async def get_embedding_status(self, project_id: int) -> EmbeddingStatus: "embedding_status", }.issubset(manifest_columns) - if not manifest_schema_current: + if not manifest_schema_current or not vector_tables_exist: # Count distinct entities in search index for the recommendation message si_result = await self.repository.execute_query( session, @@ -1128,6 +1138,17 @@ async def get_embedding_status(self, project_id: int) -> EmbeddingStatus: ) total_indexed_entities = si_result.scalar() or 0 + if manifest_exists and not manifest_schema_current: + reindex_reason = ( + "Vector manifest schema is outdated — run: bm reindex --embeddings" + ) + elif manifest_schema_current: + reindex_reason = "Vector storage not initialized — run: bm reindex --embeddings" + else: + reindex_reason = ( + "Vector manifest not initialized — run: bm reindex --embeddings" + ) + return EmbeddingStatus( semantic_search_enabled=True, embedding_provider=provider, @@ -1138,11 +1159,7 @@ async def get_embedding_status(self, project_id: int) -> EmbeddingStatus: total_indexed_entities=total_indexed_entities, vector_tables_exist=False, reindex_recommended=True, - reindex_reason=( - "Vector manifest schema is outdated — run: bm reindex --embeddings" - if vector_tables_exist - else "Vector manifest not initialized — run: bm reindex --embeddings" - ), + reindex_reason=reindex_reason, ) # --- Count queries (manifest exists) --- @@ -1151,6 +1168,9 @@ async def get_embedding_status(self, project_id: int) -> EmbeddingStatus: entity_exists = ( "AND entity_id IN (SELECT id FROM entity WHERE project_id = :project_id)" ) + chunk_entity_exists = ( + "AND c.entity_id IN (SELECT id FROM entity WHERE project_id = :project_id)" + ) si_result = await self.repository.execute_query( session, @@ -1172,6 +1192,11 @@ async def get_embedding_status(self, project_id: int) -> EmbeddingStatus: "AND embedding_model = :embedding_identity " "AND embedding_status = 'ready'" ) + current_ready_chunk = ( + "c.vector_index = :vector_index " + "AND c.embedding_model = :embedding_identity " + "AND c.embedding_status = 'ready'" + ) chunks_result = await self.repository.execute_query( session, @@ -1193,25 +1218,91 @@ async def get_embedding_status(self, project_id: int) -> EmbeddingStatus: ) total_entities_with_chunks = entities_with_chunks_result.scalar() or 0 - embeddings_result = await self.repository.execute_query( - session, - text( - "SELECT COUNT(*) FROM search_vector_chunks " - f"WHERE project_id = :project_id AND {current_ready} {entity_exists}" - ), - manifest_params, - ) - total_embeddings = embeddings_result.scalar() or 0 + if uses_builtin_vector_storage: + physical_id = "e.chunk_id" if is_postgres else "e.rowid" + embeddings_sql = text( + "SELECT COUNT(*) FROM search_vector_chunks c " + f"JOIN search_vector_embeddings e ON {physical_id} = c.id " + "AND e.source_hash = c.source_hash " + f"WHERE c.project_id = :project_id AND {current_ready_chunk} " + f"{chunk_entity_exists}" + ) + orphaned_sql = text( + "SELECT COUNT(*) FROM search_vector_chunks c " + f"LEFT JOIN search_vector_embeddings e ON {physical_id} = c.id " + "AND e.source_hash = c.source_hash " + "WHERE c.project_id = :project_id " + f"AND (NOT ({current_ready_chunk}) OR {physical_id} IS NULL) " + f"{chunk_entity_exists}" + ) - orphaned_result = await self.repository.execute_query( - session, - text( - "SELECT COUNT(*) FROM search_vector_chunks " - f"WHERE project_id = :project_id AND NOT ({current_ready}) {entity_exists}" - ), - manifest_params, - ) - orphaned_chunks = orphaned_result.scalar() or 0 + async def _physical_vector_count(query: Executable) -> int: + if is_postgres: + result = await self.repository.execute_query( + session, + query, + manifest_params, + ) + return result.scalar() or 0 + count = await self.repository.scalar_vec_query( + session, + query, + manifest_params, + ) + if count is None: + raise SAOperationalError( + str(query), + {}, + Exception("no such module: vec0"), + ) + return count + + try: + total_embeddings = await _physical_vector_count(embeddings_sql) + orphaned_chunks = await _physical_vector_count(orphaned_sql) + except SAOperationalError as exc: + # Trigger: sqlite_master can list vec0 storage even though this + # Python runtime cannot load the sqlite-vec extension. + # Why: ready manifest rows do not prove vectors are searchable. + # Outcome: surface the missing dependency and require a rebuild. + if is_postgres or "no such module: vec0" not in str(exc).lower(): + raise + return EmbeddingStatus( + semantic_search_enabled=True, + embedding_provider=provider, + embedding_model=model, + embedding_dimensions=dimensions, + embedding_document_prefix_set=document_prefix_set, + embedding_query_prefix_set=query_prefix_set, + total_indexed_entities=total_indexed_entities, + vector_tables_exist=False, + reindex_recommended=True, + reindex_reason=( + "SQLite vector tables exist but sqlite-vec is unavailable in this " + "Python environment — install/update basic-memory, then run: " + "bm reindex --embeddings" + ), + ) + else: + embeddings_result = await self.repository.execute_query( + session, + text( + "SELECT COUNT(*) FROM search_vector_chunks " + f"WHERE project_id = :project_id AND {current_ready} {entity_exists}" + ), + manifest_params, + ) + total_embeddings = embeddings_result.scalar() or 0 + + orphaned_result = await self.repository.execute_query( + session, + text( + "SELECT COUNT(*) FROM search_vector_chunks " + f"WHERE project_id = :project_id AND NOT ({current_ready}) {entity_exists}" + ), + manifest_params, + ) + orphaned_chunks = orphaned_result.scalar() or 0 # --- Reindex recommendation logic (priority order) --- reindex_recommended = False diff --git a/test-int/test_embedding_status_vec0.py b/test-int/test_embedding_status_vec0.py index d9fb3cf0c..1ccbee5cc 100644 --- a/test-int/test_embedding_status_vec0.py +++ b/test-int/test_embedding_status_vec0.py @@ -1,12 +1,12 @@ -"""Integration regression test for manifest-backed status with a real vec0 table. +"""Integration regression test for built-in storage status with a real vec0 table. Regression for #658: after a successful `bm reindex --embeddings`, `bm project info` still reported "sqlite-vec is unavailable", "Indexed 0/N", and "Chunks 0", and recommended an unnecessary reindex. -Status now reads the authoritative SQL manifest instead of reopening backend-specific -vector storage. This test builds and writes a real vec0 table, then proves a fresh project -status connection can report the ready manifest without loading sqlite-vec. +Status loads sqlite-vec on the same connection used to inspect the physical table. This +test builds and writes a real vec0 table, then proves a fresh status connection can verify +the ready manifest and its matching vector without a false unavailable result. """ import os @@ -194,7 +194,7 @@ def _config_manager_semantic_enabled() -> ConfigManager: status = await project_service.get_embedding_status(project_id) assert status.semantic_search_enabled is True - # Status reads the ready manifest, so it never needs to reopen the vec0 table. + # Status reloads sqlite-vec on this fresh connection and verifies the physical row. assert status.vector_tables_exist is True assert status.reindex_recommended is False assert status.reindex_reason is None diff --git a/tests/services/test_project_service_embedding_status.py b/tests/services/test_project_service_embedding_status.py index 42ad36831..e86d75c9c 100644 --- a/tests/services/test_project_service_embedding_status.py +++ b/tests/services/test_project_service_embedding_status.py @@ -29,6 +29,29 @@ async def _execute(project_service: ProjectService, query, params=None): return await project_service.repository.execute_query(session, query, params or {}) +async def _create_embeddings_stub(project_service: ProjectService) -> None: + """Create portable built-in storage for status tests.""" + await _execute( + project_service, + text( + "CREATE TABLE IF NOT EXISTS search_vector_embeddings (" + "chunk_id INTEGER PRIMARY KEY, source_hash TEXT NOT NULL)" + ), + {}, + ) + + +async def _drop_embeddings_stub(project_service: ProjectService) -> None: + """Remove portable built-in storage created by a status test.""" + await _execute(project_service, text("DROP TABLE IF EXISTS search_vector_embeddings"), {}) + + +async def _scalar_regular_query(session, query, params=None): + """Execute a vector count against the portable regular-table test double.""" + result = await session.execute(query, params or {}) + return result.scalar() + + @pytest.mark.asyncio async def test_embedding_status_semantic_disabled(project_service: ProjectService, test_project): """When semantic search is disabled, return minimal status with zero counts.""" @@ -177,14 +200,23 @@ async def test_embedding_status_entities_without_chunks( project_service: ProjectService, test_graph, test_project ): """When entities have search_index rows but no chunks, recommend reindex.""" - with patch.object( - type(project_service), - "config_manager", - new_callable=lambda: property( - lambda self: _config_manager_with(semantic_search_enabled=True) + await _create_embeddings_stub(project_service) + with ( + patch.object( + type(project_service), + "config_manager", + new_callable=lambda: property( + lambda self: _config_manager_with(semantic_search_enabled=True) + ), + ), + patch.object( + project_service.repository, + "scalar_vec_query", + side_effect=_scalar_regular_query, ), ): status = await project_service.get_embedding_status(test_project.id) + await _drop_embeddings_stub(project_service) assert status.semantic_search_enabled is True assert status.vector_tables_exist is True @@ -199,7 +231,7 @@ async def test_embedding_status_entities_without_chunks( async def test_embedding_status_orphaned_chunks( project_service: ProjectService, test_graph, test_project ): - """When a manifest row is pending, recommend reindex.""" + """A ready manifest row without its physical vector must recommend reindex.""" # Get a real entity_id from the test graph entity_result = await _execute( project_service, @@ -213,30 +245,39 @@ async def test_embedding_status_orphaned_chunks( entity_id=entity_id, project_id=test_project.id, chunk_key="chunk-1", - embedding_status="pending", ) + await _create_embeddings_stub(project_service) - with patch.object( - type(project_service), - "config_manager", - new_callable=lambda: property( - lambda self: _config_manager_with(semantic_search_enabled=True) + with ( + patch.object( + type(project_service), + "config_manager", + new_callable=lambda: property( + lambda self: _config_manager_with(semantic_search_enabled=True) + ), + ), + patch.object( + project_service.repository, + "scalar_vec_query", + side_effect=_scalar_regular_query, ), ): status = await project_service.get_embedding_status(test_project.id) + await _drop_embeddings_stub(project_service) assert status.vector_tables_exist is True assert status.total_chunks == 1 + assert status.total_embeddings == 0 assert status.orphaned_chunks == 1 assert status.reindex_recommended is True assert "need vector indexing" in (status.reindex_reason or "") @pytest.mark.asyncio -async def test_embedding_status_counts_only_current_ready_manifest_rows( +async def test_embedding_status_external_index_counts_only_current_ready_manifest_rows( project_service: ProjectService, test_graph, test_project ): - """Status is backend-neutral and treats pending or mismatched rows as stale.""" + """External indexes remain manifest-only because their storage is not inspectable.""" entity_result = await _execute( project_service, text("SELECT id FROM entity WHERE project_id = :project_id LIMIT 1"), @@ -249,12 +290,14 @@ async def test_embedding_status_counts_only_current_ready_manifest_rows( entity_id=entity_id, project_id=test_project.id, chunk_key="ready", + vector_index="milvus", ) await _insert_manifest_chunk( project_service, entity_id=entity_id, project_id=test_project.id, chunk_key="pending", + vector_index="milvus", embedding_status="pending", ) await _insert_manifest_chunk( @@ -262,7 +305,7 @@ async def test_embedding_status_counts_only_current_ready_manifest_rows( entity_id=entity_id, project_id=test_project.id, chunk_key="wrong-index", - vector_index="milvus", + vector_index="pgvector", ) await _insert_manifest_chunk( project_service, @@ -272,19 +315,25 @@ async def test_embedding_status_counts_only_current_ready_manifest_rows( embedding_identity="OtherProvider:other:384", ) - with patch.object( - type(project_service), - "config_manager", - new_callable=lambda: property( - lambda self: _config_manager_with(semantic_search_enabled=True) + with ( + patch.object( + type(project_service), + "config_manager", + new_callable=lambda: property( + lambda self: _config_manager_with(semantic_search_enabled=True) + ), ), - ): - with patch.object( + patch( + "basic_memory.services.project_service.resolve_semantic_vector_index_name", + return_value="milvus", + ), + patch.object( project_service.repository, "scalar_vec_query", side_effect=AssertionError("status must not query vector storage"), - ): - status = await project_service.get_embedding_status(test_project.id) + ), + ): + status = await project_service.get_embedding_status(test_project.id) assert status.semantic_search_enabled is True assert status.vector_tables_exist is True @@ -295,13 +344,94 @@ async def test_embedding_status_counts_only_current_ready_manifest_rows( assert "pending or stale" in (status.reindex_reason or "") +@pytest.mark.asyncio +async def test_embedding_status_reports_missing_builtin_storage( + project_service: ProjectService, + test_graph, + test_project, +): + """A surviving ready manifest must not hide a lost built-in storage table.""" + await _drop_embeddings_stub(project_service) + entity_result = await _execute( + project_service, + text("SELECT id FROM entity WHERE project_id = :project_id LIMIT 1"), + {"project_id": test_project.id}, + ) + await _insert_manifest_chunk( + project_service, + entity_id=entity_result.scalar(), + project_id=test_project.id, + chunk_key="ready-without-storage", + ) + + with patch.object( + type(project_service), + "config_manager", + new_callable=lambda: property( + lambda self: _config_manager_with(semantic_search_enabled=True) + ), + ): + status = await project_service.get_embedding_status(test_project.id) + + assert status.vector_tables_exist is False + assert status.reindex_recommended is True + assert "Vector storage not initialized" in (status.reindex_reason or "") + + +@pytest.mark.asyncio +async def test_embedding_status_handles_sqlite_vec_unavailable( + project_service: ProjectService, + test_graph, + test_project, +): + """Ready manifests must not look healthy when sqlite-vec cannot load.""" + if _is_postgres(): + pytest.skip("sqlite-vec unavailable handling is SQLite-specific.") + + entity_result = await _execute( + project_service, + text("SELECT id FROM entity WHERE project_id = :project_id LIMIT 1"), + {"project_id": test_project.id}, + ) + await _insert_manifest_chunk( + project_service, + entity_id=entity_result.scalar(), + project_id=test_project.id, + chunk_key="ready-without-extension", + ) + await _create_embeddings_stub(project_service) + + with ( + patch.object( + type(project_service), + "config_manager", + new_callable=lambda: property( + lambda self: _config_manager_with(semantic_search_enabled=True) + ), + ), + patch.object( + project_service.repository, + "scalar_vec_query", + return_value=None, + ), + ): + status = await project_service.get_embedding_status(test_project.id) + await _drop_embeddings_stub(project_service) + + assert status.semantic_search_enabled is True + assert status.vector_tables_exist is False + assert status.reindex_recommended is True + assert "sqlite-vec is unavailable" in (status.reindex_reason or "") + + @pytest.mark.asyncio async def test_embedding_status_healthy(project_service: ProjectService, test_graph, test_project): """When all entities have embeddings, no reindex recommended.""" # Clear any leftover data from prior tests await _execute(project_service, text("DELETE FROM search_vector_chunks"), {}) + await _create_embeddings_stub(project_service) - # Insert a current, ready manifest row for every search_index entity. + # Insert a current ready manifest row and matching physical vector for every entity. entity_result = await _execute( project_service, text("SELECT DISTINCT entity_id FROM search_index WHERE project_id = :project_id"), @@ -311,22 +441,38 @@ async def test_embedding_status_healthy(project_service: ProjectService, test_gr chunk_id = 1 for eid in entity_ids: - await _insert_manifest_chunk( + inserted_chunk_id = await _insert_manifest_chunk( project_service, entity_id=eid, project_id=test_project.id, chunk_key=f"chunk-{chunk_id}", ) + await _execute( + project_service, + text( + "INSERT INTO search_vector_embeddings (chunk_id, source_hash) " + "VALUES (:chunk_id, 'hash')" + ), + {"chunk_id": inserted_chunk_id}, + ) chunk_id += 1 - with patch.object( - type(project_service), - "config_manager", - new_callable=lambda: property( - lambda self: _config_manager_with(semantic_search_enabled=True) + with ( + patch.object( + type(project_service), + "config_manager", + new_callable=lambda: property( + lambda self: _config_manager_with(semantic_search_enabled=True) + ), + ), + patch.object( + project_service.repository, + "scalar_vec_query", + side_effect=_scalar_regular_query, ), ): status = await project_service.get_embedding_status(test_project.id) + await _drop_embeddings_stub(project_service) assert status.vector_tables_exist is True assert status.total_chunks > 0 @@ -349,6 +495,7 @@ async def test_embedding_status_excludes_stale_entity_ids( # Include 'id' column — required NOT NULL on Postgres (regular table), # ignored on SQLite (FTS5 virtual table where id is UNINDEXED). stale_entity_id = 999999 + await _create_embeddings_stub(project_service) await _execute( project_service, text( @@ -361,14 +508,22 @@ async def test_embedding_status_excludes_stale_entity_ids( {"id": stale_entity_id, "eid": stale_entity_id, "pid": test_project.id}, ) - with patch.object( - type(project_service), - "config_manager", - new_callable=lambda: property( - lambda self: _config_manager_with(semantic_search_enabled=True) + with ( + patch.object( + type(project_service), + "config_manager", + new_callable=lambda: property( + lambda self: _config_manager_with(semantic_search_enabled=True) + ), + ), + patch.object( + project_service.repository, + "scalar_vec_query", + side_effect=_scalar_regular_query, ), ): status = await project_service.get_embedding_status(test_project.id) + await _drop_embeddings_stub(project_service) # The stale entity_id should NOT be counted in total_indexed_entities. # Count real entities that have search_index rows (the stale one should be excluded). @@ -419,19 +574,20 @@ async def _insert_manifest_chunk( vector_index: str | None = None, embedding_identity: str | None = None, embedding_status: str = "ready", -) -> None: +) -> int: """Insert one manifest row with explicit backend and readiness identity.""" config = _config_manager_with(semantic_search_enabled=True).config active_vector_index = resolve_semantic_vector_index_name(config, config.database_backend) active_embedding_identity = configured_embedding_provider_identity(config) - await _execute( + result = await _execute( project_service, text( "INSERT INTO search_vector_chunks " "(entity_id, project_id, chunk_key, chunk_text, source_hash, " "entity_fingerprint, embedding_model, vector_index, embedding_status) " "VALUES (:entity_id, :project_id, :chunk_key, 'test text', 'hash', " - "'fingerprint', :embedding_identity, :vector_index, :embedding_status)" + "'fingerprint', :embedding_identity, :vector_index, :embedding_status) " + "RETURNING id" ), { "entity_id": entity_id, @@ -442,3 +598,4 @@ async def _insert_manifest_chunk( "embedding_status": embedding_status, }, ) + return int(result.scalar_one())