diff --git a/src/basic_memory/repository/postgres_search_repository.py b/src/basic_memory/repository/postgres_search_repository.py index 487fc8e9e..0d6a2c867 100644 --- a/src/basic_memory/repository/postgres_search_repository.py +++ b/src/basic_memory/repository/postgres_search_repository.py @@ -28,7 +28,10 @@ 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_sync import ( + PendingEmbeddingJob, + StagedVectorDeletion, +) from basic_memory.repository.semantic_vector_index_factory import ( build_vector_index_scope, resolve_semantic_vector_index_name, @@ -412,7 +415,7 @@ async def _upsert_scheduled_chunk_records( existing_by_key: dict[str, VectorChunkState], entity_fingerprint: str, embedding_model: str, - ) -> list[tuple[int, str]]: + ) -> list[PendingEmbeddingJob]: """Use Postgres UPSERT to rewrite only the scheduled chunk rows.""" if not scheduled_records: return [] @@ -475,7 +478,13 @@ async def _upsert_scheduled_chunk_records( str(row["chunk_key"]): int(row["id"]) for row in upsert_result.mappings().all() } return [ - (upserted_ids_by_key[record["chunk_key"]], record["chunk_text"]) + PendingEmbeddingJob( + entity_id=entity_id, + chunk_row_id=upserted_ids_by_key[record["chunk_key"]], + chunk_key=record["chunk_key"], + chunk_text=record["chunk_text"], + source_hash=record["source_hash"], + ) for record in scheduled_records ] diff --git a/src/basic_memory/repository/search_repository_base.py b/src/basic_memory/repository/search_repository_base.py index a24227eda..26bd37d19 100644 --- a/src/basic_memory/repository/search_repository_base.py +++ b/src/basic_memory/repository/search_repository_base.py @@ -48,6 +48,7 @@ VectorRecord, ) from basic_memory.repository.semantic_vector_sync import ( + EmbeddingPersistenceResult as _EmbeddingPersistenceResult, EntitySyncRuntime as _EntitySyncRuntime, EntityVectorShardPlan as _EntityVectorShardPlan, PendingEmbeddingJob as _PendingEmbeddingJob, @@ -338,12 +339,23 @@ async def _write_embeddings( async def _persist_embeddings( self, - jobs: list[tuple[int, str]], + jobs: Sequence[_PendingEmbeddingJob], embeddings: list[list[float]], - ) -> None: + ) -> _EmbeddingPersistenceResult: """Write vectors through the adapter, then make their manifest rows ready.""" if not jobs: - return + return _EmbeddingPersistenceResult() + if len(jobs) != len(embeddings): + raise RuntimeError("Embedding provider returned an unexpected number of vectors.") + + for job in jobs: + expected_source_hash = hashlib.sha256(job.chunk_text.encode("utf-8")).hexdigest() + if job.source_hash != expected_source_hash: + raise RuntimeError( + f"Embedding job source hash does not match its chunk text: {job.chunk_row_id}" + ) + + job_pairs = [(job.chunk_row_id, job.chunk_text) for job in jobs] # Compatibility: focused orchestration tests and third-party subclasses # from before the adapter contract may still override the private writer. @@ -352,11 +364,13 @@ async def _persist_embeddings( 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 self._write_embeddings(session, job_pairs, embeddings) await session.commit() - return + return _EmbeddingPersistenceResult( + persisted_row_ids=frozenset(row_id for row_id, _chunk_text in job_pairs) + ) - row_ids = [row_id for row_id, _ in jobs] + row_ids = [row_id for row_id, _chunk_text in job_pairs] 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: @@ -392,41 +406,74 @@ async def _persist_embeddings( ) 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: + missing_jobs = [job for job in jobs if job.chunk_row_id not in rows_by_id] + superseded_row_ids: set[int] = set() + missing_current_row_ids: list[int] = [] + if missing_jobs: + entity_ids = sorted({job.entity_id for job in missing_jobs}) + source_rows_by_entity = await self._fetch_prepare_window_source_rows( + session, + entity_ids, + ) + current_generations = { + (entity_id, record["chunk_key"], record["source_hash"]) + for entity_id, source_rows in source_rows_by_entity.items() + for record in self._build_chunk_records(source_rows) + } + for job in missing_jobs: + generation = (job.entity_id, job.chunk_key, job.source_hash) + if generation in current_generations: + missing_current_row_ids.append(job.chunk_row_id) + else: + superseded_row_ids.add(job.chunk_row_id) + + if missing_current_row_ids: raise RuntimeError( - f"Vector manifest rows disappeared before write: {missing_row_ids}" + "Vector manifest rows disappeared before write: " + f"{sorted(missing_current_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: + current_jobs: list[tuple[_PendingEmbeddingJob, list[float]]] = [] + for job, embedding in zip(jobs, embeddings, strict=True): + row = rows_by_id.get(job.chunk_row_id) + if row is None: continue - current_jobs.append((row_id, chunk_text, expected_source_hash, embedding)) + if int(row["entity_id"]) != job.entity_id or str(row["chunk_key"]) != job.chunk_key: + raise RuntimeError( + f"Vector manifest row identity changed before write: {job.chunk_row_id}" + ) + if str(row["source_hash"]) != job.source_hash: + superseded_row_ids.add(job.chunk_row_id) + continue + current_jobs.append((job, embedding)) if not current_jobs: - return + return _EmbeddingPersistenceResult(superseded_row_ids=frozenset(superseded_row_ids)) 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"]), + entity_id=job.entity_id, + chunk_key=job.chunk_key, ), - source_hash=source_hash, + source_hash=job.source_hash, values=tuple(embedding), ) - for row_id, _chunk_text, source_hash, embedding in current_jobs + for job, 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 + for index, (job, _embedding) in enumerate(current_jobs): + params[f"row_id_{index}"] = job.chunk_row_id + params[f"source_hash_{index}"] = job.source_hash generation_predicates.append( f"(id = :row_id_{index} AND source_hash = :source_hash_{index})" ) + persistence = _EmbeddingPersistenceResult( + persisted_row_ids=frozenset(job.chunk_row_id for job, _embedding in current_jobs), + superseded_row_ids=frozenset(superseded_row_ids), + ) + if lock_external_write: # Constraint: extension adapters use stable logical keys outside # the authoritative SQL database. Hold its manifest lock across @@ -439,7 +486,7 @@ async def _persist_embeddings( generation_predicates=generation_predicates, ) await session.commit() - return + return persistence # Built-in adapters share the authoritative database. They verify and lock # each record's source_hash inside the same transaction as their vector write. @@ -451,6 +498,7 @@ async def _persist_embeddings( generation_predicates=generation_predicates, ) await session.commit() + return persistence async def _mark_embedding_jobs_ready( self, @@ -1390,7 +1438,7 @@ async def _upsert_scheduled_chunk_records( existing_by_key: dict[str, VectorChunkState], entity_fingerprint: str, embedding_model: str, - ) -> list[tuple[int, str]]: + ) -> list[_PendingEmbeddingJob]: """Upsert scheduled chunk rows and return embedding jobs.""" return await semantic_vector_sync.upsert_scheduled_chunk_records( self, diff --git a/src/basic_memory/repository/semantic_vector_sync.py b/src/basic_memory/repository/semantic_vector_sync.py index cb07d349a..c2297d0e6 100644 --- a/src/basic_memory/repository/semantic_vector_sync.py +++ b/src/basic_memory/repository/semantic_vector_sync.py @@ -72,7 +72,7 @@ class PreparedEntityVectorSync: entity_id: int sync_start: float source_rows_count: int - embedding_jobs: list[tuple[int, str]] + embedding_jobs: list[PendingEmbeddingJob] chunks_total: int = 0 chunks_skipped: int = 0 entity_skipped: bool = False @@ -146,13 +146,23 @@ class UpsertEntityVectorPreparePlan: type EntityVectorPreparePlan = DeleteEntityVectorPreparePlan | UpsertEntityVectorPreparePlan -@dataclass +@dataclass(frozen=True, slots=True) class PendingEmbeddingJob: """Pending embedding write entry with entity ownership metadata.""" entity_id: int chunk_row_id: int + chunk_key: str chunk_text: str + source_hash: str + + +@dataclass(frozen=True, slots=True) +class EmbeddingPersistenceResult: + """Manifest generations persisted or superseded during one adapter write.""" + + persisted_row_ids: frozenset[int] = frozenset() + superseded_row_ids: frozenset[int] = frozenset() @dataclass @@ -176,6 +186,7 @@ class EntitySyncRuntime: prepare_seconds: float = 0.0 embed_seconds: float = 0.0 write_seconds: float = 0.0 + superseded: bool = False @dataclass(frozen=True) @@ -396,14 +407,7 @@ def emit_progress(entity_id: int) -> None: remaining_jobs_after_shard=prepared.remaining_jobs_after_shard, prepare_seconds=prepared.prepare_seconds, ) - pending_jobs.extend( - PendingEmbeddingJob( - entity_id=entity_id, - chunk_row_id=row_id, - chunk_text=chunk_text, - ) - for row_id, chunk_text in prepared.embedding_jobs - ) + pending_jobs.extend(prepared.embedding_jobs) while len(pending_jobs) >= repository._semantic_embedding_sync_batch_size: flush_jobs = pending_jobs[: repository._semantic_embedding_sync_batch_size] @@ -1074,7 +1078,7 @@ async def apply_entity_vector_prepare_plan( }, ) - embedding_jobs: list[tuple[int, str]] = [] + embedding_jobs: list[PendingEmbeddingJob] = [] if plan.scheduled_records: embedding_jobs = await repository._upsert_scheduled_chunk_records( session, @@ -1115,7 +1119,7 @@ async def upsert_scheduled_chunk_records( existing_by_key: dict[str, VectorChunkState], entity_fingerprint: str, embedding_model: str, -) -> list[tuple[int, str]]: +) -> list[PendingEmbeddingJob]: """Upsert scheduled chunk rows and return embedding jobs.""" repository._assert_manifest_vector_ownership( current.vector_index @@ -1123,7 +1127,7 @@ async def upsert_scheduled_chunk_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]] = [] + embedding_jobs: list[PendingEmbeddingJob] = [] for record in scheduled_records: current = existing_by_key.get(record["chunk_key"]) if current: @@ -1154,7 +1158,15 @@ async def upsert_scheduled_chunk_records( "vector_index": repository._semantic_vector_index_name, }, ) - embedding_jobs.append((current.id, record["chunk_text"])) + embedding_jobs.append( + PendingEmbeddingJob( + entity_id=entity_id, + chunk_row_id=current.id, + chunk_key=record["chunk_key"], + chunk_text=record["chunk_text"], + source_hash=record["source_hash"], + ) + ) continue inserted = await session.execute( @@ -1180,7 +1192,15 @@ async def upsert_scheduled_chunk_records( "vector_index": repository._semantic_vector_index_name, }, ) - embedding_jobs.append((int(inserted.scalar_one()), record["chunk_text"])) + embedding_jobs.append( + PendingEmbeddingJob( + entity_id=entity_id, + chunk_row_id=int(inserted.scalar_one()), + chunk_key=record["chunk_key"], + chunk_text=record["chunk_text"], + source_hash=record["source_hash"], + ) + ) return embedding_jobs @@ -1203,14 +1223,22 @@ async def flush_embedding_jobs( raise RuntimeError("Embedding provider returned an unexpected number of vectors.") write_start = time.perf_counter() - write_jobs = [(job.chunk_row_id, job.chunk_text) for job in flush_jobs] - await repository._persist_embeddings(write_jobs, embeddings) + persistence = await repository._persist_embeddings(flush_jobs, embeddings) write_seconds = time.perf_counter() - write_start + expected_row_ids = {job.chunk_row_id for job in flush_jobs} + classified_row_ids = persistence.persisted_row_ids | persistence.superseded_row_ids + if classified_row_ids != expected_row_ids: + raise RuntimeError("Embedding persistence did not classify every manifest row.") + flush_size = len(flush_jobs) entity_job_counts: dict[int, int] = {} for job in flush_jobs: entity_job_counts[job.entity_id] = entity_job_counts.get(job.entity_id, 0) + 1 + if job.chunk_row_id in persistence.superseded_row_ids: + runtime = entity_runtime.get(job.entity_id) + if runtime is not None: + runtime.superseded = True for entity_id, entity_job_count in entity_job_counts.items(): runtime = entity_runtime.get(entity_id) @@ -1223,7 +1251,7 @@ async def flush_embedding_jobs( runtime.embed_seconds += embed_seconds * flush_share runtime.write_seconds += write_seconds * flush_share - if runtime.remaining_jobs <= 0 and runtime.entity_complete: + if runtime.remaining_jobs <= 0 and runtime.entity_complete and not runtime.superseded: synced_entity_ids.add(entity_id) return embed_seconds, write_seconds @@ -1243,7 +1271,7 @@ def finalize_completed_entity_syncs( if runtime.remaining_jobs > 0: continue - if runtime.entity_complete: + if runtime.entity_complete and not runtime.superseded: synced_entity_ids.add(entity_id) else: deferred_entity_ids.add(entity_id) diff --git a/tests/repository/test_postgres_search_repository_unit.py b/tests/repository/test_postgres_search_repository_unit.py index 1303da36e..fecd64194 100644 --- a/tests/repository/test_postgres_search_repository_unit.py +++ b/tests/repository/test_postgres_search_repository_unit.py @@ -5,7 +5,9 @@ are difficult to reach in integration tests. """ +import hashlib from contextlib import asynccontextmanager +from typing import Any from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -23,7 +25,7 @@ SemanticSearchDisabledError, SemanticVectorIndexExtensionError, ) -from typing import Any +from basic_memory.repository.semantic_vector_sync import PendingEmbeddingJob # --- Helpers --------------------------------------------------------------- @@ -42,6 +44,16 @@ async def embed_documents(self, texts: list[str]) -> list[list[float]]: return [[0.0] * 4 for _ in texts] +def _pending_job(entity_id: int, row_id: int, chunk_text: str) -> PendingEmbeddingJob: + return PendingEmbeddingJob( + entity_id=entity_id, + chunk_row_id=row_id, + chunk_key=f"entity:{entity_id}:0", + chunk_text=chunk_text, + source_hash=hashlib.sha256(chunk_text.encode("utf-8")).hexdigest(), + ) + + def _make_repo( *, semantic_enabled: bool = False, @@ -424,7 +436,7 @@ async def _stub_prepare_window(entity_ids: list[int]): entity_id=entity_id, sync_start=0.0, source_rows_count=1, - embedding_jobs=[(200 + entity_id, f"chunk-{entity_id}")], + embedding_jobs=[_pending_job(entity_id, 200 + entity_id, f"chunk-{entity_id}")], prepare_seconds=1.0, ) for entity_id in entity_ids @@ -494,7 +506,10 @@ async def _stub_prepare_window(entity_ids: list[int]): entity_id=entity_id, sync_start=0.0, source_rows_count=1, - embedding_jobs=[(201, "chunk-1a"), (202, "chunk-1b")], + embedding_jobs=[ + _pending_job(1, 201, "chunk-1a"), + _pending_job(1, 202, "chunk-1b"), + ], chunks_total=5, pending_jobs_total=5, entity_complete=False, @@ -510,7 +525,7 @@ async def _stub_prepare_window(entity_ids: list[int]): entity_id=entity_id, sync_start=0.0, source_rows_count=1, - embedding_jobs=[(301, "chunk-2a")], + embedding_jobs=[_pending_job(2, 301, "chunk-2a")], chunks_total=1, pending_jobs_total=1, entity_complete=True, diff --git a/tests/repository/test_semantic_search_base.py b/tests/repository/test_semantic_search_base.py index e45caced0..0d5f83afe 100644 --- a/tests/repository/test_semantic_search_base.py +++ b/tests/repository/test_semantic_search_base.py @@ -30,12 +30,23 @@ VectorMatch, VectorRecord, ) +from basic_memory.repository.semantic_vector_sync import PendingEmbeddingJob from basic_memory.schemas.search import SearchItemType, SearchRetrievalMode # --- Helpers --- +def _pending_job(entity_id: int, row_id: int, chunk_text: str) -> PendingEmbeddingJob: + return PendingEmbeddingJob( + entity_id=entity_id, + chunk_row_id=row_id, + chunk_key=f"entity:{entity_id}:0", + chunk_text=chunk_text, + source_hash=hashlib.sha256(chunk_text.encode("utf-8")).hexdigest(), + ) + + class _ConcreteRepo(SearchRepositoryBase): """Minimal concrete subclass for testing base class methods.""" @@ -245,7 +256,10 @@ async def fake_scoped_session(_session_maker): 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]]) + await repo._persist_embeddings( + [_pending_job(41, 7, "old chunk text")], + [[1.0, 0.0, 0.0, 0.0]], + ) assert adapter.upserted_records == [] assert session.execute.await_count == 1 @@ -284,7 +298,10 @@ async def fake_scoped_session(_session_maker): 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]]) + await repo._persist_embeddings( + [_pending_job(41, 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 @@ -332,7 +349,10 @@ async def fake_scoped_session(_session_maker): 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]]) + await repo._persist_embeddings( + [_pending_job(41, 7, chunk_text)], + [[1.0, 0.0, 0.0, 0.0]], + ) project_lock_statement = session.execute.await_args_list[0].args[0] manifest_lock_statement = session.execute.await_args_list[1].args[0] @@ -386,7 +406,10 @@ async def fake_scoped_session(_session_maker): 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]]) + await repo._persist_embeddings( + [_pending_job(41, 7, chunk_text)], + [[1.0, 0.0, 0.0, 0.0]], + ) project_lock_statement = session.execute.await_args_list[0].args[0] manifest_lock_statement = session.execute.await_args_list[1].args[0] @@ -942,9 +965,9 @@ async def test_sync_entity_vectors_batch_flushes_at_configured_threshold(monkeyp repo._semantic_embedding_sync_batch_size = 2 prepared_by_entity = { - 1: _PreparedEntityVectorSync(1, 1.0, 1, [(101, "chunk-1")]), - 2: _PreparedEntityVectorSync(2, 2.0, 1, [(102, "chunk-2")]), - 3: _PreparedEntityVectorSync(3, 3.0, 1, [(103, "chunk-3")]), + 1: _PreparedEntityVectorSync(1, 1.0, 1, [_pending_job(1, 101, "chunk-1")]), + 2: _PreparedEntityVectorSync(2, 2.0, 1, [_pending_job(2, 102, "chunk-2")]), + 3: _PreparedEntityVectorSync(3, 3.0, 1, [_pending_job(3, 103, "chunk-3")]), } flush_sizes: list[int] = [] @@ -1021,8 +1044,8 @@ async def test_sync_entity_vectors_batch_progress_tracks_terminal_entities(monke prepared_by_entity = { 1: _PreparedEntityVectorSync(1, 1.0, 1, []), - 2: _PreparedEntityVectorSync(2, 2.0, 1, [(102, "chunk-2")]), - 3: _PreparedEntityVectorSync(3, 3.0, 1, [(103, "chunk-3")]), + 2: _PreparedEntityVectorSync(2, 2.0, 1, [_pending_job(2, 102, "chunk-2")]), + 3: _PreparedEntityVectorSync(3, 3.0, 1, [_pending_job(3, 103, "chunk-3")]), } progress_events: list[tuple[int, int, int]] = [] @@ -1071,7 +1094,10 @@ async def _stub_prepare_window(entity_ids: list[int]): continue prepared.append( _PreparedEntityVectorSync( - entity_id, float(entity_id), 1, [(100 + entity_id, "chunk")] + entity_id, + float(entity_id), + 1, + [_pending_job(entity_id, 100 + entity_id, "chunk")], ) ) return prepared @@ -1129,7 +1155,7 @@ async def _stub_prepare_window(entity_ids: list[int]): entity_id=2, sync_start=0.0, source_rows_count=1, - embedding_jobs=[(102, "chunk-2")], + embedding_jobs=[_pending_job(2, 102, "chunk-2")], prepare_seconds=1.0, ) ) @@ -1176,7 +1202,7 @@ async def _stub_prepare_window(entity_ids: list[int]): entity_id=entity_id, sync_start=0.0, source_rows_count=1, - embedding_jobs=[(100 + entity_id, f"chunk-{entity_id}")], + embedding_jobs=[_pending_job(entity_id, 100 + entity_id, f"chunk-{entity_id}")], prepare_seconds=1.0, ) for entity_id in entity_ids @@ -1292,7 +1318,7 @@ async def _stub_prepare_window(entity_ids: list[int]): entity_id=entity_id, sync_start=0.0, source_rows_count=1, - embedding_jobs=[(100 + entity_id, f"chunk-{entity_id}")], + embedding_jobs=[_pending_job(entity_id, 100 + entity_id, f"chunk-{entity_id}")], prepare_seconds=1.0, ) for entity_id in entity_ids diff --git a/tests/repository/test_semantic_vector_sync.py b/tests/repository/test_semantic_vector_sync.py index 9232dacbf..8b89cea2d 100644 --- a/tests/repository/test_semantic_vector_sync.py +++ b/tests/repository/test_semantic_vector_sync.py @@ -1,5 +1,6 @@ """Focused edge-case coverage for shared semantic vector synchronization.""" +import hashlib from contextlib import asynccontextmanager from datetime import datetime from types import SimpleNamespace @@ -90,10 +91,24 @@ def _distance_to_similarity(self, distance: float) -> float: return 1.0 / (1.0 + max(distance, 0.0)) +def _pending_job( + entity_id: int = 1, + row_id: int = 10, + chunk_text: str = "chunk", +) -> semantic_vector_sync.PendingEmbeddingJob: + return semantic_vector_sync.PendingEmbeddingJob( + entity_id=entity_id, + chunk_row_id=row_id, + chunk_key=f"entity:{entity_id}:0", + chunk_text=chunk_text, + source_hash=hashlib.sha256(chunk_text.encode("utf-8")).hexdigest(), + ) + + def _prepared_entity( entity_id: int = 1, *, - embedding_jobs: list[tuple[int, str]] | None = None, + embedding_jobs: list[semantic_vector_sync.PendingEmbeddingJob] | None = None, entity_complete: bool = True, ) -> semantic_vector_sync.PreparedEntityVectorSync: return semantic_vector_sync.PreparedEntityVectorSync( @@ -182,7 +197,7 @@ async def test_vector_sync_propagates_prepare_and_threshold_flush_errors( flush_repository = _batch_repository( monkeypatch, - [_prepared_entity(embedding_jobs=[(10, "chunk")])], + [_prepared_entity(embedding_jobs=[_pending_job()])], batch_size=1, ) monkeypatch.setattr( @@ -205,7 +220,7 @@ async def test_vector_sync_handles_final_flush_errors_and_orphan_runtime( ) -> None: failed_repository = _batch_repository( monkeypatch, - [_prepared_entity(embedding_jobs=[(10, "chunk")])], + [_prepared_entity(embedding_jobs=[_pending_job()])], ) monkeypatch.setattr( failed_repository, @@ -224,7 +239,7 @@ async def test_vector_sync_handles_final_flush_errors_and_orphan_runtime( strict_repository = _batch_repository( monkeypatch, - [_prepared_entity(embedding_jobs=[(10, "chunk")])], + [_prepared_entity(embedding_jobs=[_pending_job()])], ) monkeypatch.setattr( strict_repository, @@ -241,7 +256,7 @@ async def test_vector_sync_handles_final_flush_errors_and_orphan_runtime( orphan_repository = _batch_repository( monkeypatch, - [_prepared_entity(embedding_jobs=[(10, "chunk")])], + [_prepared_entity(embedding_jobs=[_pending_job()])], batch_size=1, ) orphan_result = await semantic_vector_sync.sync_entity_vectors_internal( @@ -586,11 +601,7 @@ async def test_flush_embedding_jobs_handles_empty_mismatch_and_missing_runtime(m 0.0, ) - job = semantic_vector_sync.PendingEmbeddingJob( - entity_id=1, - chunk_row_id=10, - chunk_text="chunk", - ) + job = _pending_job() with pytest.raises(RuntimeError, match="unexpected number"): await semantic_vector_sync.flush_embedding_jobs(repository, [job], {}, set()) diff --git a/tests/repository/test_sqlite_vector_search_repository.py b/tests/repository/test_sqlite_vector_search_repository.py index 2da7de361..929000332 100644 --- a/tests/repository/test_sqlite_vector_search_repository.py +++ b/tests/repository/test_sqlite_vector_search_repository.py @@ -27,6 +27,7 @@ VectorRecord, ) from basic_memory.repository.semantic_vector_sync import ( + PendingEmbeddingJob, PreparedEntityVectorSync, StagedVectorDeletion, ) @@ -127,6 +128,23 @@ async def search( return [VectorMatch(key=key, similarity=1.0) for key in list(self.records)[:limit]] +def _pending_job( + entity_id: int, + row_id: int, + chunk_text: str, + *, + chunk_key: str | None = None, + source_hash: str | None = None, +) -> PendingEmbeddingJob: + return PendingEmbeddingJob( + entity_id=entity_id, + chunk_row_id=row_id, + chunk_key=chunk_key or f"entity:{entity_id}:0", + chunk_text=chunk_text, + source_hash=source_hash or hashlib.sha256(chunk_text.encode("utf-8")).hexdigest(), + ) + + def _entity_row( *, project_id: int, @@ -668,7 +686,7 @@ async def fail_commit() -> None: ) with pytest.raises(RuntimeError, match="ready commit failed"): await search_repository._persist_embeddings( - [(row_id, "ready commit retry")], + [_pending_job(115, row_id, "ready commit retry")], [[1.0, 0.0, 0.0, 0.0]], ) monkeypatch.setattr( @@ -685,7 +703,7 @@ async def fail_commit() -> None: assert failed_status.scalar_one() == "pending" await search_repository._persist_embeddings( - [(row_id, "ready commit retry")], + [_pending_job(115, row_id, "ready commit retry")], [[1.0, 0.0, 0.0, 0.0]], ) @@ -1135,7 +1153,15 @@ async def _stub_upsert( embedding_model: str, ): await asyncio.sleep(0) - return [(entity_id * 100, scheduled_records[0]["chunk_text"])] + return [ + _pending_job( + entity_id, + entity_id * 100, + scheduled_records[0]["chunk_text"], + chunk_key=scheduled_records[0]["chunk_key"], + source_hash=scheduled_records[0]["source_hash"], + ) + ] @asynccontextmanager async def fake_scoped_session(session_maker): @@ -1204,7 +1230,15 @@ async def _stub_upsert( entity_fingerprint: str, embedding_model: str, ): - return [(entity_id * 100, scheduled_records[0]["chunk_text"])] + return [ + _pending_job( + entity_id, + entity_id * 100, + scheduled_records[0]["chunk_text"], + chunk_key=scheduled_records[0]["chunk_key"], + source_hash=scheduled_records[0]["source_hash"], + ) + ] @asynccontextmanager async def fake_scoped_session(session_maker): diff --git a/tests/repository/test_vector_manifest_generation_ownership.py b/tests/repository/test_vector_manifest_generation_ownership.py new file mode 100644 index 000000000..ea079beed --- /dev/null +++ b/tests/repository/test_vector_manifest_generation_ownership.py @@ -0,0 +1,355 @@ +"""PostgreSQL regressions for semantic vector generation ownership.""" + +from __future__ import annotations + +import asyncio +from collections.abc import Sequence +from datetime import datetime, timezone + +import pytest +from sqlalchemy import text + +from basic_memory import db +from basic_memory.config import BasicMemoryConfig, DatabaseBackend +from basic_memory.repository.postgres_search_repository import PostgresSearchRepository +from basic_memory.repository.search_index_row import SearchIndexRow +from basic_memory.repository.semantic_vector_index import ( + VectorDeletion, + VectorIndexScope, + VectorKey, + VectorMatch, + VectorRecord, +) +from basic_memory.schemas.search import SearchItemType + + +pytestmark = pytest.mark.postgres + + +class CoordinatedEmbeddingProvider: + """Deterministic provider that can pause one vector sync after prepare.""" + + model_name = "generation-ownership-test" + dimensions = 4 + + def __init__( + self, + *, + started: asyncio.Event | None = None, + resume: asyncio.Event | None = None, + ) -> None: + self.started = started + self.resume = resume + + async def embed_query(self, text: str) -> list[float]: + return self._vectorize(text) + + async def embed_documents(self, texts: list[str]) -> list[list[float]]: + if self.started is not None: + assert self.resume is not None + self.started.set() + await self.resume.wait() + return [self._vectorize(text) for text in texts] + + def runtime_log_attrs(self) -> dict[str, object]: + return {} + + @staticmethod + def _vectorize(text: str) -> list[float]: + bucket = len(text) % 4 + return [1.0 if index == bucket else 0.0 for index in range(4)] + + +class InMemoryExternalVectorIndex: + """Generation-aware external adapter used behind Core's manifest contract.""" + + def __init__(self, scope: VectorIndexScope) -> None: + self._scope = scope + self.records: dict[VectorKey, VectorRecord] = {} + + @property + def scope(self) -> VectorIndexScope: + return self._scope + + async def initialize(self) -> None: + return None + + async def upsert(self, records: Sequence[VectorRecord]) -> None: + for record in records: + self.records[record.key] = record + + async def delete(self, records: Sequence[VectorDeletion]) -> None: + for deletion in records: + current = self.records.get(deletion.key) + if current is not None and current.source_hash == deletion.source_hash: + self.records.pop(deletion.key) + + async def delete_entity(self, entity_id: int) -> None: + self.records = { + key: record for key, record in self.records.items() if key.entity_id != entity_id + } + + async def search( + self, + query: Sequence[float], + *, + limit: int, + ) -> list[VectorMatch]: + return [] + + +@pytest.fixture(autouse=True) +def _require_postgres_backend(db_backend: str) -> None: + """These concurrency regressions need PostgreSQL transaction semantics.""" + if db_backend != "postgres": + pytest.skip("Vector generation ownership tests require BASIC_MEMORY_TEST_POSTGRES=1") + + +async def _skip_if_pgvector_unavailable(session_maker) -> None: + async with db.scoped_session(session_maker) as session: + try: + await session.execute(text("CREATE EXTENSION IF NOT EXISTS vector")) + await session.commit() + except Exception: + pytest.skip("pgvector extension is unavailable in this PostgreSQL test environment") + + +def _app_config() -> BasicMemoryConfig: + return BasicMemoryConfig( + env="test", + projects={"test-project": "/tmp/basic-memory-test"}, + default_project="test-project", + database_backend=DatabaseBackend.POSTGRES, + semantic_search_enabled=True, + ) + + +async def _repositories( + session_maker, + *, + project_id: int, + vector_backend: str, + blocking_provider: CoordinatedEmbeddingProvider, +) -> tuple[ + PostgresSearchRepository, + PostgresSearchRepository, + InMemoryExternalVectorIndex | None, +]: + vector_index: InMemoryExternalVectorIndex | None = None + vector_index_name: str | None = None + if vector_backend == "pgvector": + await _skip_if_pgvector_unavailable(session_maker) + else: + vector_index_name = "test-external" + vector_index = InMemoryExternalVectorIndex( + VectorIndexScope( + namespace="generation-ownership", + project_id=project_id, + embedding_identity="test:4", + dimensions=4, + ) + ) + + app_config = _app_config() + owner = PostgresSearchRepository( + session_maker, + project_id, + app_config=app_config, + embedding_provider=blocking_provider, + vector_index_name=vector_index_name, + vector_index=vector_index, + ) + successor = PostgresSearchRepository( + session_maker, + project_id, + app_config=app_config, + embedding_provider=CoordinatedEmbeddingProvider(), + vector_index_name=vector_index_name, + vector_index=vector_index, + ) + return owner, successor, vector_index + + +async def _index_entity( + repository: PostgresSearchRepository, + *, + entity_id: int, + content: str, +) -> None: + now = datetime.now(timezone.utc) + await repository.index_item( + SearchIndexRow( + project_id=repository.project_id, + id=entity_id, + title="Vector Generation Ownership", + content_stems=content, + content_snippet=content, + permalink="specs/vector-generation-ownership", + file_path="specs/vector-generation-ownership.md", + type=SearchItemType.ENTITY.value, + entity_id=entity_id, + metadata={"note_type": "spec"}, + created_at=now, + updated_at=now, + ) + ) + + +async def _manifest_rows(session_maker, *, project_id: int, entity_id: int): + async with db.scoped_session(session_maker) as session: + result = await session.execute( + text( + "SELECT id, chunk_key, source_hash, embedding_status " + "FROM search_vector_chunks " + "WHERE project_id = :project_id AND entity_id = :entity_id " + "ORDER BY chunk_key" + ), + {"project_id": project_id, "entity_id": entity_id}, + ) + return result.mappings().all() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("vector_backend", ["pgvector", "external"]) +async def test_superseded_manifest_generation_defers_old_job_without_failure( + session_maker, + test_project, + vector_backend: str, +) -> None: + """A newer generation owns convergence after deleting job A's prepared rows.""" + embed_started = asyncio.Event() + resume_embedding = asyncio.Event() + owner, successor, external_index = await _repositories( + session_maker, + project_id=test_project.id, + vector_backend=vector_backend, + blocking_provider=CoordinatedEmbeddingProvider( + started=embed_started, + resume=resume_embedding, + ), + ) + entity_id = 1200 + await _index_entity( + owner, + entity_id=entity_id, + content="# Generation A\n- old alpha\n- old beta\n- old gamma", + ) + + owner_task = asyncio.create_task(owner.sync_entity_vectors_batch([entity_id])) + try: + await asyncio.wait_for(embed_started.wait(), timeout=10) + prepared_rows = await _manifest_rows( + session_maker, + project_id=test_project.id, + entity_id=entity_id, + ) + assert len(prepared_rows) >= 4 + assert {row["embedding_status"] for row in prepared_rows} == {"pending"} + + await _index_entity( + successor, + entity_id=entity_id, + content="# Generation B\n- current vector state", + ) + successor_result = await successor.sync_entity_vectors_batch([entity_id]) + assert successor_result.entities_synced == 1 + assert successor_result.entities_failed == 0 + + resume_embedding.set() + owner_result = await asyncio.wait_for(owner_task, timeout=10) + finally: + resume_embedding.set() + if not owner_task.done(): + owner_task.cancel() + await asyncio.gather(owner_task, return_exceptions=True) + + assert owner_result.entities_synced == 0 + assert owner_result.entities_failed == 0 + assert owner_result.entities_deferred == 1 + + final_rows = await _manifest_rows( + session_maker, + project_id=test_project.id, + entity_id=entity_id, + ) + assert final_rows + assert len(final_rows) < len(prepared_rows) + assert {row["embedding_status"] for row in final_rows} == {"ready"} + + manifest_generations = { + VectorKey(entity_id=entity_id, chunk_key=str(row["chunk_key"])): str(row["source_hash"]) + for row in final_rows + } + if external_index is not None: + assert { + key: record.source_hash for key, record in external_index.records.items() + } == manifest_generations + else: + async with db.scoped_session(session_maker) as session: + parity = await session.execute( + text( + "SELECT COUNT(*) AS manifest_count, COUNT(e.chunk_id) AS embedding_count, " + "COUNT(*) FILTER (WHERE e.source_hash != c.source_hash) AS hash_mismatches " + "FROM search_vector_chunks c " + "LEFT JOIN search_vector_embeddings e ON e.chunk_id = c.id " + "WHERE c.project_id = :project_id AND c.entity_id = :entity_id" + ), + {"project_id": test_project.id, "entity_id": entity_id}, + ) + counts = parity.mappings().one() + assert int(counts["manifest_count"]) == int(counts["embedding_count"]) + assert int(counts["hash_mismatches"]) == 0 + + retry_result = await successor.sync_entity_vectors_batch([entity_id]) + assert retry_result.entities_synced == 1 + assert retry_result.entities_failed == 0 + assert retry_result.entities_skipped == 1 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("vector_backend", ["pgvector", "external"]) +async def test_missing_current_manifest_generation_remains_a_correctness_error( + session_maker, + test_project, + vector_backend: str, +) -> None: + """A deleted row is not stale work while its prepared generation is still current.""" + embed_started = asyncio.Event() + resume_embedding = asyncio.Event() + owner, _successor, _external_index = await _repositories( + session_maker, + project_id=test_project.id, + vector_backend=vector_backend, + blocking_provider=CoordinatedEmbeddingProvider( + started=embed_started, + resume=resume_embedding, + ), + ) + entity_id = 1201 + await _index_entity( + owner, + entity_id=entity_id, + content="# Current Generation\n- this source is still authoritative", + ) + + owner_task = asyncio.create_task(owner.sync_entity_vectors(entity_id)) + try: + await asyncio.wait_for(embed_started.wait(), timeout=10) + async with db.scoped_session(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": test_project.id, "entity_id": entity_id}, + ) + await session.commit() + + resume_embedding.set() + with pytest.raises(RuntimeError, match="Vector manifest rows disappeared before write"): + await asyncio.wait_for(owner_task, timeout=10) + finally: + resume_embedding.set() + if not owner_task.done(): + owner_task.cancel() + await asyncio.gather(owner_task, return_exceptions=True)