Skip to content

fix(core): defer superseded vector generations #1203

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
phernandez merged 1 commit into from
Aug 6, 2026
Merged

fix(core): defer superseded vector generations #1203

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
15 changes: 12 additions & 3 deletions src/basic_memory/repository/postgres_search_repository.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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 []
Expand Down Expand Up @@ -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
]

Expand Down
96 changes: 72 additions & 24 deletions src/basic_memory/repository/search_repository_base.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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.
Expand All @@ -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:
Expand Down Expand Up @@ -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
Expand All @@ -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.
Expand All @@ -451,6 +498,7 @@ async def _persist_embeddings(
generation_predicates=generation_predicates,
)
await session.commit()
return persistence

async def _mark_embedding_jobs_ready(
self,
Expand Down Expand Up @@ -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,
Expand Down
66 changes: 47 additions & 19 deletions src/basic_memory/repository/semantic_vector_sync.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand All @@ -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)
Expand Down Expand Up @@ -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]
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -1115,15 +1119,15 @@ 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
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]] = []
embedding_jobs: list[PendingEmbeddingJob] = []
for record in scheduled_records:
current = existing_by_key.get(record["chunk_key"])
if current:
Expand Down Expand Up @@ -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(
Expand All @@ -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


Expand All @@ -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)
Expand All @@ -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
Expand All @@ -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)
Expand Down
Loading
Loading