diff --git a/.env.template b/.env.template index 736bf749..4170c05a 100644 --- a/.env.template +++ b/.env.template @@ -253,3 +253,4 @@ VECTOR_STORE_MIGRATED=false # LanceDB-specific settings (local embedded mode) # VECTOR_STORE_LANCEDB_PATH=./lancedb_data +#VECTOR_STORE_RECONCILIATION_INTERVAL_SECONDS diff --git a/config.toml.example b/config.toml.example index b2db5103..12d7717b 100644 --- a/config.toml.example +++ b/config.toml.example @@ -184,3 +184,4 @@ DIMENSIONS = 1536 # TURBOPUFFER_API_KEY = "your-turbopuffer-api-key" # TURBOPUFFER_REGION = "us-east-1" LANCEDB_PATH = "./lancedb_data" +RECONCILIATION_INTERVAL_SECONDS= diff --git a/migrations/versions/119a52b73c60_support_external_embeddings.py b/migrations/versions/119a52b73c60_support_external_embeddings.py index 31b1f09f..e07bc0c0 100644 --- a/migrations/versions/119a52b73c60_support_external_embeddings.py +++ b/migrations/versions/119a52b73c60_support_external_embeddings.py @@ -7,6 +7,8 @@ This migration: hybrid sync/soft delete pattern for vector store consistency. 3. Adds sync_state, last_sync_at, and sync_attempts columns to documents and message_embeddings tables for tracking vector store synchronization status. +4. Adds partial unique index on queue table for reconciler task deduplication, + ensuring only one pending reconciler task exists per work_unit_key. Revision ID: 119a52b73c60 Revises: 7c0d9a4e3b1f @@ -181,11 +183,31 @@ def upgrade() -> None: schema=schema, ) + # Add partial unique index on queue table for reconciler task deduplication + # This ensures only one pending reconciler task exists per work_unit_key + if not index_exists("queue", "uq_queue_work_unit_key", inspector): + op.create_index( + "uq_queue_work_unit_key", + "queue", + ["work_unit_key"], + unique=True, + schema=schema, + postgresql_where=sa.text("task_type = 'reconciler' AND processed = false"), + ) + def downgrade() -> None: """Remove deleted_at columns and revert embedding columns.""" inspector = sa.inspect(op.get_bind()) + # Drop reconciler queue index if it exists + if index_exists("queue", "uq_queue_work_unit_key", inspector): + op.drop_index( + "uq_queue_work_unit_key", + table_name="queue", + schema=schema, + ) + # Drop message_embeddings indexes if they exist if index_exists( "message_embeddings", "ix_message_embeddings_sync_state_last_sync_at", inspector diff --git a/src/config.py b/src/config.py index c4875aa1..5fb2769d 100644 --- a/src/config.py +++ b/src/config.py @@ -563,6 +563,10 @@ class VectorStoreSettings(HonchoSettings): # LanceDB-specific settings (local embedded mode) LANCEDB_PATH: str = "./lancedb_data" + RECONCILIATION_INTERVAL_SECONDS: Annotated[int, Field(default=300, gt=0)] = ( + 300 # 5 minutes + ) + @model_validator(mode="after") def _require_api_key_for_turbopuffer(self) -> "VectorStoreSettings": if self.TYPE == "turbopuffer" and not self.TURBOPUFFER_API_KEY: diff --git a/src/deriver/consumer.py b/src/deriver/consumer.py index 6d66f30e..ae9bb97e 100644 --- a/src/deriver/consumer.py +++ b/src/deriver/consumer.py @@ -10,12 +10,15 @@ from src.deriver.deriver import process_representation_tasks_batch from src.dreamer.dreamer import process_dream from src.exceptions import ResourceNotFoundException from src.models import Message -from src.schemas import ResolvedConfiguration +from src.reconciler.queue_cleanup import cleanup_queue_items +from src.reconciler.sync_vectors import run_vector_reconciliation_cycle +from src.schemas import ReconcilerType, ResolvedConfiguration from src.utils import summarizer from src.utils.logging import log_performance_metrics from src.utils.queue_payload import ( DeletionPayload, DreamPayload, + ReconcilerPayload, SummaryPayload, WebhookPayload, ) @@ -119,6 +122,19 @@ async def process_item(queue_item: models.QueueItem) -> None: raise ValueError(f"Invalid payload structure: {str(e)}") from e await process_deletion(validated, workspace_name) + elif task_type == "reconciler": + with sentry_sdk.start_transaction(name="process_reconciler_task", op="deriver"): + try: + validated = ReconcilerPayload(**queue_payload) + except ValidationError as e: + logger.error( + "Invalid reconciler payload received: %s. Payload: %s", + str(e), + queue_payload, + ) + raise ValueError(f"Invalid payload structure: {str(e)}") from e + await process_reconciler(validated) + else: raise ValueError(f"Invalid task type: {task_type}") @@ -220,3 +236,43 @@ async def process_deletion( else: raise ValueError(f"Unsupported deletion type: {deletion_type}") + + +async def process_reconciler(payload: ReconcilerPayload) -> None: + """ + Process a reconciler task from the queue. + + Currently supports: + - sync_vectors: Syncs pending documents/message embeddings to vector store + and cleans up soft-deleted documents. + - cleanup_queue: Removes old processed queue items. + + Args: + payload: The reconciler payload containing the reconciler type + """ + reconciler_type = payload.reconciler_type + + if reconciler_type == ReconcilerType.SYNC_VECTORS: + logger.debug("Processing sync_vectors task") + metrics = await run_vector_reconciliation_cycle() + + if ( + metrics.total_synced > 0 + or metrics.total_failed > 0 + or metrics.total_cleaned > 0 + ): + logger.info( + "Reconciliation complete: synced %s docs, %s message embeddings; failed %s docs, %s message embeddings; cleaned %s docs", + metrics.documents_synced, + metrics.message_embeddings_synced, + metrics.documents_failed, + metrics.message_embeddings_failed, + metrics.documents_cleaned, + ) + + elif reconciler_type == ReconcilerType.CLEANUP_QUEUE: + logger.debug("Processing cleanup_queue task") + await cleanup_queue_items() + + else: + raise ValueError(f"Unsupported reconciler type: {reconciler_type}") diff --git a/src/deriver/queue_manager.py b/src/deriver/queue_manager.py index 87ae5831..350ca4d8 100644 --- a/src/deriver/queue_manager.py +++ b/src/deriver/queue_manager.py @@ -23,13 +23,17 @@ from src.deriver.consumer import ( process_item, process_representation_batch, ) -from src.deriver.vector_reconciliation import run_vector_reconciliation_cycle from src.dreamer.dream_scheduler import ( DreamScheduler, get_dream_scheduler, set_dream_scheduler, ) from src.models import QueueItem +from src.reconciler import ( + ReconcilerScheduler, + get_reconciler_scheduler, + set_reconciler_scheduler, +) from src.schemas import ResolvedConfiguration from src.sentry import initialize_sentry from src.utils.work_unit import parse_work_unit_key @@ -50,17 +54,12 @@ class WorkerOwnership(NamedTuple): aqs_id: str # The ID of the ActiveQueueSession that the worker is processing -QUEUE_CLEANUP_INTERVAL_SECONDS = 43200 # 12 hours -RECONCILIATION_INTERVAL_SECONDS = 300 # 5 minutes - - class QueueManager: def __init__(self): self.shutdown_event: asyncio.Event = asyncio.Event() self.active_tasks: set[asyncio.Task[None]] = set() self.worker_ownership: dict[str, WorkerOwnership] = {} self.queue_empty_flag: asyncio.Event = asyncio.Event() - self._maintenance_task: asyncio.Task[None] | None = None # Initialize from settings self.workers: int = settings.DERIVER.WORKERS @@ -74,6 +73,14 @@ class QueueManager: else: self.dream_scheduler = existing_scheduler + # Get or create the singleton reconciler scheduler + existing_reconciler = get_reconciler_scheduler() + if existing_reconciler is None: + self.reconciler_scheduler: ReconcilerScheduler = ReconcilerScheduler() + set_reconciler_scheduler(self.reconciler_scheduler) + else: + self.reconciler_scheduler = existing_reconciler + # Initialize Sentry if enabled, using settings if settings.SENTRY.ENABLED: initialize_sentry(integrations=[AsyncioIntegration()]) @@ -116,11 +123,11 @@ class QueueManager: ) logger.debug("Signal handlers registered") - # Start background maintenance loop (handles both queue cleanup and vector cleanup) + # Start the reconciler scheduler try: - self._maintenance_task = asyncio.create_task(self._maintenance_loop()) + await self.reconciler_scheduler.start() except Exception: - logger.exception("Failed to start maintenance loop") + logger.exception("Failed to start reconciler scheduler") # Run the polling loop directly in this task logger.debug("Starting polling loop directly") @@ -137,6 +144,9 @@ class QueueManager: # Cancel all pending dreams await self.dream_scheduler.shutdown() + # Stop the reconciler scheduler + await self.reconciler_scheduler.shutdown() + if self.active_tasks: logger.info( f"Waiting for {len(self.active_tasks)} active tasks to complete..." @@ -168,14 +178,6 @@ class QueueManager: finally: self.worker_ownership.clear() - # Cancel maintenance loop if running - if self._maintenance_task is not None: - from contextlib import suppress - - self._maintenance_task.cancel() - with suppress(asyncio.CancelledError): - await self._maintenance_task - ########################## # Polling and Scheduling # ########################## @@ -323,113 +325,6 @@ class QueueManager: # Queue Worker Logic # ###################### - async def cleanup_queue_items(self) -> None: - """Delete processed queue items. - Successfully processed queue items are deleted immediately, - while errored queue items are deleted after retention window.""" - async with tracked_db("cleanup_queue_items") as db: - now = datetime.now(timezone.utc) - error_cutoff = now - timedelta( - seconds=settings.DERIVER.QUEUE_ERROR_RETENTION_SECONDS - ) - - await db.execute( - delete(models.QueueItem).where( - models.QueueItem.processed - & ( - models.QueueItem.error.is_(None) - | ( - models.QueueItem.error.is_not(None) - & (models.QueueItem.created_at < error_cutoff) - ) - ) - ) - ) - await db.commit() - - async def _maintenance_loop(self) -> None: - """ - Run periodic maintenance tasks. - - - Queue cleanup: every 12 hours (remove old processed/errored queue items) - - Reconciliation: every 5 minutes (sync vectors + clean up soft deletes) - 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) - - try: - while not self.shutdown_event.is_set(): - now = datetime.now(timezone.utc) - - # Run queue cleanup if due - if now >= next_queue_cleanup: - try: - await self.cleanup_queue_items() - except Exception: - logger.exception("Error during queue cleanup") - if settings.SENTRY.ENABLED: - sentry_sdk.capture_exception() - next_queue_cleanup = now + timedelta( - seconds=QUEUE_CLEANUP_INTERVAL_SECONDS - ) - - # Run vector store reconciliation if due - if now >= next_vector_reconciliation: - try: - logger.info("Running vector reconciliation cycle") - await self._run_reconciliation() - except Exception: - logger.exception("Error during vector reconciliation") - if settings.SENTRY.ENABLED: - sentry_sdk.capture_exception() - next_vector_reconciliation = now + timedelta( - seconds=RECONCILIATION_INTERVAL_SECONDS - ) - - # Sleep until next task is due or shutdown - # Filter out None values when computing next task time - task_times = [next_queue_cleanup] - 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() - ) - - try: - await asyncio.wait_for( - self.shutdown_event.wait(), - timeout=sleep_seconds - or 1, # At least 1 second to avoid busy loop - ) - break # Shutdown event set - except asyncio.TimeoutError: - # Timeout means it's time for next task - pass - except asyncio.CancelledError: - logger.debug("Maintenance loop cancelled") - raise - - async def _run_reconciliation(self) -> None: - """Run vector store reconciliation for sync + cleanup.""" - - metrics = await run_vector_reconciliation_cycle() - - if ( - metrics.total_synced > 0 - or metrics.total_failed > 0 - or metrics.total_cleaned > 0 - ): - logger.info( - "Reconciliation: synced %s docs, %s message embeddings; failed %s docs, %s message embeddings; cleaned %s docs", - metrics.documents_synced, - metrics.message_embeddings_synced, - metrics.documents_failed, - metrics.message_embeddings_failed, - metrics.documents_cleaned, - ) - async def _handle_processing_error( self, error: Exception, diff --git a/src/reconciler/__init__.py b/src/reconciler/__init__.py new file mode 100644 index 00000000..f49ddd57 --- /dev/null +++ b/src/reconciler/__init__.py @@ -0,0 +1,13 @@ +"""Reconciler for self-healing background tasks like vector sync and cleanup.""" + +from .scheduler import ( + ReconcilerScheduler, + get_reconciler_scheduler, + set_reconciler_scheduler, +) + +__all__ = [ + "ReconcilerScheduler", + "get_reconciler_scheduler", + "set_reconciler_scheduler", +] diff --git a/src/reconciler/queue_cleanup.py b/src/reconciler/queue_cleanup.py new file mode 100644 index 00000000..1ac63abf --- /dev/null +++ b/src/reconciler/queue_cleanup.py @@ -0,0 +1,45 @@ +""" +Queue cleanup job. + +This module provides a periodic cleanup job that removes old processed queue items. +""" + +import logging +from datetime import datetime, timedelta, timezone + +from sqlalchemy import delete + +from src import models +from src.config import settings +from src.dependencies import tracked_db + +logger = logging.getLogger(__name__) + + +async def cleanup_queue_items() -> None: + """ + Delete processed queue items. + + Successfully processed queue items are deleted immediately, + while errored queue items are deleted after retention window. + """ + async with tracked_db("cleanup_queue_items") as db: + now = datetime.now(timezone.utc) + error_cutoff = now - timedelta( + seconds=settings.DERIVER.QUEUE_ERROR_RETENTION_SECONDS + ) + + await db.execute( + delete(models.QueueItem).where( + models.QueueItem.processed + & ( + models.QueueItem.error.is_(None) + | ( + models.QueueItem.error.is_not(None) + & (models.QueueItem.created_at < error_cutoff) + ) + ) + ) + ) + await db.commit() + logger.info("Queue cleanup completed") diff --git a/src/reconciler/scheduler.py b/src/reconciler/scheduler.py new file mode 100644 index 00000000..0fdc622d --- /dev/null +++ b/src/reconciler/scheduler.py @@ -0,0 +1,283 @@ +""" +Reconciler scheduler for self-healing background tasks. + +This module provides a scheduler for running reconciliation and cleanup tasks +like vector store sync and soft-delete cleanup. It ensures only one task of each +type runs at a time across multiple deriver instances by using the queue table +for coordination. +""" + +import asyncio +import contextlib +import logging +from datetime import datetime, timedelta, timezone + +import sentry_sdk +from pydantic import BaseModel +from sqlalchemy import exists, select +from sqlalchemy.dialects.postgresql import insert as pg_insert +from sqlalchemy.exc import IntegrityError +from sqlalchemy.ext.asyncio import AsyncSession + +from src import models +from src.config import settings +from src.dependencies import tracked_db +from src.models import QueueItem + +logger = logging.getLogger(__name__) + +# System workspace used for global reconciler tasks +SYSTEM_WORKSPACE_NAME = "__system__" + + +class ReconcilerTask(BaseModel): + """Definition of a reconciler task.""" + + name: str + work_unit_key: str + interval_seconds: int + + +# Task intervals +QUEUE_CLEANUP_INTERVAL_SECONDS = 12 * 3600 # 12 hours + +# Task registry - add new tasks here +RECONCILER_TASKS: dict[str, ReconcilerTask] = { + "sync_vectors": ReconcilerTask( + name="sync_vectors", + work_unit_key="global:sync_vectors", + interval_seconds=settings.VECTOR_STORE.RECONCILIATION_INTERVAL_SECONDS, + ), + "cleanup_queue": ReconcilerTask( + name="cleanup_queue", + work_unit_key="global:cleanup_queue", + interval_seconds=QUEUE_CLEANUP_INTERVAL_SECONDS, + ), +} + + +_reconciler_scheduler: "ReconcilerScheduler | None" = None + + +def set_reconciler_scheduler(scheduler: "ReconcilerScheduler") -> None: + """Set the global reconciler scheduler reference.""" + global _reconciler_scheduler + _reconciler_scheduler = scheduler + + +def get_reconciler_scheduler() -> "ReconcilerScheduler | None": + """Get the global reconciler scheduler reference.""" + return _reconciler_scheduler + + +class ReconcilerScheduler: + """ + Scheduler for self-healing reconciliation and cleanup tasks. + + Ensures only one task of each type runs at a time across multiple deriver + instances by using the queue table for coordination. This provides: + - Vector store synchronization (syncing pending documents/embeddings) + - Soft-delete cleanup (removing soft-deleted documents from vector stores) + - Self-healing behavior (retrying failed syncs) + - Extensible task registry for adding new maintenance tasks + + Each task type has its own interval and work_unit_key, allowing multiple + different tasks to be queued simultaneously while preventing duplicates + of the same task type. + """ + + _instance: "ReconcilerScheduler | None" = None + _initialized: bool = False + + def __new__(cls): + if cls._instance is None: + cls._instance = super().__new__(cls) + return cls._instance + + def __init__(self): + # Only initialize once + if not ReconcilerScheduler._initialized: + self._scheduler_task: asyncio.Task[None] | None = None + self._shutdown_event: asyncio.Event = asyncio.Event() + # Track next run time for each task + self._next_run: dict[str, datetime] = {} + ReconcilerScheduler._initialized = True + + @classmethod + def reset_singleton(cls) -> None: + """Reset the singleton instance. Only use this in tests.""" + cls._instance = None + cls._initialized = False + + async def start(self) -> None: + """Start the reconciler scheduler loop.""" + if self._scheduler_task is not None: + logger.warning("ReconcilerScheduler already running") + return + + self._shutdown_event.clear() + # Initialize next run times to first interval + now = datetime.now(timezone.utc) + for task_name, task in RECONCILER_TASKS.items(): + self._next_run[task_name] = now + timedelta(seconds=task.interval_seconds) + + self._scheduler_task = asyncio.create_task(self._scheduler_loop()) + logger.info( + "ReconcilerScheduler started with %d tasks: %s", + len(RECONCILER_TASKS), + list(RECONCILER_TASKS.keys()), + ) + + async def shutdown(self) -> None: + """Stop the reconciler scheduler.""" + if self._scheduler_task is None: + return + + logger.info("Shutting down ReconcilerScheduler...") + self._shutdown_event.set() + + try: + await asyncio.wait_for(self._scheduler_task, timeout=5.0) + except asyncio.TimeoutError: + logger.warning("ReconcilerScheduler shutdown timed out, cancelling task") + self._scheduler_task.cancel() + with contextlib.suppress(asyncio.CancelledError): + await self._scheduler_task + + self._scheduler_task = None + self._next_run.clear() + logger.info("ReconcilerScheduler stopped") + + async def _scheduler_loop(self) -> None: + """ + Main scheduler loop that enqueues tasks based on their intervals. + + Each task has its own interval and the loop checks all tasks on each + iteration, enqueueing any that are due. + """ + try: + while not self._shutdown_event.is_set(): + now = datetime.now(timezone.utc) + + # Check each task and enqueue if due + for task_name, task in RECONCILER_TASKS.items(): + next_run = self._next_run.get(task_name, now) + if now >= next_run: + try: + enqueued = await self._try_enqueue_task(task) + if enqueued: + logger.debug("Enqueued task: %s", task_name) + except Exception as e: + logger.exception("Error enqueueing task %s", task_name) + if settings.SENTRY.ENABLED: + sentry_sdk.capture_exception(e) + + # Schedule next run regardless of whether we enqueued + # (if already pending/in-progress, we'll skip next time too) + self._next_run[task_name] = now + timedelta( + seconds=task.interval_seconds + ) + + # Calculate sleep time until next task is due + if self._next_run: + next_task_time = min(self._next_run.values()) + sleep_seconds = max( + 1.0, # At least 1 second to avoid busy loop + (next_task_time - datetime.now(timezone.utc)).total_seconds(), + ) + else: + sleep_seconds = 60.0 # Default if no tasks + + try: + await asyncio.wait_for( + self._shutdown_event.wait(), + timeout=sleep_seconds, + ) + break # Shutdown event was set + except asyncio.TimeoutError: + # Timeout means interval elapsed, continue loop + pass + + except asyncio.CancelledError: + logger.debug("ReconcilerScheduler loop cancelled") + raise + + async def _try_enqueue_task(self, task: ReconcilerTask) -> bool: + """ + Attempt to enqueue a reconciler task. + + This is idempotent - if the task is already in-progress (has an + ActiveQueueSession) or pending in the queue, the enqueue is skipped. + + Args: + task: The task definition to enqueue + + Returns: + True if a task was enqueued, False if skipped + """ + async with tracked_db("reconciler_enqueue") as db: + # Ensure the system workspace exists (for FK constraint) + await self._ensure_system_workspace(db) + + # Check if task is already in progress + in_progress_check = select( + exists( + select(models.ActiveQueueSession.id).where( + models.ActiveQueueSession.work_unit_key == task.work_unit_key + ) + ) + ) + is_in_progress = await db.scalar(in_progress_check) + + if is_in_progress: + logger.debug("Task %s already in progress, skipping enqueue", task.name) + return False + + # Check if there's already a pending task + pending_check = select( + exists( + select(QueueItem.id).where( + QueueItem.work_unit_key == task.work_unit_key, + QueueItem.processed == False, # noqa: E712 + ) + ) + ) + is_pending = await db.scalar(pending_check) + + if is_pending: + logger.debug( + "Task %s already pending in queue, skipping enqueue", task.name + ) + return False + + # Enqueue the task using ORM + queue_item = QueueItem( + work_unit_key=task.work_unit_key, + payload={ + "reconciler_type": task.name, + }, + session_id=None, + task_type="reconciler", + workspace_name=SYSTEM_WORKSPACE_NAME, + message_id=None, + ) + db.add(queue_item) + try: + await db.commit() + except IntegrityError: + # Another instance already enqueued this task (unique constraint) + await db.rollback() + logger.debug( + "Task %s already enqueued by another instance, skipping", task.name + ) + return False + + logger.info("Enqueued reconciler task: %s", task.name) + return True + + async def _ensure_system_workspace(self, db: AsyncSession) -> None: + """Ensure the system workspace exists for reconciler tasks.""" + # Use upsert to create workspace if it doesn't exist + stmt = pg_insert(models.Workspace).values(name=SYSTEM_WORKSPACE_NAME) + stmt = stmt.on_conflict_do_nothing(index_elements=["name"]) + await db.execute(stmt) diff --git a/src/deriver/vector_reconciliation.py b/src/reconciler/sync_vectors.py similarity index 100% rename from src/deriver/vector_reconciliation.py rename to src/reconciler/sync_vectors.py diff --git a/src/schemas.py b/src/schemas.py index 9b482c7f..c747d3bc 100644 --- a/src/schemas.py +++ b/src/schemas.py @@ -27,6 +27,13 @@ class DreamType(str, Enum): OMNI = "omni" +class ReconcilerType(str, Enum): + """Types of reconciler tasks that can be performed.""" + + SYNC_VECTORS = "sync_vectors" + CLEANUP_QUEUE = "cleanup_queue" + + class ReasoningConfiguration(BaseModel): enabled: bool | None = Field( default=None, diff --git a/src/utils/queue_payload.py b/src/utils/queue_payload.py index 9deefb72..002a5f8f 100644 --- a/src/utils/queue_payload.py +++ b/src/utils/queue_payload.py @@ -3,7 +3,7 @@ from typing import Any, Literal from pydantic import BaseModel, ConfigDict -from src.schemas import DreamType, ResolvedConfiguration +from src.schemas import DreamType, ReconcilerType, ResolvedConfiguration class BasePayload(BaseModel): @@ -67,6 +67,13 @@ class DeletionPayload(BasePayload): resource_id: str +class ReconcilerPayload(BasePayload): + """Payload for reconciler tasks (vector sync, queue cleanup, self-healing).""" + + task_type: Literal["reconciler"] = "reconciler" + reconciler_type: ReconcilerType + + def create_webhook_payload( event_type: str, data: dict[str, Any], diff --git a/src/utils/types.py b/src/utils/types.py index 8e464be4..f9efffd1 100644 --- a/src/utils/types.py +++ b/src/utils/types.py @@ -11,6 +11,8 @@ class GetOrCreateResult(NamedTuple, Generic[T]): SupportedProviders = Literal["anthropic", "openai", "google", "groq", "custom", "vllm"] -TaskType = Literal["webhook", "summary", "representation", "dream", "deletion"] +TaskType = Literal[ + "webhook", "summary", "representation", "dream", "deletion", "reconciler" +] VectorSyncState = Literal["synced", "pending", "failed"] DocumentLevel = Literal["explicit", "deductive", "inductive", "contradiction"] diff --git a/tests/alembic/revisions/test_119a52b73c60_support_external_embeddings.py b/tests/alembic/revisions/test_119a52b73c60_support_external_embeddings.py index ee87d61b..8a88a4a3 100644 --- a/tests/alembic/revisions/test_119a52b73c60_support_external_embeddings.py +++ b/tests/alembic/revisions/test_119a52b73c60_support_external_embeddings.py @@ -12,6 +12,7 @@ INDEXES = ( ("documents", "ix_documents_sync_state_last_sync_at"), ("message_embeddings", "ix_message_embeddings_sync_state"), ("message_embeddings", "ix_message_embeddings_sync_state_last_sync_at"), + ("queue", "uq_queue_work_unit_key"), ) diff --git a/tests/deriver/test_vector_reconciliation.py b/tests/deriver/test_vector_reconciliation.py index 728f99d5..bd3ff28d 100644 --- a/tests/deriver/test_vector_reconciliation.py +++ b/tests/deriver/test_vector_reconciliation.py @@ -15,7 +15,7 @@ from sqlalchemy import select from sqlalchemy.ext.asyncio import AsyncSession from src import models -from src.deriver.vector_reconciliation import ( +from src.reconciler.sync_vectors import ( MAX_SYNC_ATTEMPTS, ReconciliationMetrics, _get_documents_needing_sync, # pyright: ignore[reportPrivateUsage]