diff --git a/src/crud/document.py b/src/crud/document.py index b423c77a..edc8d9a1 100644 --- a/src/crud/document.py +++ b/src/crud/document.py @@ -151,7 +151,32 @@ async def query_documents( f"Query exceeds maximum token limit of {settings.MAX_EMBEDDING_TOKENS}." ) from e - # Get vector store and namespace for this collection + # 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": + stmt = ( + select(models.Document) + .where(models.Document.workspace_name == workspace_name) + .where(models.Document.observer == observer) + .where(models.Document.observed == observed) + .where(models.Document.embedding.isnot(None)) + .where(models.Document.deleted_at.is_(None)) + ) + + if max_distance is not None: + stmt = stmt.where( + models.Document.embedding.cosine_distance(embedding) <= max_distance + ) + + stmt = apply_filter(stmt, models.Document, filters) + stmt = stmt.order_by( + models.Document.embedding.cosine_distance(embedding) + ).limit(top_k) + + 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( "document", workspace_name, observer, observed diff --git a/src/utils/search.py b/src/utils/search.py index 9431aedf..1d17ed4a 100644 --- a/src/utils/search.py +++ b/src/utils/search.py @@ -91,7 +91,40 @@ async def _semantic_search( f"Query exceeds maximum token limit of {settings.MAX_EMBEDDING_TOKENS}." ) from e - # Get vector store and namespace for this workspace's messages + # 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": + # Join message_embeddings with messages to get full message objects + distance_expr = models.MessageEmbedding.embedding.cosine_distance( + embedding_query + ) + + stmt = ( + select(models.Message) + .join( + models.MessageEmbedding, + models.Message.public_id == models.MessageEmbedding.message_id, + ) + .where(models.MessageEmbedding.embedding.isnot(None)) + .where(models.MessageEmbedding.workspace_name == workspace_name) + ) + + # Apply all additional filters using the standard filter utility + # filters dict uses external names (session_id, peer_id) which apply_filter will map + # to internal column names (session_name, peer_name) + if filters: + # Create a copy with workspace added + internal_filters = filters.copy() + internal_filters["workspace_id"] = workspace_name + stmt = apply_filter(stmt, models.Message, internal_filters) + + # Order by cosine distance and limit + stmt = stmt.order_by(distance_expr).limit(limit) + + 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) diff --git a/src/vector_store/pgvector.py b/src/vector_store/pgvector.py index bc0b4701..d46ebde5 100644 --- a/src/vector_store/pgvector.py +++ b/src/vector_store/pgvector.py @@ -8,7 +8,7 @@ using the existing embedding columns on documents and message_embeddings tables. import logging from typing import Any -from sqlalchemy import select, update +from sqlalchemy import delete, select from sqlalchemy.ext.asyncio import AsyncSession from src import models @@ -280,7 +280,10 @@ class PgVectorStore(VectorStore): async def delete_many(self, namespace: str, ids: list[str]) -> None: """ - Delete vectors by setting embedding to NULL. + Delete vectors by removing the rows from the database. + + For pgvector, since the vector data is stored in the same table as the entities, + deleting from the vector store means deleting the actual rows. Args: namespace: The namespace containing the vectors @@ -294,11 +297,7 @@ class PgVectorStore(VectorStore): async with tracked_db("pgvector_delete") as db: try: if table_type == "documents": - stmt = ( - update(models.Document) - .where(models.Document.id.in_(ids)) - .values(embedding=None) - ) + stmt = delete(models.Document).where(models.Document.id.in_(ids)) await db.execute(stmt) elif table_type == "message_embeddings": @@ -310,28 +309,29 @@ class PgVectorStore(VectorStore): f"Invalid message vector id format: {vector_id}" ) from exc - stmt = ( - update(models.MessageEmbedding) - .where(models.MessageEmbedding.id == embedding_id) - .values(embedding=None) + stmt = delete(models.MessageEmbedding).where( + models.MessageEmbedding.id == embedding_id ) await db.execute(stmt) await db.commit() logger.debug( - f"Deleted {len(ids)} vectors from {table_type} in namespace {namespace}" + f"Deleted {len(ids)} rows from {table_type} in namespace {namespace}" ) except Exception: await db.rollback() logger.exception( - f"Failed to delete {len(ids)} vectors from namespace {namespace}" + f"Failed to delete {len(ids)} rows from namespace {namespace}" ) raise async def delete_namespace(self, namespace: str) -> None: """ - Delete all vectors in a namespace by setting embedding to NULL. + Delete all vectors in a namespace by removing rows from the database. + + For pgvector, since the vector data is stored in the same table as the entities, + deleting a namespace means deleting the actual rows. Args: namespace: The namespace to delete @@ -342,29 +342,24 @@ class PgVectorStore(VectorStore): try: if table_type == "documents": stmt = ( - update(models.Document) + delete(models.Document) .where( models.Document.workspace_name == context["workspace_name"] ) .where(models.Document.observer == context["observer"]) .where(models.Document.observed == context["observed"]) - .values(embedding=None) ) await db.execute(stmt) elif table_type == "message_embeddings": - stmt = ( - update(models.MessageEmbedding) - .where( - models.MessageEmbedding.workspace_name - == context["workspace_name"] - ) - .values(embedding=None) + stmt = delete(models.MessageEmbedding).where( + models.MessageEmbedding.workspace_name + == context["workspace_name"] ) await db.execute(stmt) await db.commit() - logger.debug(f"Deleted all vectors from namespace {namespace}") + logger.debug(f"Deleted all rows from namespace {namespace}") except Exception: await db.rollback()