diff --git a/src/crud/document.py b/src/crud/document.py index d7264f8a..b423c77a 100644 --- a/src/crud/document.py +++ b/src/crud/document.py @@ -8,13 +8,6 @@ from sqlalchemy.exc import IntegrityError from sqlalchemy.ext.asyncio import AsyncSession from sqlalchemy.sql import Select from sqlalchemy.sql.functions import func -from tenacity import ( - AsyncRetrying, - retry_if_exception_type, - retry_if_result, - stop_after_attempt, - wait_exponential, -) from src import models, schemas from src.config import settings @@ -24,7 +17,12 @@ from src.crud.session import get_session from src.embedding_client import embedding_client from src.exceptions import ResourceNotFoundException, ValidationException from src.utils.filter import apply_filter -from src.vector_store import VectorRecord, VectorStore, get_vector_store +from src.vector_store import ( + VectorRecord, + VectorStore, + get_vector_store, + upsert_with_retry, +) logger = getLogger(__name__) @@ -334,26 +332,13 @@ async def create_documents( ) ) - # Retry vector upsert with exponential backoff (3 attempts) + # Upsert to vector store with retry and update sync state try: - result = None - async for attempt in AsyncRetrying( - stop=stop_after_attempt(3), - wait=wait_exponential(multiplier=0.5, min=0.5, max=2.0), - retry=retry_if_exception_type(Exception) - | retry_if_result( - lambda res: res is not None and res.secondary_ok is False - ), - reraise=True, - ): - with attempt: - result = await vector_store.upsert_many( - namespace, vector_records - ) - + result = await upsert_with_retry( + vector_store, namespace, vector_records + ) if result is not None and result.secondary_ok is False: # Partial success: primary has data but secondary doesn't - # Keep as "pending" for reconciliation to sync secondary logger.warning( f"Partial sync for namespace {namespace}: {result.secondary_error}" ) @@ -380,9 +365,8 @@ async def create_documents( await db.commit() except Exception as e: - # Total failure: primary write failed - # Keep as "pending" for reconciliation to retry - logger.error(f"Failed to upsert vectors after 3 retries: {e}") + # Total failure: primary write failed after retries + logger.error(f"Failed to upsert vectors after retries: {e}") await db.execute( update(models.Document) .where(models.Document.id.in_(doc_ids)) @@ -414,9 +398,12 @@ async def delete_document( """ Delete a single document by ID using hybrid sync/soft delete pattern. - Tries to delete from vector store first, then hard deletes from DB. - If vector store delete fails, soft deletes (sets deleted_at) and lets - cleanup job handle vector deletion later. + Soft deletes first (sets deleted_at), then tries to delete from vector store. + If vector store delete succeeds, hard deletes from DB. + If vector store delete fails, leaves soft-deleted for cleanup job. + + This order ensures crash safety: if the process crashes at any point, + the document is either fully deleted or soft-deleted (never orphaned). Args: db: Database session @@ -450,7 +437,16 @@ async def delete_document( f"Document {document_id} not found or does not belong to the specified collection/session" ) - # Try to delete from vector store first + # Step 1: Soft delete first (crash-safe - ensures document is marked for deletion) + update_stmt = ( + update(models.Document) + .where(models.Document.id == document_id) + .values(deleted_at=func.now()) + ) + await db.execute(update_stmt) + await db.commit() + + # Step 2: Try to delete from vector store vector_store = get_vector_store() namespace = vector_store.get_vector_namespace( "document", workspace_name, observer, observed @@ -463,20 +459,12 @@ async def delete_document( except Exception as e: logger.warning(f"Failed to delete vector for document {document_id}: {e}") + # Step 3: If vector deleted successfully, hard delete from DB if vector_deleted: - # Happy path: hard delete from DB delete_stmt = delete(models.Document).where(models.Document.id == document_id) await db.execute(delete_stmt) - else: - # Fallback: soft delete, let cleanup job handle vector - update_stmt = ( - update(models.Document) - .where(models.Document.id == document_id) - .values(deleted_at=func.now()) - ) - await db.execute(update_stmt) - - await db.commit() + await db.commit() + # If vector delete failed, document stays soft-deleted for cleanup job async def delete_document_by_id( @@ -487,9 +475,12 @@ async def delete_document_by_id( """ Delete a single document by ID and workspace using hybrid sync/soft delete pattern. - Tries to delete from vector store first, then hard deletes from DB. - If vector store delete fails, soft deletes (sets deleted_at) and lets - cleanup job handle vector deletion later. + Soft deletes first (sets deleted_at), then tries to delete from vector store. + If vector store delete succeeds, hard deletes from DB. + If vector store delete fails, leaves soft-deleted for cleanup job. + + This order ensures crash safety: if the process crashes at any point, + the document is either fully deleted or soft-deleted (never orphaned). Args: db: Database session @@ -513,7 +504,16 @@ async def delete_document_by_id( f"Document {document_id} not found or does not belong to workspace {workspace_name}" ) - # Try to delete from vector store first + # Step 1: Soft delete first (crash-safe - ensures document is marked for deletion) + update_stmt = ( + update(models.Document) + .where(models.Document.id == document_id) + .values(deleted_at=func.now()) + ) + await db.execute(update_stmt) + await db.commit() + + # Step 2: Try to delete from vector store vector_store = get_vector_store() namespace = vector_store.get_vector_namespace( "document", @@ -529,20 +529,12 @@ async def delete_document_by_id( except Exception as e: logger.warning(f"Failed to delete vector for document {document_id}: {e}") + # Step 3: If vector deleted successfully, hard delete from DB if vector_deleted: - # Happy path: hard delete from DB delete_stmt = delete(models.Document).where(models.Document.id == document_id) await db.execute(delete_stmt) - else: - # Fallback: soft delete, let cleanup job handle vector - update_stmt = ( - update(models.Document) - .where(models.Document.id == document_id) - .values(deleted_at=func.now()) - ) - await db.execute(update_stmt) - - await db.commit() + await db.commit() + # If vector delete failed, document stays soft-deleted for cleanup job async def create_observations( @@ -688,26 +680,13 @@ async def create_observations( ) ) - # Retry vector upsert with exponential backoff (3 attempts) + # Upsert to vector store with retry and update sync state try: - result = None - async for attempt in AsyncRetrying( - stop=stop_after_attempt(3), - wait=wait_exponential(multiplier=0.5, min=0.5, max=2.0), - retry=retry_if_exception_type(Exception) - | retry_if_result( - lambda res: res is not None and res.secondary_ok is False - ), - reraise=True, - ): - with attempt: - result = await vector_store.upsert_many( - namespace, vector_records - ) - + result = await upsert_with_retry( + vector_store, namespace, vector_records + ) if result is not None and result.secondary_ok is False: # Partial success: primary has data but secondary doesn't - # Keep as "pending" for reconciliation to sync secondary logger.warning( f"Partial sync for namespace {namespace}: {result.secondary_error}" ) @@ -734,10 +713,9 @@ async def create_observations( await db.commit() except Exception as e: - # Total failure: primary write failed - # Keep as "pending" for reconciliation to retry + # Total failure: primary write failed after retries logger.error( - f"Failed to upsert vectors for {namespace} after 3 retries: {e}" + f"Failed to upsert vectors for {namespace} after retries: {e}" ) await db.execute( update(models.Document) diff --git a/src/crud/message.py b/src/crud/message.py index 58d226db..31f7e6e8 100644 --- a/src/crud/message.py +++ b/src/crud/message.py @@ -4,19 +4,12 @@ from typing import Any from nanoid import generate as generate_nanoid from sqlalchemy import ColumnElement, Select, and_, func, select, text, update from sqlalchemy.ext.asyncio import AsyncSession -from tenacity import ( - AsyncRetrying, - retry_if_exception_type, - retry_if_result, - stop_after_attempt, - wait_exponential, -) from src import models, schemas from src.config import settings from src.embedding_client import embedding_client from src.utils.filter import apply_filter -from src.vector_store import VectorRecord, get_vector_store +from src.vector_store import VectorRecord, get_vector_store, upsert_with_retry from .session import get_or_create_session @@ -223,28 +216,14 @@ async def create_messages( ) await db.commit() - # Retry vector upsert with exponential backoff + # Upsert to vector store with retry and update sync state if vector_records: try: - result = None - async for attempt in AsyncRetrying( - stop=stop_after_attempt(3), - wait=wait_exponential(multiplier=0.5, min=0.5, max=2.0), - retry=retry_if_exception_type(Exception) - | retry_if_result( - lambda res: res is not None - and res.secondary_ok is False - ), - reraise=True, - ): - with attempt: - result = await vector_store.upsert_many( - namespace, vector_records - ) - + result = await upsert_with_retry( + vector_store, namespace, vector_records + ) if result is not None and result.secondary_ok is False: # Partial success: primary has data but secondary doesn't - # Keep as "pending" for reconciliation to sync secondary logger.warning( "Partial sync for message embeddings: %s", result.secondary_error, @@ -260,7 +239,7 @@ async def create_messages( ) await db.commit() else: - # Success: primary succeeded and (secondary succeeded or no secondary configured) + # Success: both primary and secondary stores have the data await db.execute( update(models.MessageEmbedding) .where(models.MessageEmbedding.id.in_(embedding_ids)) @@ -273,9 +252,9 @@ async def create_messages( await db.commit() except Exception as e: - # Total failure: primary write failed + # Total failure: primary write failed after retries logger.error( - f"Failed to upsert message vectors after 3 retries: {e}" + f"Failed to upsert message vectors after retries: {e}" ) await db.execute( update(models.MessageEmbedding) diff --git a/src/deriver/vector_reconciliation.py b/src/deriver/vector_reconciliation.py index 1f2375d4..7f1e7042 100644 --- a/src/deriver/vector_reconciliation.py +++ b/src/deriver/vector_reconciliation.py @@ -23,9 +23,7 @@ from src.vector_store import VectorRecord, VectorStore, get_vector_store logger = logging.getLogger(__name__) # Constants -RECONCILIATION_BATCH_SIZE = ( - 30 # Keep batch size small to avoid exceeding embedding API limits -) +RECONCILIATION_BATCH_SIZE = 50 RECONCILIATION_TIME_BUDGET_SECONDS = 240 # Leave headroom for other maintenance work MAX_SYNC_ATTEMPTS = 5 # After this many failures, mark as permanently_failed diff --git a/src/vector_store/__init__.py b/src/vector_store/__init__.py index 8ff785dd..3b1c95ea 100644 --- a/src/vector_store/__init__.py +++ b/src/vector_store/__init__.py @@ -177,6 +177,7 @@ from src.vector_store.composite import CompositeVectorStore # noqa: E402 from src.vector_store.lancedb import LanceDBVectorStore # noqa: E402 from src.vector_store.pgvector import PgVectorStore # noqa: E402 from src.vector_store.turbopuffer import TurbopufferVectorStore # noqa: E402 +from src.vector_store.utils import upsert_with_retry # noqa: E402 def _create_store_by_type(store_type: str) -> VectorStore: @@ -256,4 +257,5 @@ __all__ = [ "VectorUpsertResult", "get_vector_store", "close_vector_store", + "upsert_with_retry", ] diff --git a/src/vector_store/pgvector.py b/src/vector_store/pgvector.py index d3854eaf..bc0b4701 100644 --- a/src/vector_store/pgvector.py +++ b/src/vector_store/pgvector.py @@ -12,7 +12,7 @@ from sqlalchemy import select, update from sqlalchemy.ext.asyncio import AsyncSession from src import models -from src.db import SessionLocal +from src.dependencies import tracked_db from . import VectorQueryResult, VectorRecord, VectorStore, VectorUpsertResult @@ -71,10 +71,6 @@ class PgVectorStore(VectorStore): else: raise ValueError(f"Invalid namespace format: {namespace}") - async def _get_session(self) -> AsyncSession: - """Get a database session.""" - return SessionLocal() - async def upsert_many( self, namespace: str, @@ -126,29 +122,27 @@ class PgVectorStore(VectorStore): """ table_type, context = self._parse_namespace(namespace) - db = await self._get_session() - try: - if table_type == "documents": - results = await self._query_documents( - db, context, embedding, top_k, filters, max_distance - ) - elif table_type == "message_embeddings": - results = await self._query_message_embeddings( - db, context, embedding, top_k, filters, max_distance - ) - else: - results = [] + async with tracked_db("pgvector_query") as db: + try: + if table_type == "documents": + results = await self._query_documents( + db, context, embedding, top_k, filters, max_distance + ) + elif table_type == "message_embeddings": + results = await self._query_message_embeddings( + db, context, embedding, top_k, filters, max_distance + ) + else: + results = [] - logger.debug( - f"Query returned {len(results)} results from namespace {namespace}" - ) - return results + logger.debug( + f"Query returned {len(results)} results from namespace {namespace}" + ) + return results - except Exception: - logger.exception(f"Failed to query namespace {namespace}") - raise - finally: - await db.close() + except Exception: + logger.exception(f"Failed to query namespace {namespace}") + raise async def _query_documents( self, @@ -297,45 +291,43 @@ class PgVectorStore(VectorStore): table_type, _ = self._parse_namespace(namespace) - db = await self._get_session() - try: - if table_type == "documents": - stmt = ( - update(models.Document) - .where(models.Document.id.in_(ids)) - .values(embedding=None) - ) - await db.execute(stmt) - - elif table_type == "message_embeddings": - for vector_id in ids: - try: - embedding_id = int(vector_id) - except ValueError as exc: - raise ValueError( - f"Invalid message vector id format: {vector_id}" - ) from exc - + async with tracked_db("pgvector_delete") as db: + try: + if table_type == "documents": stmt = ( - update(models.MessageEmbedding) - .where(models.MessageEmbedding.id == embedding_id) + update(models.Document) + .where(models.Document.id.in_(ids)) .values(embedding=None) ) await db.execute(stmt) - await db.commit() - logger.debug( - f"Deleted {len(ids)} vectors from {table_type} in namespace {namespace}" - ) + elif table_type == "message_embeddings": + for vector_id in ids: + try: + embedding_id = int(vector_id) + except ValueError as exc: + raise ValueError( + f"Invalid message vector id format: {vector_id}" + ) from exc - except Exception: - await db.rollback() - logger.exception( - f"Failed to delete {len(ids)} vectors from namespace {namespace}" - ) - raise - finally: - await db.close() + stmt = ( + update(models.MessageEmbedding) + .where(models.MessageEmbedding.id == embedding_id) + .values(embedding=None) + ) + await db.execute(stmt) + + await db.commit() + logger.debug( + f"Deleted {len(ids)} vectors from {table_type} in namespace {namespace}" + ) + + except Exception: + await db.rollback() + logger.exception( + f"Failed to delete {len(ids)} vectors from namespace {namespace}" + ) + raise async def delete_namespace(self, namespace: str) -> None: """ @@ -346,38 +338,38 @@ class PgVectorStore(VectorStore): """ table_type, context = self._parse_namespace(namespace) - db = await self._get_session() - try: - if table_type == "documents": - stmt = ( - update(models.Document) - .where(models.Document.workspace_name == context["workspace_name"]) - .where(models.Document.observer == context["observer"]) - .where(models.Document.observed == context["observed"]) - .values(embedding=None) - ) - await db.execute(stmt) - - elif table_type == "message_embeddings": - stmt = ( - update(models.MessageEmbedding) - .where( - models.MessageEmbedding.workspace_name - == context["workspace_name"] + async with tracked_db("pgvector_delete_namespace") as db: + try: + if table_type == "documents": + stmt = ( + update(models.Document) + .where( + models.Document.workspace_name == context["workspace_name"] + ) + .where(models.Document.observer == context["observer"]) + .where(models.Document.observed == context["observed"]) + .values(embedding=None) ) - .values(embedding=None) - ) - await db.execute(stmt) + await db.execute(stmt) - await db.commit() - logger.debug(f"Deleted all vectors from namespace {namespace}") + elif table_type == "message_embeddings": + stmt = ( + update(models.MessageEmbedding) + .where( + models.MessageEmbedding.workspace_name + == context["workspace_name"] + ) + .values(embedding=None) + ) + await db.execute(stmt) - except Exception: - await db.rollback() - logger.exception(f"Failed to delete namespace {namespace}") - raise - finally: - await db.close() + await db.commit() + logger.debug(f"Deleted all vectors from namespace {namespace}") + + except Exception: + await db.rollback() + logger.exception(f"Failed to delete namespace {namespace}") + raise async def close(self) -> None: """Close the pgvector store (no-op for pgvector)""" diff --git a/src/vector_store/utils.py b/src/vector_store/utils.py new file mode 100644 index 00000000..8462f4ff --- /dev/null +++ b/src/vector_store/utils.py @@ -0,0 +1,61 @@ +""" +Vector store utility functions. +""" + +from __future__ import annotations + +import logging +from typing import TYPE_CHECKING + +from tenacity import ( + AsyncRetrying, + retry_if_exception_type, + retry_if_result, + stop_after_attempt, + wait_exponential, +) + +if TYPE_CHECKING: + from src.vector_store import VectorRecord, VectorStore, VectorUpsertResult + +logger = logging.getLogger(__name__) + + +async def upsert_with_retry( + vector_store: VectorStore, + namespace: str, + vector_records: list[VectorRecord], + max_attempts: int = 3, +) -> VectorUpsertResult | None: + """ + Upsert vectors with exponential backoff retry. + + Retries on any exception or when secondary store fails (partial success). + + Args: + vector_store: The vector store to upsert into + namespace: The namespace for the vectors + vector_records: List of VectorRecord objects to upsert + max_attempts: Maximum number of retry attempts (default 3) + + Returns: + VectorUpsertResult on success, or None if vector_records is empty + + Raises: + Exception: If all retries fail + """ + if not vector_records: + return None + + result: VectorUpsertResult | None = None + async for attempt in AsyncRetrying( + stop=stop_after_attempt(max_attempts), + wait=wait_exponential(multiplier=0.5, min=0.5, max=2.0), + retry=retry_if_exception_type(Exception) + | retry_if_result(lambda res: res is not None and res.secondary_ok is False), + reraise=True, + ): + with attempt: + result = await vector_store.upsert_many(namespace, vector_records) + + return result