RomanNum3ral
/
honcho
mirror of https://github.com/plastic-labs/honcho.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

fix: centralize retry logic; bump reconciliation batch; use tracked db; fix soft-delete race condition

This commit is contained in:
Rajat Ahuja 2026-01-09 10:47:48 -05:00
parent ed2b8c46f3
commit 708f0d461d
6 changed files with 208 additions and 198 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

134
src/crud/document.py
View File

@ -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)

37
src/crud/message.py
View File

@ -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)

4
src/deriver/vector_reconciliation.py
View File

@ -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

2
src/vector_store/__init__.py
View File

@ -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",
]

168
src/vector_store/pgvector.py
View File

@ -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)"""

61
src/vector_store/utils.py Normal file
View File

@ -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