feat: sync docs without embeddings

This commit is contained in:
Rajat Ahuja 2026-01-08 14:59:03 -05:00
parent 448905cf4b
commit a8075b9133
4 changed files with 116 additions and 142 deletions

View File

@ -404,20 +404,6 @@ class VectorStoreSettings(HonchoSettings):
)
return self
@property
def should_run_reconciliation(self) -> bool:
"""
Determine if vector reconciliation should run.
Reconciliation syncs embeddings from postgres (used by pgvector) to
external vector stores. It only runs when:
1. A secondary store is configured AND
2. pgvector is involved as either primary or secondary
"""
return self.SECONDARY_TYPE is not None and (
self.PRIMARY_TYPE == "pgvector" or self.SECONDARY_TYPE == "pgvector"
)
class AppSettings(HonchoSettings):
# No env_prefix for app-level settings

View File

@ -51,7 +51,7 @@ class WorkerOwnership(NamedTuple):
QUEUE_CLEANUP_INTERVAL_SECONDS = 43200 # 12 hours
RECONCILIATION_INTERVAL_SECONDS = 300 # 5 minutes
RECONCILIATION_INTERVAL_SECONDS = 900 # 15 minutes
class QueueManager:
@ -353,15 +353,11 @@ class QueueManager:
- Queue cleanup: every 12 hours (remove old processed/errored queue items)
- Reconciliation: every 5 minutes (sync vectors + clean up soft deletes)
Only runs when pgvector is involved in a dual-store configuration
Always runs to keep vector stores consistent with the DB
"""
# Track when each task should next run
next_queue_cleanup = datetime.now(timezone.utc)
next_vector_reconciliation = (
datetime.now(timezone.utc)
if settings.VECTOR_STORE.should_run_reconciliation
else None
)
next_vector_reconciliation = datetime.now(timezone.utc)
try:
while not self.shutdown_event.is_set():
@ -379,11 +375,8 @@ class QueueManager:
seconds=QUEUE_CLEANUP_INTERVAL_SECONDS
)
# Run vector store reconciliation if enabled and due
if (
next_vector_reconciliation is not None
and now >= next_vector_reconciliation
):
# Run vector store reconciliation if due
if now >= next_vector_reconciliation:
try:
logger.info("Running vector reconciliation cycle")
await self._run_reconciliation()
@ -398,8 +391,7 @@ class QueueManager:
# Sleep until next task is due or shutdown
# Filter out None values when computing next task time
task_times = [next_queue_cleanup]
if next_vector_reconciliation is not None:
task_times.append(next_vector_reconciliation)
task_times.append(next_vector_reconciliation)
next_task_time = min(task_times)
sleep_seconds = max(
0, (next_task_time - datetime.now(timezone.utc)).total_seconds()
@ -867,20 +859,6 @@ class QueueManager:
async def main():
logger.debug("Starting queue manager")
# Log reconciliation status
if settings.VECTOR_STORE.should_run_reconciliation:
logger.info(
"Vector reconciliation: ENABLED (primary=%s, secondary=%s)",
settings.VECTOR_STORE.PRIMARY_TYPE,
settings.VECTOR_STORE.SECONDARY_TYPE,
)
else:
logger.info(
"Vector reconciliation: DISABLED (primary=%s, secondary=%s)",
settings.VECTOR_STORE.PRIMARY_TYPE,
settings.VECTOR_STORE.SECONDARY_TYPE or "None",
)
try:
await init_cache()
except Exception as e:

View File

@ -8,13 +8,16 @@ embeddings to the vector store on a rolling basis, healing any missed writes.
import logging
import time
from dataclasses import dataclass
from typing import cast
from sqlalchemy import and_, select, update
from sqlalchemy.ext.asyncio import AsyncSession
from sqlalchemy.sql.functions import func
from src import models
from src.config import settings
from src.dependencies import tracked_db
from src.embedding_client import embedding_client
from src.vector_store import VectorRecord, VectorStore, get_vector_store
logger = logging.getLogger(__name__)
@ -57,8 +60,8 @@ async def _get_documents_needing_sync(
Finds documents where:
- not soft-deleted (deleted_at is NULL)
- has an embedding stored in the database
- sync_state is "pending" (never synced or retry needed)
- embedding may be NULL (will be re-embedded during reconciliation)
- Note: "synced" = done forever, "failed" = permanent failure (manual intervention)
Uses FOR UPDATE SKIP LOCKED to prevent concurrent processing.
@ -68,7 +71,6 @@ async def _get_documents_needing_sync(
.where(
and_(
models.Document.deleted_at.is_(None),
models.Document.embedding.isnot(None), # Must have embedding to sync
models.Document.sync_state == "pending", # Only pending items
)
)
@ -114,6 +116,27 @@ async def _get_message_embeddings_needing_sync(
return list(result.scalars().all())
async def _bump_document_sync_attempts(
db: AsyncSession,
documents: list[models.Document],
) -> None:
if not documents:
return
for doc in documents:
new_attempts = doc.sync_attempts + 1
new_state = "failed" if new_attempts >= MAX_SYNC_ATTEMPTS else "pending"
await db.execute(
update(models.Document)
.where(models.Document.id == doc.id)
.values(
sync_state=new_state,
sync_attempts=new_attempts,
last_sync_at=func.now(),
)
)
async def _sync_documents(
db: AsyncSession,
documents: list[models.Document],
@ -130,6 +153,56 @@ async def _sync_documents(
synced_count = 0
failed_count = 0
pgvector_in_use = (
settings.VECTOR_STORE.PRIMARY_TYPE == "pgvector"
or settings.VECTOR_STORE.SECONDARY_TYPE == "pgvector"
)
missing_docs: list[models.Document] = []
for doc in documents:
if cast(list[float] | None, doc.embedding) is None:
missing_docs.append(doc)
reembedded_by_id: dict[str, list[float]] = {}
if missing_docs:
try:
contents = [doc.content for doc in missing_docs]
embeddings = await embedding_client.simple_batch_embed(contents)
if len(embeddings) != len(missing_docs):
logger.warning(
"Re-embedded %s/%s documents; remaining will be retried",
len(embeddings),
len(missing_docs),
)
for doc, embedding in zip(missing_docs, embeddings, strict=False):
reembedded_by_id[doc.id] = embedding
if pgvector_in_use and reembedded_by_id:
for doc_id, embedding in reembedded_by_id.items():
await db.execute(
update(models.Document)
.where(models.Document.id == doc_id)
.values(embedding=embedding)
)
except Exception as e:
logger.warning(
"Failed to re-embed %s documents for reconciliation: %s",
len(missing_docs),
e,
)
missing_after_embed: list[models.Document] = []
for doc in documents:
if (
cast(list[float] | None, doc.embedding) is None
and doc.id not in reembedded_by_id
):
missing_after_embed.append(doc)
if missing_after_embed:
await _bump_document_sync_attempts(db, missing_after_embed)
failed_count += len(missing_after_embed)
# Group documents by namespace (workspace/observer/observed)
by_namespace: dict[str, list[models.Document]] = {}
for doc in documents:
@ -140,25 +213,39 @@ async def _sync_documents(
# Sync each namespace batch
for namespace, docs in by_namespace.items():
doc_ids = [doc.id for doc in docs]
docs_with_vectors: list[models.Document] = []
try:
# Build vector records
vector_records = [
VectorRecord(
id=doc.id,
embedding=[float(x) for x in doc.embedding],
metadata={
"workspace_name": doc.workspace_name,
"observer": doc.observer,
"observed": doc.observed,
"session_name": doc.session_name,
"level": doc.level,
},
vector_records: list[VectorRecord] = []
for doc in docs:
doc_embedding = cast(list[float] | None, doc.embedding)
embedding = (
doc_embedding
if doc_embedding is not None
else reembedded_by_id.get(doc.id)
)
for doc in docs
if doc.embedding is not None
]
if embedding is None:
continue
vector_records.append(
VectorRecord(
id=doc.id,
embedding=[float(x) for x in embedding],
metadata={
"workspace_name": doc.workspace_name,
"observer": doc.observer,
"observed": doc.observed,
"session_name": doc.session_name,
"level": doc.level,
},
)
)
docs_with_vectors.append(doc)
if not vector_records:
continue
doc_ids = [doc.id for doc in docs_with_vectors]
result = None
if vector_records:
@ -171,22 +258,8 @@ async def _sync_documents(
result.secondary_error,
)
# Increment attempts and mark as failed if we've hit max attempts
for doc in docs:
new_attempts = doc.sync_attempts + 1
new_state = (
"failed" if new_attempts >= MAX_SYNC_ATTEMPTS else "pending"
)
await db.execute(
update(models.Document)
.where(models.Document.id == doc.id)
.values(
sync_state=new_state,
sync_attempts=new_attempts,
last_sync_at=func.now(),
)
)
failed_count += len(docs)
await _bump_document_sync_attempts(db, docs_with_vectors)
failed_count += len(docs_with_vectors)
continue
# Mark as synced
@ -199,25 +272,13 @@ async def _sync_documents(
sync_attempts=0,
)
)
synced_count += len(docs)
synced_count += len(docs_with_vectors)
except Exception as e:
logger.warning(f"Failed to sync documents to {namespace}: {e}")
# Increment attempts and mark as failed if we've hit max attempts
for doc in docs:
new_attempts = doc.sync_attempts + 1
new_state = "failed" if new_attempts >= MAX_SYNC_ATTEMPTS else "pending"
await db.execute(
update(models.Document)
.where(models.Document.id == doc.id)
.values(
sync_state=new_state,
sync_attempts=new_attempts,
last_sync_at=func.now(),
)
)
failed_count += len(docs)
await _bump_document_sync_attempts(db, docs_with_vectors)
failed_count += len(docs_with_vectors)
return synced_count, failed_count

View File

@ -1,51 +0,0 @@
"""Tests for vector reconciliation configuration logic.
These are unit tests that test configuration logic without requiring database
or vector store fixtures.
"""
from src.config import VectorStoreSettings
def test_reconciliation_enabled_pgvector_primary():
"""Reconciliation enabled when pgvector is primary with secondary"""
settings = VectorStoreSettings(
PRIMARY_TYPE="pgvector",
SECONDARY_TYPE="turbopuffer",
TURBOPUFFER_API_KEY="test-key",
)
assert settings.should_run_reconciliation is True
def test_reconciliation_enabled_pgvector_secondary():
"""Reconciliation enabled when pgvector is secondary"""
settings = VectorStoreSettings(
PRIMARY_TYPE="turbopuffer",
SECONDARY_TYPE="pgvector",
TURBOPUFFER_API_KEY="test-key",
)
assert settings.should_run_reconciliation is True
def test_reconciliation_disabled_single_store():
"""Reconciliation disabled when no secondary configured"""
settings = VectorStoreSettings(
PRIMARY_TYPE="turbopuffer", TURBOPUFFER_API_KEY="test-key", SECONDARY_TYPE=None
)
assert settings.should_run_reconciliation is False
def test_reconciliation_disabled_no_pgvector():
"""Reconciliation disabled when both stores are non-pgvector"""
settings = VectorStoreSettings(
PRIMARY_TYPE="turbopuffer",
SECONDARY_TYPE="lancedb",
TURBOPUFFER_API_KEY="test-key",
)
assert settings.should_run_reconciliation is False
def test_reconciliation_disabled_pgvector_only():
"""Reconciliation disabled when pgvector is primary but no secondary"""
settings = VectorStoreSettings(PRIMARY_TYPE="pgvector", SECONDARY_TYPE=None)
assert settings.should_run_reconciliation is False