diff --git a/src/basic_memory/repository/semantic_vector_sync.py b/src/basic_memory/repository/semantic_vector_sync.py index b6d656d41..851d9a985 100644 --- a/src/basic_memory/repository/semantic_vector_sync.py +++ b/src/basic_memory/repository/semantic_vector_sync.py @@ -2,7 +2,6 @@ from __future__ import annotations -import asyncio import math import time from collections.abc import Callable @@ -79,6 +78,38 @@ class EntityVectorShardPlan: entity_complete: bool +@dataclass(frozen=True, slots=True) +class DeleteEntityVectorPreparePlan: + """Delete stale vector state for an entity with no indexable source rows.""" + + entity_id: int + sync_start: float + prepare_start: float + source_rows_count: int + + +@dataclass(frozen=True, slots=True) +class UpsertEntityVectorPreparePlan: + """Write-side mutations planned from one entity's prefetched vector state.""" + + entity_id: int + sync_start: float + prepare_start: float + source_rows_count: int + existing_by_key: dict[str, VectorChunkState] + stale_ids: list[int] + metadata_update_ids: list[int] + scheduled_records: list[VectorChunkRecord] + entity_fingerprint: str + embedding_model: str + chunks_total: int + chunks_skipped: int + shard_plan: EntityVectorShardPlan + + +type EntityVectorPreparePlan = DeleteEntityVectorPreparePlan | UpsertEntityVectorPreparePlan + + @dataclass class PendingEmbeddingJob: """Pending embedding write entry with entity ownership metadata.""" @@ -630,7 +661,7 @@ async def prepare_entity_vector_jobs_window( repository: SearchRepositoryBase, entity_ids: list[int], ) -> list[PreparedEntityVectorSync | BaseException]: - """Prepare one window of entity vector jobs with shared read-side batching.""" + """Prepare one entity window with batched reads and one write transaction.""" if not entity_ids: return [] @@ -650,24 +681,50 @@ async def prepare_entity_vector_jobs_window( # Outcome: every entity in the window gets the same failure object. return [exc for _ in entity_ids] - # Trigger: prepare now does one shared read pass per window instead of - # paying the same select/join round-trips per entity. - # Why: both SQLite and Postgres were still burning wall clock in read-side - # fingerprint/orphan checks even when every entity ended up skipped. - # Outcome: batch the reads once, close that shared read session, and - # then fan back out over entities while preserving input order. - prepared_window = await asyncio.gather( - *( - repository._prepare_entity_vector_jobs_prefetched( + prepared_by_index: dict[int, PreparedEntityVectorSync | BaseException] = {} + mutation_plans: list[tuple[int, EntityVectorPreparePlan]] = [] + for index, entity_id in enumerate(entity_ids): + try: + planned = plan_entity_vector_jobs_prefetched( + repository, entity_id=entity_id, source_rows=source_rows_by_entity.get(entity_id, []), existing_rows=existing_rows_by_entity.get(entity_id, []), ) - for entity_id in entity_ids - ), - return_exceptions=True, - ) - return list(prepared_window) + except Exception as exc: + prepared_by_index[index] = exc + continue + + if isinstance(planned, PreparedEntityVectorSync): + prepared_by_index[index] = planned + else: + mutation_plans.append((index, planned)) + + if mutation_plans: + try: + # Trigger: every entity in this prepare window has already been + # diffed against one shared read snapshot. + # Why: opening and committing one transaction per entity adds a + # Neon round-trip and repeated writer setup for the same window. + # Outcome: apply the planned mutations in input order and commit + # the whole window once; skip-only entities never enter the write. + 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: + prepared_by_index[index] = await apply_entity_vector_prepare_plan( + repository, + session, + plan, + ) + await session.commit() + except Exception as exc: + # The mutation plans share one transaction, so a failed write + # invalidates every entity whose result depended on that commit. + for index, _plan in mutation_plans: + prepared_by_index[index] = exc + + return [prepared_by_index[index] for index in range(len(entity_ids))] async def prepare_entity_vector_jobs( @@ -689,34 +746,52 @@ async def prepare_entity_vector_jobs_prefetched( source_rows: list[Any], existing_rows: list[VectorChunkState], ) -> PreparedEntityVectorSync: - """Prepare one entity using prefetched window rows.""" + """Prepare one entity using prefetched rows and its own write transaction.""" + planned = plan_entity_vector_jobs_prefetched( + repository, + entity_id=entity_id, + source_rows=source_rows, + existing_rows=existing_rows, + ) + if isinstance(planned, PreparedEntityVectorSync): + return planned + + 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 session.commit() + return prepared + + +def plan_entity_vector_jobs_prefetched( + repository: SearchRepositoryBase, + *, + entity_id: int, + source_rows: list[Any], + existing_rows: list[VectorChunkState], +) -> PreparedEntityVectorSync | EntityVectorPreparePlan: + """Plan one entity from prefetched rows without opening a write transaction.""" sync_start = time.perf_counter() prepare_start = sync_start source_rows_count = len(source_rows) - async def delete_entity_chunks_and_finish() -> PreparedEntityVectorSync: - """Delete derived rows and return the empty prepare result.""" - async with repository._prepare_entity_write_scope(): - async with db.scoped_session(repository.session_maker) as session: - await repository._prepare_vector_session(session) - await repository._delete_entity_chunks(session, entity_id) - await session.commit() - prepare_seconds = time.perf_counter() - prepare_start - return PreparedEntityVectorSync( + def delete_entity_chunks() -> DeleteEntityVectorPreparePlan: + """Plan cleanup for an entity without indexable semantic source rows.""" + return DeleteEntityVectorPreparePlan( entity_id=entity_id, sync_start=sync_start, + prepare_start=prepare_start, source_rows_count=source_rows_count, - embedding_jobs=[], - prepare_seconds=prepare_seconds, ) if not source_rows: - return await delete_entity_chunks_and_finish() + return delete_entity_chunks() chunk_records = repository._build_chunk_records(source_rows) built_chunk_records_count = len(chunk_records) if not chunk_records: - return await delete_entity_chunks_and_finish() + return delete_entity_chunks() current_entity_fingerprint = repository._build_entity_fingerprint(chunk_records) current_embedding_model = repository._embedding_model_key() @@ -756,7 +831,6 @@ async def delete_entity_chunks_and_finish() -> PreparedEntityVectorSync: prepare_seconds=prepare_seconds, ) - timestamp_expr = repository._timestamp_now_expr() metadata_update_ids: list[int] = [] pending_records: list[VectorChunkRecord] = [] skipped_chunks_count = 0 @@ -791,58 +865,83 @@ async def delete_entity_chunks_and_finish() -> PreparedEntityVectorSync: if record["chunk_key"] in shard_plan.scheduled_chunk_keys ] - embedding_jobs: list[tuple[int, str]] = [] - if stale_ids or metadata_update_ids or scheduled_records: - # Trigger: prepare needs to mutate chunk rows for this entity. - # Why: Postgres can keep these write-side steps concurrent, while - # SQLite should funnel them through one writer even after shared reads. - # Outcome: backends share the batched read path without forcing - # SQLite into unnecessary concurrent write transactions. - async with repository._prepare_entity_write_scope(): - async with db.scoped_session(repository.session_maker) as session: - await repository._prepare_vector_session(session) - if stale_ids: - await repository._delete_stale_chunks(session, stale_ids, entity_id) - for row_id in metadata_update_ids: - await session.execute( - text( - "UPDATE search_vector_chunks " - "SET entity_fingerprint = :entity_fingerprint, " - "embedding_model = :embedding_model, " - f"updated_at = {timestamp_expr} " - "WHERE id = :id" - ), - { - "id": row_id, - "entity_fingerprint": current_entity_fingerprint, - "embedding_model": current_embedding_model, - }, - ) - if scheduled_records: - embedding_jobs = await repository._upsert_scheduled_chunk_records( - session, - entity_id=entity_id, - scheduled_records=scheduled_records, - existing_by_key=existing_by_key, - entity_fingerprint=current_entity_fingerprint, - embedding_model=current_embedding_model, - ) - await session.commit() - - prepare_seconds = time.perf_counter() - prepare_start - return PreparedEntityVectorSync( + return UpsertEntityVectorPreparePlan( entity_id=entity_id, sync_start=sync_start, + prepare_start=prepare_start, source_rows_count=source_rows_count, - embedding_jobs=embedding_jobs, + existing_by_key=existing_by_key, + stale_ids=stale_ids, + metadata_update_ids=metadata_update_ids, + scheduled_records=scheduled_records, + entity_fingerprint=current_entity_fingerprint, + embedding_model=current_embedding_model, chunks_total=built_chunk_records_count, chunks_skipped=skipped_chunks_count, - entity_complete=shard_plan.entity_complete, - oversized_entity=shard_plan.oversized_entity, - pending_jobs_total=shard_plan.pending_jobs_total, - shard_index=shard_plan.shard_index, - shard_count=shard_plan.shard_count, - remaining_jobs_after_shard=shard_plan.remaining_jobs_after_shard, + shard_plan=shard_plan, + ) + + +async def apply_entity_vector_prepare_plan( + repository: SearchRepositoryBase, + session: AsyncSession, + plan: EntityVectorPreparePlan, +) -> PreparedEntityVectorSync: + """Apply one planned entity mutation inside the caller-owned transaction.""" + if isinstance(plan, DeleteEntityVectorPreparePlan): + await repository._delete_entity_chunks(session, plan.entity_id) + return PreparedEntityVectorSync( + entity_id=plan.entity_id, + sync_start=plan.sync_start, + source_rows_count=plan.source_rows_count, + embedding_jobs=[], + prepare_seconds=time.perf_counter() - plan.prepare_start, + ) + + timestamp_expr = repository._timestamp_now_expr() + if plan.stale_ids: + await repository._delete_stale_chunks(session, plan.stale_ids, plan.entity_id) + for row_id in plan.metadata_update_ids: + await session.execute( + text( + "UPDATE search_vector_chunks " + "SET entity_fingerprint = :entity_fingerprint, " + "embedding_model = :embedding_model, " + f"updated_at = {timestamp_expr} " + "WHERE id = :id" + ), + { + "id": row_id, + "entity_fingerprint": plan.entity_fingerprint, + "embedding_model": plan.embedding_model, + }, + ) + + embedding_jobs: list[tuple[int, str]] = [] + if plan.scheduled_records: + embedding_jobs = await repository._upsert_scheduled_chunk_records( + session, + entity_id=plan.entity_id, + scheduled_records=plan.scheduled_records, + existing_by_key=plan.existing_by_key, + entity_fingerprint=plan.entity_fingerprint, + embedding_model=plan.embedding_model, + ) + + prepare_seconds = time.perf_counter() - plan.prepare_start + return PreparedEntityVectorSync( + entity_id=plan.entity_id, + sync_start=plan.sync_start, + source_rows_count=plan.source_rows_count, + embedding_jobs=embedding_jobs, + chunks_total=plan.chunks_total, + chunks_skipped=plan.chunks_skipped, + entity_complete=plan.shard_plan.entity_complete, + oversized_entity=plan.shard_plan.oversized_entity, + pending_jobs_total=plan.shard_plan.pending_jobs_total, + shard_index=plan.shard_plan.shard_index, + shard_count=plan.shard_plan.shard_count, + remaining_jobs_after_shard=plan.shard_plan.remaining_jobs_after_shard, prepare_seconds=prepare_seconds, queue_start=time.perf_counter(), ) diff --git a/tests/repository/test_postgres_search_repository_unit.py b/tests/repository/test_postgres_search_repository_unit.py index b124a0744..a119129e3 100644 --- a/tests/repository/test_postgres_search_repository_unit.py +++ b/tests/repository/test_postgres_search_repository_unit.py @@ -5,7 +5,6 @@ are difficult to reach in integration tests. """ -import asyncio from contextlib import asynccontextmanager from unittest.mock import AsyncMock, MagicMock, patch @@ -288,7 +287,7 @@ class TestBatchPrepareWindow: """Cover the shared batched prepare window used by Postgres.""" @pytest.mark.asyncio - async def test_sync_entity_vectors_batch_uses_shared_prepare_window(self, monkeypatch): + async def test_sync_entity_vectors_batch_uses_shared_prepare_transactions(self, monkeypatch): repo = _make_repo( semantic_enabled=True, embedding_provider=StubEmbeddingProvider(), @@ -298,45 +297,50 @@ async def test_sync_entity_vectors_batch_uses_shared_prepare_window(self, monkey repo._vector_tables_initialized = True fetched_windows: list[list[int]] = [] - prepared_windows: list[list[int]] = [] - active_prepares = 0 - max_active_prepares = 0 + upserted_entity_ids: list[int] = [] + sessions: list[AsyncMock] = [] + write_scope_entries = 0 async def _stub_fetch_source_rows(session, entity_ids: list[int]): fetched_windows.append(list(entity_ids)) - return {entity_id: [object()] for entity_id in entity_ids} + return {entity_id: [entity_id] for entity_id in entity_ids} async def _stub_fetch_existing_rows(session, entity_ids: list[int]): return {entity_id: [] for entity_id in entity_ids} - async def _stub_prepare_prefetched( + def _stub_build_chunk_records(source_rows): + entity_id = source_rows[0] + return [ + { + "chunk_key": f"entity:{entity_id}:0", + "chunk_text": f"chunk {entity_id}", + "source_hash": f"hash-{entity_id}", + } + ] + + async def _stub_upsert( + session, *, entity_id: int, - source_rows, - existing_rows, - ) -> _PreparedEntityVectorSync: - nonlocal active_prepares, max_active_prepares - assert len(source_rows) == 1 - assert existing_rows == [] - active_prepares += 1 - max_active_prepares = max(max_active_prepares, active_prepares) - await asyncio.sleep(0) - active_prepares -= 1 - prepared_windows.append([entity_id]) - return _PreparedEntityVectorSync( - entity_id=entity_id, - sync_start=float(entity_id), - source_rows_count=1, - embedding_jobs=[], - entity_skipped=True, - chunks_total=1, - chunks_skipped=1, - prepare_seconds=0.1, - ) + scheduled_records, + existing_by_key, + entity_fingerprint: str, + embedding_model: str, + ): + upserted_entity_ids.append(entity_id) + return [] + + @asynccontextmanager + async def _track_write_scope(): + nonlocal write_scope_entries + write_scope_entries += 1 + yield @asynccontextmanager async def fake_scoped_session(session_maker): - yield AsyncMock() + session = AsyncMock() + sessions.append(session) + yield session monkeypatch.setattr(repo, "_ensure_vector_tables", AsyncMock()) monkeypatch.setattr( @@ -345,9 +349,10 @@ async def fake_scoped_session(session_maker): ) monkeypatch.setattr(repo, "_fetch_prepare_window_source_rows", _stub_fetch_source_rows) monkeypatch.setattr(repo, "_fetch_prepare_window_existing_rows", _stub_fetch_existing_rows) - monkeypatch.setattr( - repo, "_prepare_entity_vector_jobs_prefetched", _stub_prepare_prefetched - ) + monkeypatch.setattr(repo, "_prepare_vector_session", AsyncMock()) + monkeypatch.setattr(repo, "_build_chunk_records", _stub_build_chunk_records) + monkeypatch.setattr(repo, "_prepare_entity_write_scope", _track_write_scope) + monkeypatch.setattr(repo, "_upsert_scheduled_chunk_records", _stub_upsert) result = await repo.sync_entity_vectors_batch([1, 2, 3, 4]) @@ -355,8 +360,9 @@ async def fake_scoped_session(session_maker): assert result.entities_synced == 4 assert result.entities_failed == 0 assert fetched_windows == [[1, 2], [3, 4]] - assert prepared_windows == [[1], [2], [3], [4]] - assert max_active_prepares == 2 + assert upserted_entity_ids == [1, 2, 3, 4] + assert write_scope_entries == 2 + assert [session.commit.await_count for session in sessions] == [0, 1, 0, 1] @pytest.mark.asyncio diff --git a/tests/repository/test_semantic_vector_sync.py b/tests/repository/test_semantic_vector_sync.py index 978cc5da4..8be7da334 100644 --- a/tests/repository/test_semantic_vector_sync.py +++ b/tests/repository/test_semantic_vector_sync.py @@ -337,6 +337,61 @@ async def scoped_session(_session_maker): assert all(isinstance(result, RuntimeError) for result in prepared) +@pytest.mark.asyncio +async def test_prepare_window_reports_shared_transaction_failure_for_mutation_plans( + monkeypatch: pytest.MonkeyPatch, +) -> None: + repository = _TestRepository() + skip_result = _prepared_entity(entity_id=1) + sessions = iter([AsyncMock(), AsyncMock()]) + + @asynccontextmanager + async def scoped_session(_session_maker): + yield next(sessions) + + @asynccontextmanager + async def write_scope(): + yield + + def _stub_plan(repository, *, entity_id, source_rows, existing_rows): + if entity_id == 1: + return skip_result + if entity_id == 4: + raise ValueError("planning failed") + return semantic_vector_sync.DeleteEntityVectorPreparePlan( + entity_id=entity_id, + sync_start=0.0, + prepare_start=0.0, + source_rows_count=0, + ) + + 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={})) + monkeypatch.setattr( + repository, "_fetch_prepare_window_existing_rows", AsyncMock(return_value={}) + ) + monkeypatch.setattr(repository, "_prepare_entity_write_scope", write_scope) + monkeypatch.setattr(semantic_vector_sync, "plan_entity_vector_jobs_prefetched", _stub_plan) + monkeypatch.setattr( + repository, + "_delete_entity_chunks", + AsyncMock(side_effect=[None, RuntimeError("write failed")]), + ) + + prepared = await semantic_vector_sync.prepare_entity_vector_jobs_window( + repository, + [1, 2, 3, 4], + ) + + assert prepared[0] is skip_result + assert isinstance(prepared[1], RuntimeError) + assert prepared[1] is prepared[2] + assert str(prepared[1]) == "write failed" + assert isinstance(prepared[3], ValueError) + assert str(prepared[3]) == "planning failed" + + @pytest.mark.asyncio async def test_prepare_single_entity_propagates_window_failure( monkeypatch: pytest.MonkeyPatch, @@ -435,6 +490,41 @@ async def write_scope(): delete_stale_chunks.assert_awaited_once_with(session, [7], 1) +@pytest.mark.asyncio +async def test_prefetched_prepare_returns_unchanged_entity_without_write(monkeypatch) -> None: + repository = _TestRepository() + record = { + "chunk_key": "existing", + "chunk_text": "text", + "source_hash": "source-hash", + } + existing_row = semantic_vector_sync.VectorChunkState( + id=7, + chunk_key="existing", + source_hash="source-hash", + entity_fingerprint="fingerprint", + embedding_model="model", + has_embedding=True, + ) + monkeypatch.setattr(repository, "_build_chunk_records", Mock(return_value=[record])) + monkeypatch.setattr( + repository, + "_build_entity_fingerprint", + Mock(return_value="fingerprint"), + ) + monkeypatch.setattr(repository, "_embedding_model_key", Mock(return_value="model")) + + prepared = await semantic_vector_sync.prepare_entity_vector_jobs_prefetched( + repository, + entity_id=1, + source_rows=[object()], + existing_rows=[existing_row], + ) + + assert prepared.entity_skipped is True + assert prepared.embedding_jobs == [] + + @pytest.mark.asyncio async def test_flush_embedding_jobs_handles_empty_mismatch_and_missing_runtime(monkeypatch) -> None: repository = _TestRepository() diff --git a/tests/repository/test_sqlite_vector_search_repository.py b/tests/repository/test_sqlite_vector_search_repository.py index 751b78462..f6efa5376 100644 --- a/tests/repository/test_sqlite_vector_search_repository.py +++ b/tests/repository/test_sqlite_vector_search_repository.py @@ -409,6 +409,8 @@ async def test_sqlite_prepare_window_uses_shared_reads_and_serialized_write_scop fetched_windows: list[list[int]] = [] active_write_scopes = 0 max_active_write_scopes = 0 + write_scope_entries = 0 + sessions: list[AsyncMock] = [] async def _stub_fetch_source_rows(session, entity_ids: list[int]): fetched_windows.append(list(entity_ids)) @@ -428,8 +430,9 @@ def _stub_build_chunk_records(source_rows): @asynccontextmanager async def _track_write_scope(): - nonlocal active_write_scopes, max_active_write_scopes + nonlocal active_write_scopes, max_active_write_scopes, write_scope_entries async with repo._sqlite_prepare_write_lock: + write_scope_entries += 1 active_write_scopes += 1 max_active_write_scopes = max(max_active_write_scopes, active_write_scopes) try: @@ -451,7 +454,9 @@ async def _stub_upsert( @asynccontextmanager async def fake_scoped_session(session_maker): - yield AsyncMock() + session = AsyncMock() + sessions.append(session) + yield session monkeypatch.setattr( "basic_memory.repository.search_repository_base.db.scoped_session", @@ -470,6 +475,8 @@ async def fake_scoped_session(session_maker): assert fetched_windows == [[1, 2]] assert [result.entity_id for result in prepared_results] == [1, 2] assert max_active_write_scopes == 1 + assert write_scope_entries == 1 + assert [session.commit.await_count for session in sessions] == [0, 1] @pytest.mark.asyncio