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

fix: remove composite and pgvector store -> make more specific

This commit is contained in:
Rajat Ahuja 2026-01-12 18:11:22 -05:00
parent 6c12e6f413
commit b04770d602
18 changed files with 429 additions and 1035 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

11
.env.template
View File

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

6
config.toml.example
View File

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

16
src/config.py
View File

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

283
src/crud/document.py
View File

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

196
src/crud/message.py
View File

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

24
src/crud/session.py
View File

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

65
src/crud/workspace.py
View File

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

131
src/deriver/vector_reconciliation.py
View File

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

4
src/main.py
View File

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

22
src/utils/search.py
View File

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

67
src/vector_store/__init__.py
View File

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

261
src/vector_store/composite.py
View File

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

4
src/vector_store/lancedb.py
View File

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

315
src/vector_store/pgvector.py
View File

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

4
src/vector_store/turbopuffer.py
View File

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

6
src/vector_store/utils.py
View File

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

23
tests/conftest.py
View File

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

26
tests/deriver/test_vector_reconciliation.py
View File

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