diff --git a/src/config.py b/src/config.py index 3fa3f7e1..b2b24654 100644 --- a/src/config.py +++ b/src/config.py @@ -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 diff --git a/src/deriver/queue_manager.py b/src/deriver/queue_manager.py index e67ffe9a..6f681618 100644 --- a/src/deriver/queue_manager.py +++ b/src/deriver/queue_manager.py @@ -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: diff --git a/src/deriver/vector_reconciliation.py b/src/deriver/vector_reconciliation.py index da076bb8..9cd5ca17 100644 --- a/src/deriver/vector_reconciliation.py +++ b/src/deriver/vector_reconciliation.py @@ -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 diff --git a/tests/test_vector_reconciliation_config.py b/tests/test_vector_reconciliation_config.py deleted file mode 100644 index fd2dd8b2..00000000 --- a/tests/test_vector_reconciliation_config.py +++ /dev/null @@ -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