From b04770d602f33f48766f6b929615e37165030081 Mon Sep 17 00:00:00 2001 From: Rajat Ahuja Date: Mon, 12 Jan 2026 18:11:22 -0500 Subject: [PATCH] fix: remove composite and pgvector store -> make more specific --- .env.template | 11 +- config.toml.example | 6 +- src/config.py | 16 +- src/crud/document.py | 283 +++++++++--------- src/crud/message.py | 196 +++++------- src/crud/session.py | 24 +- src/crud/workspace.py | 65 ++-- src/deriver/vector_reconciliation.py | 131 ++++---- src/main.py | 4 +- src/utils/search.py | 22 +- src/vector_store/__init__.py | 67 ++--- src/vector_store/composite.py | 261 ---------------- src/vector_store/lancedb.py | 4 +- src/vector_store/pgvector.py | 315 -------------------- src/vector_store/turbopuffer.py | 4 +- src/vector_store/utils.py | 6 +- tests/conftest.py | 23 +- tests/deriver/test_vector_reconciliation.py | 26 +- 18 files changed, 429 insertions(+), 1035 deletions(-) delete mode 100644 src/vector_store/composite.py delete mode 100644 src/vector_store/pgvector.py diff --git a/.env.template b/.env.template index 1b7a139b..b6b17d9e 100644 --- a/.env.template +++ b/.env.template @@ -191,12 +191,11 @@ LLM_ANTHROPIC_API_KEY=your-anthropic-api-key-here # VECTOR_STORE_LANCEDB_PATH="./lancedb_data" # Vector store settings -# Primary vector store type: "pgvector", "turbopuffer", or "lancedb" -VECTOR_STORE_PRIMARY_TYPE=pgvector +# Vector store type: "pgvector", "turbopuffer", or "lancedb" +VECTOR_STORE_TYPE=pgvector -# Secondary vector store type (optional) -# When set, enables dual-write and fallback read -# VECTOR_STORE_SECONDARY_TYPE= +# Migration flag: set to true when migration from pgvector is complete +VECTOR_STORE_MIGRATED=false # Global namespace prefix for all vector namespaces # Namespaces follow the pattern: @@ -207,7 +206,7 @@ VECTOR_STORE_NAMESPACE=honcho # Embedding dimensions (default: 1536 for OpenAI text-embedding-3-small) VECTOR_STORE_DIMENSIONS=1536 -# Turbopuffer-specific settings (required if PRIMARY_TYPE or SECONDARY_TYPE is "turbopuffer") +# Turbopuffer-specific settings (required if TYPE is "turbopuffer") # VECTOR_STORE_TURBOPUFFER_API_KEY=your-turbopuffer-api-key # VECTOR_STORE_TURBOPUFFER_REGION=us-east-1 diff --git a/config.toml.example b/config.toml.example index 762ca703..b124d8f4 100644 --- a/config.toml.example +++ b/config.toml.example @@ -144,8 +144,10 @@ DEFAULT_LOCK_TTL_SECONDS = 5 # Vector store settings [vector_store] -PRIMARY_TYPE = "pgvector" -# SECONDARY_TYPE = "lancedb" +# Vector store type: "pgvector", "turbopuffer", or "lancedb" +TYPE = "pgvector" +# Migration flag: set to true when migration from pgvector is complete +MIGRATED = false NAMESPACE = "honcho" DIMENSIONS = 1536 # TURBOPUFFER_API_KEY = "your-turbopuffer-api-key" diff --git a/src/config.py b/src/config.py index b67179a3..a3ead2c0 100644 --- a/src/config.py +++ b/src/config.py @@ -362,14 +362,10 @@ class VectorStoreSettings(HonchoSettings): model_config = SettingsConfigDict(env_prefix="VECTOR_STORE_", extra="ignore") # pyright: ignore - # Primary vector store type - PRIMARY_TYPE: Literal["pgvector", "turbopuffer", "lancedb"] = "pgvector" + # Vector store type to use + TYPE: Literal["pgvector", "turbopuffer", "lancedb"] = "pgvector" - # Secondary vector store type (optional) - # When set, enables: - # - Dual-write: writes go to both primary and secondary - # - Fallback read: reads try primary first, fall back to secondary if empty - SECONDARY_TYPE: Literal["pgvector", "turbopuffer", "lancedb"] | None = None + MIGRATED: bool = False # Global namespace prefix for all vector namespaces # Namespaces follow the pattern: @@ -394,11 +390,9 @@ class VectorStoreSettings(HonchoSettings): @model_validator(mode="after") def _require_api_key_for_turbopuffer(self) -> "VectorStoreSettings": - if ( - self.PRIMARY_TYPE == "turbopuffer" or self.SECONDARY_TYPE == "turbopuffer" - ) and not self.TURBOPUFFER_API_KEY: + if self.TYPE == "turbopuffer" and not self.TURBOPUFFER_API_KEY: raise ValueError( - "VECTOR_STORE_TURBOPUFFER_API_KEY must be set when PRIMARY_TYPE or SECONDARY_TYPE is 'turbopuffer'" + "VECTOR_STORE_TURBOPUFFER_API_KEY must be set when TYPE is 'turbopuffer'" ) return self diff --git a/src/crud/document.py b/src/crud/document.py index 2260ac28..d88c1f89 100644 --- a/src/crud/document.py +++ b/src/crud/document.py @@ -20,7 +20,7 @@ from src.utils.filter import apply_filter from src.vector_store import ( VectorRecord, VectorStore, - get_vector_store, + get_external_vector_store, upsert_with_retry, ) @@ -151,9 +151,9 @@ async def query_documents( f"Query exceeds maximum token limit of {settings.MAX_EMBEDDING_TOKENS}." ) from e - # If pgvector is primary, query Postgres directly with similarity + filters - # This avoids duplicate fetches from the same database - if settings.VECTOR_STORE.PRIMARY_TYPE == "pgvector": + # Query Postgres directly when using pgvector OR during migration (not yet migrated) + # This ensures we use pgvector as source of truth until migration is complete + if settings.VECTOR_STORE.TYPE == "pgvector" or not settings.VECTOR_STORE.MIGRATED: stmt = ( select(models.Document) .where(models.Document.workspace_name == workspace_name) @@ -176,9 +176,12 @@ async def query_documents( result = await db.execute(stmt) return list(result.scalars().all()) - # FALLBACK: Use vector store abstraction for external stores (Turbopuffer, LanceDB) - vector_store = get_vector_store() - namespace = vector_store.get_vector_namespace( + # FALLBACK: Use external vector store (Turbopuffer, LanceDB) + external_vector_store = get_external_vector_store() + if external_vector_store is None: + return [] + + namespace = external_vector_store.get_vector_namespace( "document", workspace_name, observer, observed ) @@ -193,8 +196,8 @@ async def query_documents( if key in filters: vector_filters[key] = filters[key] - # Query vector store for similar documents with filters applied - vector_results = await vector_store.query( + # Query external vector store for similar documents with filters applied + vector_results = await external_vector_store.query( namespace, embedding, top_k=top_k, @@ -273,15 +276,14 @@ async def create_documents( metadata_dict = doc.metadata.model_dump(exclude_none=True) - # Check if pgvector is being used (primary or secondary) - # If so, write embeddings to ORM since pgvector relies on postgres - pgvector_in_use = ( - settings.VECTOR_STORE.PRIMARY_TYPE == "pgvector" - or settings.VECTOR_STORE.SECONDARY_TYPE == "pgvector" + # Determine if we need to persist embeddings to postgres + # True when: TYPE=pgvector OR still migrating (dual-write to both stores) + store_embeddings_in_postgres = ( + settings.VECTOR_STORE.TYPE == "pgvector" + or not settings.VECTOR_STORE.MIGRATED ) - if pgvector_in_use and doc.embedding: - # pgvector in use: write embedding to ORM (postgres) + if store_embeddings_in_postgres and doc.embedding: new_doc = models.Document( workspace_name=workspace_name, observer=observer, @@ -294,7 +296,6 @@ async def create_documents( embedding=doc.embedding, ) else: - # pgvector not in use or no embedding: don't write embedding to postgres new_doc = models.Document( workspace_name=workspace_name, observer=observer, @@ -328,56 +329,55 @@ async def create_documents( # The reconciliation job will automatically re-embed and sync these documents, await db.commit() - # Store embeddings in vector store after documents are committed (IDs now available) + # Store embeddings in external vector store after documents are committed (IDs now available) if docs_with_embeddings: - vector_store = get_vector_store() - namespace = vector_store.get_vector_namespace( - "document", - workspace_name, - observer, - observed, - ) + doc_ids = [doc.id for doc, _ in docs_with_embeddings] + external_vector_store = get_external_vector_store() - # Build vector records with metadata for filtering - vector_records: list[VectorRecord] = [] - doc_ids: list[str] = [] - for doc, embedding in docs_with_embeddings: - doc_ids.append(doc.id) - vector_records.append( - VectorRecord( - id=doc.id, - embedding=embedding, - metadata={ - "workspace_name": workspace_name, - "observer": observer, - "observed": observed, - "session_name": doc.session_name, - "level": doc.level, - }, + # If no external vector store (pgvector mode), mark as synced immediately + if external_vector_store is None: + await db.execute( + update(models.Document) + .where(models.Document.id.in_(doc_ids)) + .values( + sync_state="synced", + last_sync_at=func.now(), + sync_attempts=0, ) ) - - # Upsert to vector store with retry and update sync state - try: - result = await upsert_with_retry( - vector_store, namespace, vector_records + await db.commit() + else: + # External vector store - upsert and track sync state + namespace = external_vector_store.get_vector_namespace( + "document", + workspace_name, + observer, + observed, ) - if result is not None and result.secondary_ok is False: - # Partial success: primary has data but secondary doesn't - logger.warning( - f"Partial sync for namespace {namespace}: {result.secondary_error}" - ) - await db.execute( - update(models.Document) - .where(models.Document.id.in_(doc_ids)) - .values( - sync_attempts=models.Document.sync_attempts + 1, - last_sync_at=func.now(), + + # Build vector records with metadata for filtering + vector_records: list[VectorRecord] = [] + for doc, embedding in docs_with_embeddings: + vector_records.append( + VectorRecord( + id=doc.id, + embedding=embedding, + metadata={ + "workspace_name": workspace_name, + "observer": observer, + "observed": observed, + "session_name": doc.session_name, + "level": doc.level, + }, ) ) - await db.commit() - else: - # Success: both primary and secondary stores have the data + + # Upsert to external vector store with retry and update sync state + try: + await upsert_with_retry( + external_vector_store, namespace, vector_records + ) + # Success: mark as synced await db.execute( update(models.Document) .where(models.Document.id.in_(doc_ids)) @@ -389,18 +389,18 @@ async def create_documents( ) await db.commit() - except Exception as 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)) - .values( - sync_attempts=models.Document.sync_attempts + 1, - last_sync_at=func.now(), + except Exception as e: + # Failed after retries - increment sync_attempts for reconciliation + logger.error(f"Failed to upsert vectors after retries: {e}") + await db.execute( + update(models.Document) + .where(models.Document.id.in_(doc_ids)) + .values( + sync_attempts=models.Document.sync_attempts + 1, + last_sync_at=func.now(), + ) ) - ) - await db.commit() + await db.commit() except IntegrityError as e: await db.rollback() @@ -563,16 +563,14 @@ async def create_observations( tuple[str, str], list[tuple[models.Document, list[float]]] ] = {} - # Check if pgvector is being used (primary or secondary) - # If so, write embeddings to ORM since pgvector relies on postgres - pgvector_in_use = ( - settings.VECTOR_STORE.PRIMARY_TYPE == "pgvector" - or settings.VECTOR_STORE.SECONDARY_TYPE == "pgvector" + # Determine if we need to persist embeddings to postgres + # True when: TYPE=pgvector OR still migrating (dual-write to both stores) + store_embeddings_in_postgres = ( + settings.VECTOR_STORE.TYPE == "pgvector" or not settings.VECTOR_STORE.MIGRATED ) for obs, embedding in zip(observations, embeddings, strict=True): - if pgvector_in_use: - # pgvector in use: write embedding to ORM (postgres) + if store_embeddings_in_postgres: doc = models.Document( workspace_name=workspace_name, observer=obs.observer_id, @@ -585,7 +583,6 @@ async def create_observations( embedding=embedding, ) else: - # pgvector not in use: don't write embedding to postgres doc = models.Document( workspace_name=workspace_name, observer=obs.observer_id, @@ -612,56 +609,60 @@ async def create_observations( for doc in honcho_documents: await db.refresh(doc) - # Store embeddings in vector store after documents are committed (IDs now available) - vector_store = get_vector_store() - for (observer, observed), docs_with_embeddings in collection_embeddings.items(): - namespace = vector_store.get_vector_namespace( - "document", - workspace_name, + # Store embeddings in external vector store after documents are committed (IDs now available) + external_vector_store = get_external_vector_store() + all_doc_ids = [doc.id for doc in honcho_documents] + + # If no external vector store (pgvector mode), mark as synced immediately + if external_vector_store is None: + await db.execute( + update(models.Document) + .where(models.Document.id.in_(all_doc_ids)) + .values( + sync_state="synced", + last_sync_at=func.now(), + sync_attempts=0, + ) + ) + await db.commit() + else: + # External vector store - upsert each collection's embeddings + for ( observer, observed, - ) - - # Build vector records with metadata for filtering - vector_records: list[VectorRecord] = [] - doc_ids: list[str] = [] - for doc, embedding in docs_with_embeddings: - doc_ids.append(doc.id) - vector_records.append( - VectorRecord( - id=doc.id, - embedding=embedding, - metadata={ - "workspace_name": workspace_name, - "observer": observer, - "observed": observed, - "session_name": doc.session_name, - "level": doc.level, - }, - ) + ), docs_with_embeddings in collection_embeddings.items(): + namespace = external_vector_store.get_vector_namespace( + "document", + workspace_name, + observer, + observed, ) - # Upsert to vector store with retry and update sync state - try: - 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 - logger.warning( - f"Partial sync for namespace {namespace}: {result.secondary_error}" - ) - await db.execute( - update(models.Document) - .where(models.Document.id.in_(doc_ids)) - .values( - sync_attempts=models.Document.sync_attempts + 1, - last_sync_at=func.now(), + # Build vector records with metadata for filtering + vector_records: list[VectorRecord] = [] + doc_ids: list[str] = [] + for doc, embedding in docs_with_embeddings: + doc_ids.append(doc.id) + vector_records.append( + VectorRecord( + id=doc.id, + embedding=embedding, + metadata={ + "workspace_name": workspace_name, + "observer": observer, + "observed": observed, + "session_name": doc.session_name, + "level": doc.level, + }, ) ) - await db.commit() - else: - # Success: both primary and secondary stores have the data + + # Upsert to external vector store with retry and update sync state + try: + await upsert_with_retry( + external_vector_store, namespace, vector_records + ) + # Success: mark as synced await db.execute( update(models.Document) .where(models.Document.id.in_(doc_ids)) @@ -673,20 +674,20 @@ async def create_observations( ) await db.commit() - except Exception as e: - # Total failure: primary write failed after retries - logger.error( - f"Failed to upsert vectors for {namespace} after retries: {e}" - ) - await db.execute( - update(models.Document) - .where(models.Document.id.in_(doc_ids)) - .values( - sync_attempts=models.Document.sync_attempts + 1, - last_sync_at=func.now(), + except Exception as e: + # Failed after retries - increment sync_attempts for reconciliation + logger.error( + f"Failed to upsert vectors for {namespace} after retries: {e}" ) - ) - await db.commit() + await db.execute( + update(models.Document) + .where(models.Document.id.in_(doc_ids)) + .values( + sync_attempts=models.Document.sync_attempts + 1, + last_sync_at=func.now(), + ) + ) + await db.commit() except IntegrityError as e: await db.rollback() @@ -770,7 +771,7 @@ async def is_rejected_duplicate( async def cleanup_soft_deleted_documents( db: AsyncSession, - vector_store: VectorStore, + external_vector_store: VectorStore, batch_size: int = 100, older_than_minutes: int = 5, ) -> int: @@ -782,7 +783,7 @@ async def cleanup_soft_deleted_documents( Args: db: Database session for executing queries - vector_store: Vector store instance for deleting vectors + external_vector_store: External vector store instance for deleting vectors batch_size: Maximum number of documents to process per call (default 100) older_than_minutes: Only process documents soft-deleted more than this many minutes ago (default 5). @@ -813,7 +814,7 @@ async def cleanup_soft_deleted_documents( # Group by namespace for batch vector deletion by_namespace: dict[str, list[str]] = {} for doc in documents: - namespace = vector_store.get_vector_namespace( + namespace = external_vector_store.get_vector_namespace( "document", doc.workspace_name, doc.observer, @@ -821,11 +822,11 @@ async def cleanup_soft_deleted_documents( ) by_namespace.setdefault(namespace, []).append(doc.id) - # Delete from vector store (per namespace) and track successful deletions + # Delete from external vector store (per namespace) and track successful deletions successfully_deleted_ids: set[str] = set() for namespace, ids in by_namespace.items(): try: - await vector_store.delete_many(namespace, ids) + await external_vector_store.delete_many(namespace, ids) # Only add to successfully_deleted_ids if vector deletion succeeded successfully_deleted_ids.update(ids) except Exception as e: diff --git a/src/crud/message.py b/src/crud/message.py index cfa9e0a6..7f7815fa 100644 --- a/src/crud/message.py +++ b/src/crud/message.py @@ -9,7 +9,7 @@ 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, upsert_with_retry +from src.vector_store import VectorRecord, get_external_vector_store, upsert_with_retry from .session import get_or_create_session @@ -137,120 +137,86 @@ async def create_messages( } embedding_dict = await embedding_client.batch_embed(id_resource_dict) - # Get vector store and namespace for this workspace's messages - vector_store = get_vector_store() - namespace = vector_store.get_vector_namespace("message", workspace_name) + external_vector_store = get_external_vector_store() + + # Determine if we need to persist embeddings to postgres + # True when: TYPE=pgvector OR still migrating (dual-write to both stores) + store_embeddings_in_postgres = ( + settings.VECTOR_STORE.TYPE == "pgvector" + or not settings.VECTOR_STORE.MIGRATED + ) # Create MessageEmbedding entries embedding_objects: list[models.MessageEmbedding] = [] - - # Check if pgvector is being used (primary or secondary) - # If so, write embeddings to ORM since pgvector relies on postgres - # Otherwise, store in memory for vector store upsert only - pgvector_in_use = ( - settings.VECTOR_STORE.PRIMARY_TYPE == "pgvector" - or settings.VECTOR_STORE.SECONDARY_TYPE == "pgvector" - ) - for message_obj in message_objects: embeddings = embedding_dict.get(message_obj.public_id, []) for chunk_idx, embedding in enumerate(embeddings): - if pgvector_in_use: - # Create MessageEmbedding record for pgvector storage - embedding_obj = models.MessageEmbedding( - content=message_obj.content, - message_id=message_obj.public_id, - workspace_name=workspace_name, - session_name=session_name, - peer_name=message_obj.peer_name, - sync_state="pending", - embedding=embedding, - ) - # Track chunk index in-memory (not persisted to avoid HNSW recomputation) - embedding_obj._chunk_index = chunk_idx - embedding_objects.append(embedding_obj) - else: - # pgvector not in use: don't create MessageEmbedding, track in-memory only - # Create a minimal object to hold metadata for vector store upsert - embedding_obj = models.MessageEmbedding( - content=message_obj.content, - message_id=message_obj.public_id, - workspace_name=workspace_name, - session_name=session_name, - peer_name=message_obj.peer_name, - sync_state="pending", - ) - embedding_obj._pending_embedding = embedding - embedding_obj._chunk_index = chunk_idx - embedding_objects.append(embedding_obj) + embedding_obj = models.MessageEmbedding( + content=message_obj.content, + message_id=message_obj.public_id, + workspace_name=workspace_name, + session_name=session_name, + peer_name=message_obj.peer_name, + sync_state="pending", + embedding=embedding if store_embeddings_in_postgres else None, + ) + embedding_obj._chunk_index = chunk_idx + embedding_obj._pending_embedding = embedding + embedding_objects.append(embedding_obj) - # Add MessageEmbedding rows to database only if pgvector in use - if embedding_objects and pgvector_in_use: + # Add MessageEmbedding rows to database only if storing in postgres + embedding_ids: list[int] = [] + if embedding_objects and store_embeddings_in_postgres: db.add_all(embedding_objects) await db.flush() - - # Track embedding IDs for sync state updates embedding_ids = [emb.id for emb in embedding_objects] - else: - embedding_ids = [] - - # Build vector records with {message_id}_{chunk_index} as vector ID - vector_records: list[VectorRecord] = [] - for emb in embedding_objects: - # Always use {message_id}_{chunk_index} as vector ID (all stores) - vector_id = f"{emb.message_id}_{emb._chunk_index}" - - # Get embedding from appropriate source - if pgvector_in_use and emb.embedding is not None: - embedding_data = [float(x) for x in emb.embedding] - elif ( - hasattr(emb, "_pending_embedding") - and emb._pending_embedding is not None - ): - embedding_data = list(emb._pending_embedding) - else: - continue - - vector_records.append( - VectorRecord( - id=vector_id, - embedding=embedding_data, - metadata={ - "message_id": emb.message_id, - "session_name": emb.session_name, - "peer_name": emb.peer_name, - }, - ) - ) await db.commit() - # Upsert to vector store with retry and update sync state - if vector_records: - try: - result = await upsert_with_retry( - vector_store, namespace, vector_records + # If no external vector store (pgvector-only mode), mark as synced immediately + if external_vector_store is None: + if embedding_ids: + await db.execute( + update(models.MessageEmbedding) + .where(models.MessageEmbedding.id.in_(embedding_ids)) + .values( + sync_state="synced", + last_sync_at=func.now(), + sync_attempts=0, + ) ) - # Only update MessageEmbedding sync state if pgvector is in use - if pgvector_in_use and embedding_ids: - if result is not None and result.secondary_ok is False: - # Partial success: primary has data but secondary doesn't - logger.warning( - "Partial sync for message embeddings: %s", - result.secondary_error, - ) - await db.execute( - update(models.MessageEmbedding) - .where(models.MessageEmbedding.id.in_(embedding_ids)) - .values( - sync_attempts=models.MessageEmbedding.sync_attempts - + 1, - last_sync_at=func.now(), - ) - ) - await db.commit() - else: - # Success: both primary and secondary stores have the data + await db.commit() + else: + # External vector store - build and upsert vector records + namespace = external_vector_store.get_vector_namespace( + "message", workspace_name + ) + + # Build vector records with {message_id}_{chunk_index} as vector ID + vector_records: list[VectorRecord] = [] + for emb in embedding_objects: + vector_id = f"{emb.message_id}_{emb._chunk_index}" + embedding_data = list(emb._pending_embedding) + vector_records.append( + VectorRecord( + id=vector_id, + embedding=embedding_data, + metadata={ + "message_id": emb.message_id, + "session_name": emb.session_name, + "peer_name": emb.peer_name, + }, + ) + ) + + # Upsert to external vector store with retry and update sync state + if vector_records: + try: + await upsert_with_retry( + external_vector_store, namespace, vector_records + ) + # Success: mark as synced if we have DB rows + if embedding_ids: await db.execute( update(models.MessageEmbedding) .where(models.MessageEmbedding.id.in_(embedding_ids)) @@ -262,20 +228,22 @@ async def create_messages( ) await db.commit() - except Exception as e: - # Total failure: primary write failed after retries - logger.error(f"Failed to upsert message vectors after retries: {e}") - # Only update MessageEmbedding sync state if pgvector is in use - if pgvector_in_use and embedding_ids: - await db.execute( - update(models.MessageEmbedding) - .where(models.MessageEmbedding.id.in_(embedding_ids)) - .values( - sync_attempts=models.MessageEmbedding.sync_attempts + 1, - last_sync_at=func.now(), - ) + except Exception as e: + # Failed after retries - increment sync_attempts for reconciliation + logger.error( + f"Failed to upsert message vectors after retries: {e}" ) - await db.commit() + if embedding_ids: + await db.execute( + update(models.MessageEmbedding) + .where(models.MessageEmbedding.id.in_(embedding_ids)) + .values( + sync_attempts=models.MessageEmbedding.sync_attempts + + 1, + last_sync_at=func.now(), + ) + ) + await db.commit() except Exception: logger.exception( diff --git a/src/crud/session.py b/src/crud/session.py index 30f5d505..62e6f6a2 100644 --- a/src/crud/session.py +++ b/src/crud/session.py @@ -18,7 +18,7 @@ from src.exceptions import ( ResourceNotFoundException, ) from src.utils.filter import apply_filter -from src.vector_store import get_vector_store +from src.vector_store import get_external_vector_store from .peer import get_or_create_peers, get_peer from .workspace import get_or_create_workspace @@ -430,9 +430,10 @@ async def delete_session( ) ) embeddings = list(embedding_result.scalars().all()) - vector_store = get_vector_store() + external_vector_store = get_external_vector_store() - if embeddings: + # Only delete from external vector store if one exists + if external_vector_store is not None and embeddings: # Compute chunk_index for each embedding based on message_id ordering message_chunks: dict[str, list[models.MessageEmbedding]] = {} for emb in embeddings: @@ -445,10 +446,12 @@ async def delete_session( for chunk_idx, chunk in enumerate(chunks): vector_ids.append(f"{chunk.message_id}_{chunk_idx}") - # Try to delete from vector store (best effort) + # Try to delete from external vector store (best effort) try: - namespace = vector_store.get_vector_namespace("message", workspace_name) - await vector_store.delete_many(namespace, vector_ids) + namespace = external_vector_store.get_vector_namespace( + "message", workspace_name + ) + await external_vector_store.delete_many(namespace, vector_ids) logger.debug( f"Deleted {len(vector_ids)} message vectors for session {session_name}" ) @@ -483,11 +486,12 @@ async def delete_session( ) documents = doc_result.all() - if documents: + # Only delete from external vector store if one exists + if external_vector_store is not None and documents: # Group document IDs by namespace (observer/observed) docs_by_namespace: dict[str, list[str]] = {} for doc in documents: - namespace = vector_store.get_vector_namespace( + namespace = external_vector_store.get_vector_namespace( "document", workspace_name, doc.observer, @@ -495,10 +499,10 @@ async def delete_session( ) docs_by_namespace.setdefault(namespace, []).append(doc.id) - # Try to delete from vector store (best effort, per namespace) + # Try to delete from external vector store (best effort, per namespace) for namespace, doc_ids in docs_by_namespace.items(): try: - await vector_store.delete_many(namespace, doc_ids) + await external_vector_store.delete_many(namespace, doc_ids) logger.debug( f"Deleted {len(doc_ids)} document vectors from {namespace}" ) diff --git a/src/crud/workspace.py b/src/crud/workspace.py index fce64804..eabdd57f 100644 --- a/src/crud/workspace.py +++ b/src/crud/workspace.py @@ -11,7 +11,7 @@ from src.cache.client import cache, get_cache_namespace from src.config import settings from src.exceptions import ConflictException, ResourceNotFoundException from src.utils.filter import apply_filter -from src.vector_store import get_vector_store +from src.vector_store import get_external_vector_store logger = getLogger(__name__) @@ -323,47 +323,50 @@ async def delete_workspace(db: AsyncSession, workspace_name: str) -> schemas.Wor await db.commit() # Delete vector store namespaces for this workspace - vector_store = get_vector_store() + external_vector_store = get_external_vector_store() # Delete message embeddings namespace for this workspace - message_namespace = vector_store.get_vector_namespace("message", workspace_name) - try: - await vector_store.delete_namespace(message_namespace) - logger.debug( - "Deleted message embeddings namespace %s for workspace %s", - message_namespace, - workspace_name, - ) - except Exception as e: - logger.warning( - "Failed to delete message embeddings namespace %s: %s", - message_namespace, - e, - ) - - # Delete document embeddings namespaces for each collection - for collection in collections: - doc_namespace = vector_store.get_vector_namespace( - "document", - workspace_name, - collection.observer, - collection.observed, + if external_vector_store: + message_namespace = external_vector_store.get_vector_namespace( + "message", workspace_name ) try: - await vector_store.delete_namespace(doc_namespace) + await external_vector_store.delete_namespace(message_namespace) logger.debug( - "Deleted document namespace %s for collection %s/%s", - doc_namespace, - collection.observer, - collection.observed, + "Deleted message embeddings namespace %s for workspace %s", + message_namespace, + workspace_name, ) except Exception as e: logger.warning( - "Failed to delete document namespace %s: %s", - doc_namespace, + "Failed to delete message embeddings namespace %s: %s", + message_namespace, e, ) + # Delete document embeddings namespaces for each collection + for collection in collections: + doc_namespace = external_vector_store.get_vector_namespace( + "document", + workspace_name, + collection.observer, + collection.observed, + ) + try: + await external_vector_store.delete_namespace(doc_namespace) + logger.debug( + "Deleted document namespace %s for collection %s/%s", + doc_namespace, + collection.observer, + collection.observed, + ) + except Exception as e: + logger.warning( + "Failed to delete document namespace %s: %s", + doc_namespace, + e, + ) + cache_key = workspace_cache_key(workspace_name) workspace_pattern = f"{cache_key}*" await cache.delete_match(workspace_pattern) diff --git a/src/deriver/vector_reconciliation.py b/src/deriver/vector_reconciliation.py index 1dea0aa1..8b3dcddf 100644 --- a/src/deriver/vector_reconciliation.py +++ b/src/deriver/vector_reconciliation.py @@ -5,12 +5,13 @@ This module provides a periodic reconciliation job that syncs documents and mess embeddings to the vector store on a rolling basis, healing any missed writes. """ +import datetime import logging import time from dataclasses import dataclass from typing import cast -from sqlalchemy import and_, select, update +from sqlalchemy import and_, delete, select, update from sqlalchemy.ext.asyncio import AsyncSession from sqlalchemy.sql.functions import func @@ -18,7 +19,7 @@ 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 +from src.vector_store import VectorRecord, VectorStore, get_external_vector_store logger = logging.getLogger(__name__) @@ -157,10 +158,10 @@ async def _bump_message_embedding_sync_attempts( async def _sync_documents( db: AsyncSession, documents: list[models.Document], - vector_store: VectorStore, + external_vector_store: VectorStore, ) -> tuple[int, int]: """ - Sync a batch of documents to the vector store. + Sync a batch of documents to the external vector store. Returns (synced_count, failed_count). """ @@ -170,9 +171,10 @@ 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" + # Determine if we need to persist embeddings to postgres + # True when: TYPE=pgvector OR still migrating (dual-write to both stores) + store_embeddings_in_postgres = ( + settings.VECTOR_STORE.TYPE == "pgvector" or not settings.VECTOR_STORE.MIGRATED ) missing_docs: list[models.Document] = [] @@ -197,8 +199,8 @@ async def _sync_documents( for doc, embedding in zip(missing_docs, embeddings, strict=False): reembedded_by_id[doc.id] = embedding - # Write re-embedded vectors to postgres if pgvector is in use - if pgvector_in_use and reembedded_by_id: + # Write re-embedded vectors to postgres if needed + if store_embeddings_in_postgres and reembedded_by_id: for doc_id, embedding in reembedded_by_id.items(): await db.execute( update(models.Document) @@ -226,7 +228,7 @@ async def _sync_documents( # Group documents by namespace (workspace/observer/observed) by_namespace: dict[str, list[models.Document]] = {} for doc in documents: - namespace = vector_store.get_vector_namespace( + namespace = external_vector_store.get_vector_namespace( "document", doc.workspace_name, doc.observer, doc.observed ) by_namespace.setdefault(namespace, []).append(doc) @@ -267,20 +269,8 @@ async def _sync_documents( doc_ids = [doc.id for doc in docs_with_vectors] - result = None if vector_records: - result = await vector_store.upsert_many(namespace, vector_records) - - if result is not None and result.secondary_ok is False: - logger.warning( - "Partial sync for namespace %s: %s", - namespace, - result.secondary_error, - ) - # Increment attempts and mark as failed if we've hit max attempts - await _bump_document_sync_attempts(db, docs_with_vectors) - failed_count += len(docs_with_vectors) - continue + await external_vector_store.upsert_many(namespace, vector_records) # Mark as synced await db.execute( @@ -295,7 +285,9 @@ async def _sync_documents( synced_count += len(docs_with_vectors) except Exception as e: - logger.warning(f"Failed to sync documents to {namespace}: {e}") + logger.warning( + f"Failed to sync documents to external vector store {namespace}: {e}" + ) # Increment attempts and mark as failed if we've hit max attempts await _bump_document_sync_attempts(db, docs_with_vectors) failed_count += len(docs_with_vectors) @@ -306,10 +298,10 @@ async def _sync_documents( async def _sync_message_embeddings( db: AsyncSession, embeddings: list[models.MessageEmbedding], - vector_store: VectorStore, + external_vector_store: VectorStore, ) -> tuple[int, int]: """ - Sync a batch of message embeddings to the vector store. + Sync a batch of message embeddings to the external vector store. Returns (synced_count, failed_count). """ @@ -319,9 +311,10 @@ async def _sync_message_embeddings( synced_count = 0 failed_count = 0 - pgvector_in_use = ( - settings.VECTOR_STORE.PRIMARY_TYPE == "pgvector" - or settings.VECTOR_STORE.SECONDARY_TYPE == "pgvector" + # Determine if we need to persist embeddings to postgres + # True when: TYPE=pgvector OR still migrating (dual-write to both stores) + store_embeddings_in_postgres = ( + settings.VECTOR_STORE.TYPE == "pgvector" or not settings.VECTOR_STORE.MIGRATED ) # Re-embed missing payloads so reconciliation can heal non-pgvector writes @@ -334,8 +327,8 @@ async def _sync_message_embeddings( new_embedding = await embedding_client.embed(emb.content) reembedded_by_id[emb.id] = new_embedding - # Only persist embeddings to postgres when pgvector is in play - if pgvector_in_use: + # Only persist embeddings to postgres when needed + if store_embeddings_in_postgres: await db.execute( update(models.MessageEmbedding) .where(models.MessageEmbedding.id == emb.id) @@ -359,7 +352,9 @@ async def _sync_message_embeddings( # Group by namespace (workspace) by_namespace: dict[str, list[models.MessageEmbedding]] = {} for emb in embeddings: - namespace = vector_store.get_vector_namespace("message", emb.workspace_name) + namespace = external_vector_store.get_vector_namespace( + "message", emb.workspace_name + ) by_namespace.setdefault(namespace, []).append(emb) # Compute chunk_index for each embedding based on message_id ordering @@ -410,20 +405,8 @@ async def _sync_message_embeddings( emb_ids = [emb.id for emb in embs_with_vectors] - result = None if vector_records: - result = await vector_store.upsert_many(namespace, vector_records) - - if result is not None and result.secondary_ok is False: - logger.warning( - "Partial sync for namespace %s: %s", - namespace, - result.secondary_error, - ) - # Increment attempts and mark as failed if we've hit max attempts - await _bump_message_embedding_sync_attempts(db, embs_with_vectors) - failed_count += len(embs_with_vectors) - continue + await external_vector_store.upsert_many(namespace, vector_records) # Mark as synced await db.execute( @@ -438,7 +421,9 @@ async def _sync_message_embeddings( synced_count += len(embs_with_vectors) except Exception as e: - logger.warning(f"Failed to sync message embeddings to {namespace}: {e}") + logger.warning( + f"Failed to sync message embeddings to external vector store {namespace}: {e}" + ) # Increment attempts and mark as failed if we've hit max attempts await _bump_message_embedding_sync_attempts(db, embs_with_vectors) failed_count += len(embs_with_vectors) @@ -446,6 +431,39 @@ async def _sync_message_embeddings( return synced_count, failed_count +async def _cleanup_soft_deleted_documents_pgvector( + db: AsyncSession, + batch_size: int = RECONCILIATION_BATCH_SIZE, + older_than_minutes: int = 5, +) -> int: + """ + Cleanup soft-deleted documents + """ + + cutoff = datetime.datetime.now(datetime.timezone.utc) - datetime.timedelta( + minutes=older_than_minutes + ) + + # Find soft-deleted documents ready for cleanup + stmt = ( + select(models.Document.id) + .where(models.Document.deleted_at.is_not(None)) + .where(models.Document.deleted_at < cutoff) + .limit(batch_size) + .with_for_update(skip_locked=True) + ) + result = await db.execute(stmt) + doc_ids = [row[0] for row in result.all()] + + if not doc_ids: + return 0 + + # Hard delete directly (no vector store cleanup needed in pgvector mode) + await db.execute(delete(models.Document).where(models.Document.id.in_(doc_ids))) + logger.debug(f"Cleaned up {len(doc_ids)} soft-deleted documents (pgvector mode)") + return len(doc_ids) + + async def run_vector_reconciliation_cycle() -> ReconciliationMetrics: """ Run a complete reconciliation cycle. @@ -456,20 +474,33 @@ async def run_vector_reconciliation_cycle() -> ReconciliationMetrics: Returns metrics about what was synced. """ metrics = ReconciliationMetrics() - vector_store = get_vector_store() + external_vector_store = get_external_vector_store() deadline = time.monotonic() + RECONCILIATION_TIME_BUDGET_SECONDS from src.crud.document import cleanup_soft_deleted_documents print("Running vector reconciliation cycle") async with tracked_db("reconciliation") as db: + # If no external vector store (pgvector mode), only clean up soft-deleted documents + if external_vector_store is None: + while time.monotonic() < deadline: + cleaned = await _cleanup_soft_deleted_documents_pgvector( + db, batch_size=RECONCILIATION_BATCH_SIZE + ) + if cleaned: + metrics.documents_cleaned += cleaned + await db.commit() + else: + break + return metrics + while time.monotonic() < deadline: did_work = False # Reconcile documents docs = await _get_documents_needing_sync(db) if docs: - synced, failed = await _sync_documents(db, docs, vector_store) + synced, failed = await _sync_documents(db, docs, external_vector_store) metrics.documents_synced += synced metrics.documents_failed += failed await db.commit() @@ -483,7 +514,7 @@ async def run_vector_reconciliation_cycle() -> ReconciliationMetrics: if embs: try: synced, failed = await _sync_message_embeddings( - db, embs, vector_store + db, embs, external_vector_store ) except Exception as e: logger.warning( @@ -505,7 +536,7 @@ async def run_vector_reconciliation_cycle() -> ReconciliationMetrics: # Clean up soft-deleted documents cleaned = await cleanup_soft_deleted_documents( db, - vector_store, + external_vector_store, batch_size=RECONCILIATION_BATCH_SIZE, ) if cleaned: diff --git a/src/main.py b/src/main.py index e83d538b..6b69824b 100644 --- a/src/main.py +++ b/src/main.py @@ -126,9 +126,9 @@ async def lifespan(_: FastAPI): yield finally: # Import here to avoid circular import at module load time - from src.vector_store import close_vector_store + from src.vector_store import close_external_vector_store - await close_vector_store() + await close_external_vector_store() await close_cache() await engine.dispose() diff --git a/src/utils/search.py b/src/utils/search.py index 1d17ed4a..68717e43 100644 --- a/src/utils/search.py +++ b/src/utils/search.py @@ -17,7 +17,7 @@ from src.embedding_client import embedding_client from src.exceptions import ValidationException from src.models import session_peers_table from src.utils.filter import apply_filter -from src.vector_store import get_vector_store +from src.vector_store import get_external_vector_store T = TypeVar("T") @@ -91,9 +91,10 @@ async def _semantic_search( f"Query exceeds maximum token limit of {settings.MAX_EMBEDDING_TOKENS}." ) from e - # If pgvector is primary, query Postgres directly with similarity + filters - # This avoids duplicate fetches from the same database - if settings.EMBED_MESSAGES and settings.VECTOR_STORE.PRIMARY_TYPE == "pgvector": + # Query Postgres / pgvector directly + if settings.EMBED_MESSAGES and ( + settings.VECTOR_STORE.TYPE == "pgvector" or not settings.VECTOR_STORE.MIGRATED + ): # Join message_embeddings with messages to get full message objects distance_expr = models.MessageEmbedding.embedding.cosine_distance( embedding_query @@ -124,9 +125,12 @@ async def _semantic_search( result = await db.execute(stmt) return list(result.scalars().all()) - # FALLBACK: Use vector store abstraction for external stores (Turbopuffer, LanceDB) - vector_store = get_vector_store() - namespace = vector_store.get_vector_namespace("message", workspace_name) + # FALLBACK: Use external vector store (Turbopuffer, LanceDB) + external_vector_store = get_external_vector_store() + if external_vector_store is None: + return [] + + namespace = external_vector_store.get_vector_namespace("message", workspace_name) # Build vector store filters from the provided filters vector_filters: dict[str, Any] = {} @@ -137,9 +141,9 @@ async def _semantic_search( if "peer_id" in filters: vector_filters["peer_name"] = filters["peer_id"] - # Query vector store for similar message embeddings + # Query external vector store for similar message embeddings # Since all filters are applied at the vector store level, we don't need to oversample - vector_results = await vector_store.query( + vector_results = await external_vector_store.query( namespace, embedding_query, top_k=limit, diff --git a/src/vector_store/__init__.py b/src/vector_store/__init__.py index 3b1c95ea..3d19b6e8 100644 --- a/src/vector_store/__init__.py +++ b/src/vector_store/__init__.py @@ -37,12 +37,9 @@ class VectorUpsertResult(BaseModel): model_config: ClassVar[ConfigDict] = ConfigDict( extra="forbid", frozen=True, - arbitrary_types_allowed=True, ) - primary_ok: bool - secondary_ok: bool | None = None - secondary_error: Exception | None = None + ok: bool class VectorStore(ABC): @@ -173,9 +170,7 @@ class VectorStore(ABC): # Import implementations after base classes are defined to avoid circular imports -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 @@ -186,68 +181,46 @@ def _create_store_by_type(store_type: str) -> VectorStore: return TurbopufferVectorStore() elif store_type == "lancedb": return LanceDBVectorStore() - elif store_type == "pgvector": - return PgVectorStore() else: raise ValueError(f"Unknown vector store type: {store_type}") -def _create_vector_store() -> VectorStore: - """ - Create a new vector store instance based on configuration. - - If SECONDARY_TYPE is set, returns a CompositeVectorStore that: - - Writes to both primary and secondary stores - - Reads from primary only, falls back to secondary on failure - - Returns: - The vector store instance based on configuration. - - Raises: - ValueError: If the configured vector store type is invalid. - """ - primary = _create_store_by_type(settings.VECTOR_STORE.PRIMARY_TYPE) - - if settings.VECTOR_STORE.SECONDARY_TYPE: - secondary = _create_store_by_type(settings.VECTOR_STORE.SECONDARY_TYPE) - return CompositeVectorStore(primary=primary, secondary=secondary) - - return primary - - @cache -def get_vector_store() -> VectorStore: +def get_external_vector_store() -> VectorStore | None: """ - Get the configured vector store instance (singleton). + Get the configured external vector store instance (singleton). - Uses functools.cache to ensure only one instance is created per process. - This is asyncio-safe since there are no await points in the creation path. + Returns None if TYPE='pgvector' since pgvector operations happen via ORM directly. + External vector stores include Turbopuffer and LanceDB. Returns: - The vector store instance based on configuration. + The external vector store instance, or None if using pgvector (ORM handles it). Raises: ValueError: If the configured vector store type is invalid. """ - return _create_vector_store() + if settings.VECTOR_STORE.TYPE == "pgvector": + return None + return _create_store_by_type(settings.VECTOR_STORE.TYPE) -async def close_vector_store() -> None: +async def close_external_vector_store() -> None: """ - Close the vector store and release resources. + Close the external vector store and release resources. Call this during application shutdown to cleanly close connections. - After calling this, you must call get_vector_store.cache_clear() if you + After calling this, you must call get_external_vector_store.cache_clear() if you want to create a new instance. """ # Check if an instance was ever created if ( - get_vector_store.cache_info().hits > 0 - or get_vector_store.cache_info().misses > 0 + get_external_vector_store.cache_info().hits > 0 + or get_external_vector_store.cache_info().misses > 0 ): - store = get_vector_store() - await store.close() - get_vector_store.cache_clear() + store = get_external_vector_store() + if store is not None: + await store.close() + get_external_vector_store.cache_clear() __all__ = [ @@ -255,7 +228,7 @@ __all__ = [ "VectorRecord", "VectorQueryResult", "VectorUpsertResult", - "get_vector_store", - "close_vector_store", + "get_external_vector_store", + "close_external_vector_store", "upsert_with_retry", ] diff --git a/src/vector_store/composite.py b/src/vector_store/composite.py deleted file mode 100644 index 46638d74..00000000 --- a/src/vector_store/composite.py +++ /dev/null @@ -1,261 +0,0 @@ -""" -Composite vector store implementation. - -This module provides a composite VectorStore that writes to two stores -and reads from primary with fallback to secondary. -""" - -import asyncio -import logging -from typing import Any - -from . import VectorQueryResult, VectorRecord, VectorStore, VectorUpsertResult - -logger = logging.getLogger(__name__) - - -class CompositeVectorStore(VectorStore): - """ - Composite vector store with dual-write and fallback-read. - - Behavior: - - Writes go to BOTH primary and secondary stores - - Reads try primary only, fall back to secondary on failure - - Migration Strategy: - - Primary = source of truth (pgvector) - - Secondary = target being populated (turbopuffer) - - Reconciliation job syncs data from primary to secondary via dual-writes - - Reads use primary, fallback to secondary only on exception (network errors, etc.) - - After migration completes, remove secondary from config - """ - - primary: VectorStore - secondary: VectorStore - - def __init__(self, primary: VectorStore, secondary: VectorStore): - """ - Initialize the composite vector store. - - Args: - primary: The primary vector store (reads prefer this) - secondary: The secondary vector store (fallback for reads) - """ - super().__init__() - self.primary = primary - self.secondary = secondary - - async def upsert_many( - self, - namespace: str, - vectors: list[VectorRecord], - ) -> VectorUpsertResult: - """ - Upsert multiple vectors to both stores. - - Success cases: - - Primary ✓, Secondary ✓ → Success (fully synced) - - Primary ✓, Secondary ✗ → Returns partial result (not fully synced) - - Failure cases: - - Primary ✗, Secondary ✓ → Raises primary exception (weird, shouldn't happen) - - Primary ✗, Secondary ✗ → Raises primary exception (total failure) - - Args: - namespace: The namespace to store the vectors in - vectors: List of VectorRecord objects to upsert - - Returns: - Result describing primary/secondary outcomes. - - Raises: - Exception: Primary failed (secondary state doesn't matter) - """ - if not vectors: - return VectorUpsertResult(primary_ok=True, secondary_ok=True) - - # Write to both stores concurrently - primary_task = asyncio.create_task(self.primary.upsert_many(namespace, vectors)) - secondary_task = asyncio.create_task( - self.secondary.upsert_many(namespace, vectors) - ) - - # Wait for both, gathering exceptions - results: tuple[ - VectorUpsertResult | BaseException, VectorUpsertResult | BaseException - ] = await asyncio.gather(primary_task, secondary_task, return_exceptions=True) - - primary_result, secondary_result = results - - # Case 1: Both failed → raise primary exception - if isinstance(primary_result, Exception) and isinstance( - secondary_result, Exception - ): - logger.error( - f"Both primary and secondary upsert failed for namespace {namespace}. Primary: {primary_result}, Secondary: {secondary_result}" - ) - raise primary_result - - # Case 2: Primary failed, secondary succeeded → raise primary exception (weird case) - if isinstance(primary_result, Exception): - logger.error( - f"Primary upsert failed but secondary succeeded for namespace {namespace}: {primary_result}" - ) - raise primary_result - - # Case 3: Primary succeeded, secondary failed → return partial result - if isinstance(secondary_result, Exception): - logger.warning( - f"Primary upsert succeeded but secondary failed for namespace {namespace}: {secondary_result}" - ) - return VectorUpsertResult( - primary_ok=True, - secondary_ok=False, - secondary_error=secondary_result, - ) - - # Case 4: Both succeeded → log success - logger.debug( - f"Dual-write upserted {len(vectors)} vectors to namespace {namespace}" - ) - return VectorUpsertResult(primary_ok=True, secondary_ok=True) - - async def query( - self, - namespace: str, - embedding: list[float], - *, - top_k: int = 10, - filters: dict[str, Any] | None = None, - max_distance: float | None = None, - ) -> list[VectorQueryResult]: - """ - Query for similar vectors, trying primary first then falling back to secondary on failure. - - Primary is the source of truth. Secondary is only used if primary query raises an - exception (network errors, timeouts, etc.). If primary returns empty results ([]), - that's considered success and we don't query secondary. - - Args: - namespace: The namespace to query - embedding: The query embedding vector - top_k: Maximum number of results to return - filters: Optional metadata filters - max_distance: Optional maximum distance threshold (cosine distance) - - Returns: - List of VectorQueryResult objects, ordered by similarity (most similar first) - - Raises: - Exception: If both primary and secondary queries fail - """ - try: - results = await self.primary.query( - namespace, - embedding, - top_k=top_k, - filters=filters, - max_distance=max_distance, - ) - logger.debug( - f"Primary query returned {len(results)} results from namespace {namespace}" - ) - return results - except Exception as primary_error: - logger.warning( - f"Primary query failed for namespace {namespace}: {primary_error}, attempting fallback to secondary" - ) - try: - results = await self.secondary.query( - namespace, - embedding, - top_k=top_k, - filters=filters, - max_distance=max_distance, - ) - logger.warning( - f"Secondary query returned {len(results)} results from namespace {namespace}" - ) - return results - except Exception as secondary_error: - logger.error( - f"Both primary and secondary queries failed for namespace {namespace}. Primary: {primary_error}, Secondary: {secondary_error}" - ) - raise primary_error from secondary_error - - async def delete_many(self, namespace: str, ids: list[str]) -> None: - """ - Delete vectors from both stores. - - Args: - namespace: The namespace containing the vectors - ids: List of vector identifiers to delete - """ - if not ids: - return - - # Delete from both stores concurrently - primary_task = asyncio.create_task(self.primary.delete_many(namespace, ids)) - secondary_task = asyncio.create_task(self.secondary.delete_many(namespace, ids)) - - results: tuple[ - None | BaseException, None | BaseException - ] = await asyncio.gather(primary_task, secondary_task, return_exceptions=True) - - primary_result, secondary_result = results - - # Primary failure is critical - raise it - if isinstance(primary_result, Exception): - logger.error( - f"Primary vector store delete failed for namespace {namespace}: {primary_result}" - ) - raise primary_result - - if isinstance(secondary_result, Exception): - logger.warning( - f"Secondary vector store delete failed for namespace {namespace}: {secondary_result}" - ) - - logger.debug( - f"Dual-delete removed {len(ids)} vectors from namespace {namespace}" - ) - - async def delete_namespace(self, namespace: str) -> None: - """ - Delete an entire namespace from both stores. - - Args: - namespace: The namespace to delete - """ - # Delete from both stores concurrently - primary_task = asyncio.create_task(self.primary.delete_namespace(namespace)) - secondary_task = asyncio.create_task(self.secondary.delete_namespace(namespace)) - - results: tuple[ - None | BaseException, None | BaseException - ] = await asyncio.gather(primary_task, secondary_task, return_exceptions=True) - - primary_result, secondary_result = results - - # Primary failure is critical - raise it - if isinstance(primary_result, Exception): - logger.error( - f"Primary vector store namespace delete failed for {namespace}: {primary_result}" - ) - raise primary_result - - if isinstance(secondary_result, Exception): - logger.warning( - f"Secondary vector store namespace delete failed for {namespace}: {secondary_result}" - ) - - logger.debug(f"Dual-delete removed namespace {namespace}") - - async def close(self) -> None: - """Close both primary and secondary vector stores.""" - await asyncio.gather( - self.primary.close(), - self.secondary.close(), - return_exceptions=True, - ) - logger.debug("Composite vector store closed") diff --git a/src/vector_store/lancedb.py b/src/vector_store/lancedb.py index 2ea1cdff..fe4d3a58 100644 --- a/src/vector_store/lancedb.py +++ b/src/vector_store/lancedb.py @@ -150,7 +150,7 @@ class LanceDBVectorStore(VectorStore): vectors: List of VectorRecord objects to upsert """ if not vectors: - return VectorUpsertResult(primary_ok=True) + return VectorUpsertResult(ok=True) try: rows = [self._row_to_dict(v) for v in vectors] @@ -165,7 +165,7 @@ class LanceDBVectorStore(VectorStore): ) logger.debug(f"Upserted {len(vectors)} vectors to namespace {namespace}") - return VectorUpsertResult(primary_ok=True) + return VectorUpsertResult(ok=True) except Exception: logger.exception( f"Failed to upsert {len(vectors)} vectors to namespace {namespace}" diff --git a/src/vector_store/pgvector.py b/src/vector_store/pgvector.py deleted file mode 100644 index 7b113f33..00000000 --- a/src/vector_store/pgvector.py +++ /dev/null @@ -1,315 +0,0 @@ -""" -PostgreSQL pgvector vector store implementation. - -This module provides a pgvector-based implementation of the VectorStore interface -using the existing embedding columns on documents and message_embeddings tables. -""" - -import logging -from typing import Any - -from sqlalchemy import select -from sqlalchemy.ext.asyncio import AsyncSession - -from src import models -from src.dependencies import tracked_db - -from . import VectorQueryResult, VectorRecord, VectorStore, VectorUpsertResult - -logger = logging.getLogger(__name__) - - -class PgVectorStore(VectorStore): - """ - PostgreSQL pgvector implementation of the VectorStore interface. - - Uses the existing embedding columns on documents and message_embeddings tables, - providing transactional consistency with document/message metadata. - - Namespace mapping: - - {prefix}.{workspace}.{observer}.{observed} -> documents.embedding - - {prefix}.{workspace}.messages -> message_embeddings.embedding - """ - - def __init__(self): - """Initialize the pgvector store.""" - super().__init__() - - def _parse_namespace(self, namespace: str) -> tuple[str, dict[str, str]]: - """ - Parse a namespace string to determine the table and filter context. - - Args: - namespace: Namespace string like "{prefix}.{workspace}.{observer}.{observed}" - or "{prefix}.{workspace}.messages" - - Returns: - Tuple of (table_type, context_dict) where: - - table_type is "documents" or "message_embeddings" - - context_dict contains workspace_name and optionally observer/observed - """ - parts = namespace.split(".") - - # Expected formats: - # {prefix}.{workspace}.messages -> message_embeddings - # {prefix}.{workspace}.{observer}.{observed} -> documents - if len(parts) < 3: - raise ValueError(f"Invalid namespace format: {namespace}") - - workspace = parts[1] - - if parts[2] == "messages": - return "message_embeddings", {"workspace_name": workspace} - elif len(parts) >= 4: - observer = parts[2] - observed = parts[3] - return "documents", { - "workspace_name": workspace, - "observer": observer, - "observed": observed, - } - else: - raise ValueError(f"Invalid namespace format: {namespace}") - - async def upsert_many( - self, - namespace: str, - vectors: list[VectorRecord], - ) -> VectorUpsertResult: - """ - Upsert multiple vectors into the database. - - NOTE: This is a no-op. When pgvector is being used (as primary or secondary), - embeddings are written directly to postgres via the ORM (in message.py/document.py). - This method exists only to satisfy the VectorStore interface. - - The vector store abstraction is used for: - - Queries (which still go through pgvector) - - Writing to secondary stores (e.g., turbopuffer during migration) - - Args: - namespace: The namespace (determines table) - vectors: List of VectorRecord objects to upsert - """ - # No-op: embeddings are already written to postgres via ORM - if vectors: - logger.debug( - f"PgVectorStore.upsert_many() no-op for {len(vectors)} vectors in {namespace} (embeddings written via ORM)" - ) - return VectorUpsertResult(primary_ok=True) - - async def query( - self, - namespace: str, - embedding: list[float], - *, - top_k: int = 10, - filters: dict[str, Any] | None = None, - max_distance: float | None = None, - ) -> list[VectorQueryResult]: - """ - Query for similar vectors using pgvector cosine distance. - - Args: - namespace: The namespace to query - embedding: The query embedding vector - top_k: Maximum number of results to return - filters: Optional metadata filters - max_distance: Optional maximum distance threshold (cosine distance) - - Returns: - List of VectorQueryResult objects, ordered by similarity (most similar first) - """ - table_type, context = self._parse_namespace(namespace) - - 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 - - except Exception: - logger.exception(f"Failed to query namespace {namespace}") - raise - - async def _query_documents( - self, - db: AsyncSession, - context: dict[str, str], - embedding: list[float], - top_k: int, - filters: dict[str, Any] | None, - max_distance: float | None, - ) -> list[VectorQueryResult]: - """Query documents table for similar vectors.""" - # Build the query with cosine distance - # pgvector uses <=> for cosine distance - stmt = ( - select( - models.Document.id, - models.Document.embedding.cosine_distance(embedding).label("distance"), - models.Document.workspace_name, - models.Document.observer, - models.Document.observed, - models.Document.session_name, - models.Document.level, - ) - .where(models.Document.embedding.isnot(None)) - .where(models.Document.deleted_at.is_(None)) - .where(models.Document.workspace_name == context["workspace_name"]) - .where(models.Document.observer == context["observer"]) - .where(models.Document.observed == context["observed"]) - ) - - # Apply additional filters - if filters: - if "session_name" in filters: - stmt = stmt.where( - models.Document.session_name == filters["session_name"] - ) - if "level" in filters: - stmt = stmt.where(models.Document.level == filters["level"]) - - # Apply max_distance filter - if max_distance is not None: - stmt = stmt.where( - models.Document.embedding.cosine_distance(embedding) <= max_distance - ) - - # Order by distance and limit - stmt = stmt.order_by("distance").limit(top_k) - - result = await db.execute(stmt) - rows = result.all() - - return [ - VectorQueryResult( - id=str(row.id), - score=float(row.distance), - metadata={ - "workspace_name": row.workspace_name, - "observer": row.observer, - "observed": row.observed, - "session_name": row.session_name, - "level": row.level, - }, - ) - for row in rows - ] - - async def _query_message_embeddings( - self, - db: AsyncSession, - context: dict[str, str], - embedding: list[float], - top_k: int, - filters: dict[str, Any] | None, - max_distance: float | None, - ) -> list[VectorQueryResult]: - """Query message_embeddings table for similar vectors.""" - # Build the query with cosine distance - stmt = ( - select( - models.MessageEmbedding.id, - models.MessageEmbedding.embedding.cosine_distance(embedding).label( - "distance" - ), - models.MessageEmbedding.message_id, - models.MessageEmbedding.workspace_name, - models.MessageEmbedding.session_name, - models.MessageEmbedding.peer_name, - ) - .where(models.MessageEmbedding.embedding.isnot(None)) - .where(models.MessageEmbedding.workspace_name == context["workspace_name"]) - ) - - # Apply additional filters - if filters: - if "session_name" in filters: - stmt = stmt.where( - models.MessageEmbedding.session_name == filters["session_name"] - ) - if "peer_name" in filters: - stmt = stmt.where( - models.MessageEmbedding.peer_name == filters["peer_name"] - ) - if "message_id" in filters: - stmt = stmt.where( - models.MessageEmbedding.message_id == filters["message_id"] - ) - - # Apply max_distance filter - if max_distance is not None: - stmt = stmt.where( - models.MessageEmbedding.embedding.cosine_distance(embedding) - <= max_distance - ) - - # Order by distance and limit - stmt = stmt.order_by("distance").limit(top_k) - - result = await db.execute(stmt) - rows = result.all() - - return [ - VectorQueryResult( - id=str(row.id), - score=float(row.distance), - metadata={ - "embedding_id": row.id, - "message_id": row.message_id, - "workspace_name": row.workspace_name, - "session_name": row.session_name, - "peer_name": row.peer_name, - }, - ) - for row in rows - ] - - async def delete_many(self, namespace: str, ids: list[str]) -> None: - """ - No-op for pgvector. Vector deletion is handled by row deletion. - - For pgvector, vectors are stored in the same postgres rows as documents/ - message_embeddings. The reconciliation job handles hard-deleting rows - after calling this method, which removes both the data and the embedding. - - Args: - namespace: The namespace containing the vectors - ids: List of vector identifiers to delete - """ - if ids: - logger.debug( - f"PgVectorStore.delete_many() no-op for {len(ids)} vectors in {namespace} (row deletion handles embedding removal)" - ) - - async def delete_namespace(self, namespace: str) -> None: - """ - No-op for pgvector. Namespace deletion is handled by row deletion. - - For pgvector, vectors are stored in the same postgres rows as documents/ - message_embeddings. Deleting a collection or workspace should delete - the rows directly via ORM/SQL, which removes both data and embeddings. - - Args: - namespace: The namespace to delete - """ - logger.debug( - f"PgVectorStore.delete_namespace() no-op for {namespace} (row deletion handles embedding removal)" - ) - - async def close(self) -> None: - """Close the pgvector store (no-op for pgvector)""" - pass diff --git a/src/vector_store/turbopuffer.py b/src/vector_store/turbopuffer.py index fc4e26be..ba5c3612 100644 --- a/src/vector_store/turbopuffer.py +++ b/src/vector_store/turbopuffer.py @@ -71,7 +71,7 @@ class TurbopufferVectorStore(VectorStore): vectors: List of VectorRecord objects to upsert """ if not vectors: - return VectorUpsertResult(primary_ok=True) + return VectorUpsertResult(ok=True) ns = self._get_namespace(namespace) @@ -89,7 +89,7 @@ class TurbopufferVectorStore(VectorStore): upsert_rows=rows, distance_metric=DISTANCE_METRIC, ) - return VectorUpsertResult(primary_ok=True) + return VectorUpsertResult(ok=True) except Exception: logger.exception( f"Failed to upsert {len(vectors)} vectors to namespace {namespace}" diff --git a/src/vector_store/utils.py b/src/vector_store/utils.py index 8462f4ff..ae613ada 100644 --- a/src/vector_store/utils.py +++ b/src/vector_store/utils.py @@ -10,7 +10,6 @@ from typing import TYPE_CHECKING from tenacity import ( AsyncRetrying, retry_if_exception_type, - retry_if_result, stop_after_attempt, wait_exponential, ) @@ -30,8 +29,6 @@ async def upsert_with_retry( """ 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 @@ -51,8 +48,7 @@ async def upsert_with_retry( 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), + retry=retry_if_exception_type(Exception), reraise=True, ): with attempt: diff --git a/tests/conftest.py b/tests/conftest.py index a531d867..b042a7b6 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -393,7 +393,7 @@ def mock_vector_store(): vector_storage[namespace] = {} for vector in vectors: vector_storage[namespace][vector.id] = (vector.embedding, vector.metadata) - return VectorUpsertResult(primary_ok=True) + return VectorUpsertResult(ok=True) async def mock_query( namespace: str, embedding: list[float], **kwargs: Any @@ -423,10 +423,10 @@ def mock_vector_store(): async def mock_delete_namespace(namespace: str) -> None: vector_storage.pop(namespace, None) - # Clear the cache on get_vector_store before patching - from src.vector_store import get_vector_store + # Clear the cache on get_external_vector_store before patching + from src.vector_store import get_external_vector_store - get_vector_store.cache_clear() # type: ignore + get_external_vector_store.cache_clear() # type: ignore # Create the mock vector store mock_vs = MagicMock() @@ -454,19 +454,20 @@ def mock_vector_store(): mock_vs.get_vector_namespace = mock_get_vector_namespace with ( - patch("src.crud.document.get_vector_store", return_value=mock_vs), - patch("src.crud.workspace.get_vector_store", return_value=mock_vs), - patch("src.crud.session.get_vector_store", return_value=mock_vs), - patch("src.crud.message.get_vector_store", return_value=mock_vs), + patch("src.crud.document.get_external_vector_store", return_value=mock_vs), + patch("src.crud.workspace.get_external_vector_store", return_value=mock_vs), + patch("src.crud.session.get_external_vector_store", return_value=mock_vs), + patch("src.crud.message.get_external_vector_store", return_value=mock_vs), patch( - "src.deriver.vector_reconciliation.get_vector_store", return_value=mock_vs + "src.deriver.vector_reconciliation.get_external_vector_store", + return_value=mock_vs, ), - patch("src.utils.search.get_vector_store", return_value=mock_vs), + patch("src.utils.search.get_external_vector_store", return_value=mock_vs), ): yield mock_vs # Clear cache after test as well for cleanliness - get_vector_store.cache_clear() # type: ignore + get_external_vector_store.cache_clear() # type: ignore @pytest.fixture(autouse=True) diff --git a/tests/deriver/test_vector_reconciliation.py b/tests/deriver/test_vector_reconciliation.py index 901e5b9a..728f99d5 100644 --- a/tests/deriver/test_vector_reconciliation.py +++ b/tests/deriver/test_vector_reconciliation.py @@ -80,7 +80,7 @@ class TestStateTransitions: return_value=f"honcho.{workspace.name}.{peer1.name}.{peer1.name}" ) mock_vector_store.upsert_many = AsyncMock( - return_value=VectorUpsertResult(primary_ok=True, secondary_ok=True) + return_value=VectorUpsertResult(ok=True) ) # Run sync @@ -136,23 +136,19 @@ class TestStateTransitions: await db_session.commit() await db_session.refresh(doc) - # Mock vector store to have partial failure (secondary fails) + # Mock vector store to fail with exception mock_vector_store = MagicMock(spec=VectorStore) mock_vector_store.get_vector_namespace = MagicMock( return_value=f"honcho.{workspace.name}.{peer1.name}.{peer1.name}" ) mock_vector_store.upsert_many = AsyncMock( - return_value=VectorUpsertResult( - primary_ok=True, - secondary_ok=False, - secondary_error=Exception("Secondary failed"), - ) + side_effect=Exception("Vector store failed") ) # Run sync synced, failed = await _sync_documents(db_session, [doc], mock_vector_store) - # Verify partial failure recorded + # Verify failure recorded assert synced == 0 assert failed == 1 @@ -201,15 +197,13 @@ class TestStateTransitions: await db_session.commit() await db_session.refresh(doc) - # Mock vector store to fail + # Mock vector store to fail with exception mock_vector_store = MagicMock(spec=VectorStore) mock_vector_store.get_vector_namespace = MagicMock( return_value=f"honcho.{workspace.name}.{peer1.name}.{peer1.name}" ) mock_vector_store.upsert_many = AsyncMock( - return_value=VectorUpsertResult( - primary_ok=True, secondary_ok=False, secondary_error=Exception("Failed") - ) + side_effect=Exception("Vector store failed") ) # Run sync - this should be the final attempt @@ -316,7 +310,7 @@ class TestBatchProcessing: if namespace not in namespace_calls: namespace_calls[namespace] = [] namespace_calls[namespace].extend(vectors) - return VectorUpsertResult(primary_ok=True, secondary_ok=True) + return VectorUpsertResult(ok=True) mock_vector_store.get_vector_namespace = mock_get_namespace mock_vector_store.upsert_many = mock_upsert @@ -444,7 +438,7 @@ class TestReEmbedding: return_value=f"honcho.{workspace.name}.{peer1.name}.{peer1.name}" ) mock_vector_store.upsert_many = AsyncMock( - return_value=VectorUpsertResult(primary_ok=True, secondary_ok=True) + return_value=VectorUpsertResult(ok=True) ) # Run sync @@ -518,7 +512,7 @@ class TestReEmbedding: return_value=f"honcho.{workspace.name}.{peer1.name}.{peer1.name}" ) mock_vector_store.upsert_many = AsyncMock( - return_value=VectorUpsertResult(primary_ok=True, secondary_ok=True) + return_value=VectorUpsertResult(ok=True) ) # Run sync @@ -792,7 +786,7 @@ class TestEndToEndReconciliation: # For now, we verify the function signature and return type with ( patch("src.deriver.vector_reconciliation.tracked_db") as mock_tracked_db, - patch("src.deriver.vector_reconciliation.get_vector_store"), + patch("src.deriver.vector_reconciliation.get_external_vector_store"), patch( "src.deriver.vector_reconciliation._get_documents_needing_sync" ) as mock_get_docs,