feat: stop fetching embedding vectors on vector store query - DEV-1727 (#682)

* feat: stop fetching embedding vectors on vector store query

* fix: add similar filtering for lancedb

* fix: add lancedb tests

---------

Co-authored-by: Vineeth Voruganti <13438633+VVoruganti@users.noreply.github.com>

src/crud/document.py

@@ -242,6 +242,7 @@ async def query_external_vector_document_ids(
         top_k=top_k,
         max_distance=max_distance,
         filters=vector_filters if vector_filters else None,
+        include_attributes=False,
     )
 
     if not vector_results:

src/crud/message.py

@@ -694,6 +694,7 @@ async def _search_messages_external(
         query_embedding,
         top_k=limit * oversample,
         filters=vector_filters if vector_filters else None,
+        include_attributes=["message_id"],
     )
 
     if not vector_results:

src/utils/search.py

@@ -100,6 +100,7 @@ async def query_external_vector_message_ids(
         embedding_query,
         top_k=limit * 3,
         filters=vector_filters if vector_filters else None,
+        include_attributes=["message_id"],
     )
 
     if not vector_results:

src/vector_store/__init__.py

@@ -134,6 +134,7 @@ class VectorStore(ABC):
         top_k: int = 10,
         filters: dict[str, Any] | None = None,
         max_distance: float | None = None,
+        include_attributes: bool | list[str] = True,
     ) -> list[VectorQueryResult]:
         """
         Query for similar vectors.
@@ -144,6 +145,8 @@ class VectorStore(ABC):
             top_k: Maximum number of results to return
             filters: Optional metadata filters
             max_distance: Optional maximum distance threshold (cosine distance)
+            include_attributes: Attributes to return with each result. Use False when
+                callers only need IDs/scores, or a list for selected metadata.
 
         Returns:
             List of VectorQueryResult objects, ordered by similarity (most similar first)

src/vector_store/lancedb.py

@@ -197,6 +197,7 @@ class LanceDBVectorStore(VectorStore):
         top_k: int = 10,
         filters: dict[str, Any] | None = None,
         max_distance: float | None = None,
+        include_attributes: bool | list[str] = True,
     ) -> list[VectorQueryResult]:
         """
         Query for similar vectors in LanceDB.
@@ -207,6 +208,8 @@ class LanceDBVectorStore(VectorStore):
             top_k: Maximum number of results to return
             filters: Optional metadata filters
             max_distance: Optional maximum distance threshold (cosine distance)
+            include_attributes: Attributes to return with each result. False returns
+                no metadata; a list returns only those metadata fields.
 
         Returns:
             List of VectorQueryResult objects, ordered by similarity (most similar first)
@@ -217,9 +220,15 @@ class LanceDBVectorStore(VectorStore):
             return []
 
         try:
-            # Build query
             query = table.vector_search(embedding).distance_type("cosine").limit(top_k)
 
+            if include_attributes is False:
+                # Caller only needs id/score. Don't fetch any metadata or the vector.
+                query = query.select(["id"])
+            elif isinstance(include_attributes, list):
+                projection = ["id", *(c for c in include_attributes if c != "id")]
+                query = query.select(projection)
+
             # Apply filters if provided
             if filters:
                 where_clause = self._build_where_clause(filters)

src/vector_store/turbopuffer.py

@@ -116,6 +116,7 @@ class TurbopufferVectorStore(VectorStore):
         top_k: int = 10,
         filters: dict[str, Any] | None = None,
         max_distance: float | None = None,
+        include_attributes: bool | list[str] = True,
     ) -> list[VectorQueryResult]:
         """
         Query for similar vectors in Turbopuffer.
@@ -126,6 +127,8 @@ class TurbopufferVectorStore(VectorStore):
             top_k: Maximum number of results to return
             filters: Optional metadata filters
             max_distance: Optional maximum distance threshold (cosine distance)
+            include_attributes: Attributes to include in the response. Passing False
+                avoids parsing unused row attributes.
 
         Returns:
             List of VectorQueryResult objects, ordered by similarity (most similar first)
@@ -149,7 +152,7 @@ class TurbopufferVectorStore(VectorStore):
                 "rank_by": rank_by,
                 "top_k": top_k,
                 "distance_metric": DISTANCE_METRIC,
-                "include_attributes": True,
+                "include_attributes": include_attributes,
             }
             if filter_condition is not None:
                 query_kwargs["filters"] = filter_condition

tests/vector_store/test_lancedb.py

@@ -0,0 +1,151 @@
+"""Tests for LanceDBVectorStore query projection behavior."""
+
+from __future__ import annotations
+
+from typing import Any
+from unittest.mock import AsyncMock, MagicMock
+
+import pytest
+
+from src.vector_store.lancedb import LanceDBVectorStore
+
+
+def _build_query_chain(rows: list[dict[str, Any]]) -> MagicMock:
+    """Build a chainable mock that mirrors LanceDB's async query builder."""
+    chain = MagicMock()
+    chain.distance_type.return_value = chain
+    chain.limit.return_value = chain
+    chain.select.return_value = chain
+    chain.where.return_value = chain
+    chain.to_list = AsyncMock(return_value=rows)
+    return chain
+
+
+def _patch_table(
+    store: LanceDBVectorStore, rows: list[dict[str, Any]]
+) -> tuple[MagicMock, MagicMock]:
+    """Patch _get_table to return a mock whose vector_search yields the chain."""
+    chain = _build_query_chain(rows)
+    table = MagicMock()
+    table.vector_search = MagicMock(return_value=chain)
+    store._get_table = AsyncMock(return_value=table)  # pyright: ignore[reportPrivateUsage]
+    return table, chain
+
+
+@pytest.fixture
+def store() -> LanceDBVectorStore:
+    return LanceDBVectorStore()
+
+
+@pytest.mark.asyncio
+async def test_query_returns_empty_when_table_missing(
+    store: LanceDBVectorStore,
+) -> None:
+    store._get_table = AsyncMock(return_value=None)  # pyright: ignore[reportPrivateUsage]
+
+    results = await store.query("honcho.msg.missing", [0.1, 0.2, 0.3, 0.4])
+
+    assert results == []
+
+
+@pytest.mark.asyncio
+async def test_query_default_does_not_project(store: LanceDBVectorStore) -> None:
+    _table, chain = _patch_table(store, rows=[])
+
+    await store.query("honcho.msg.test", [0.1, 0.2, 0.3, 0.4])
+
+    chain.select.assert_not_called()
+
+
+@pytest.mark.asyncio
+async def test_query_with_include_attributes_false_selects_only_id(
+    store: LanceDBVectorStore,
+) -> None:
+    _table, chain = _patch_table(store, rows=[])
+
+    await store.query(
+        "honcho.doc.test",
+        [0.1, 0.2, 0.3, 0.4],
+        include_attributes=False,
+    )
+
+    chain.select.assert_called_once_with(["id"])
+
+
+@pytest.mark.asyncio
+async def test_query_with_attribute_list_projects_id_plus_listed(
+    store: LanceDBVectorStore,
+) -> None:
+    _table, chain = _patch_table(store, rows=[])
+
+    await store.query(
+        "honcho.msg.test",
+        [0.1, 0.2, 0.3, 0.4],
+        include_attributes=["message_id"],
+    )
+
+    chain.select.assert_called_once_with(["id", "message_id"])
+
+
+@pytest.mark.asyncio
+async def test_query_attribute_list_dedupes_explicit_id(
+    store: LanceDBVectorStore,
+) -> None:
+    _table, chain = _patch_table(store, rows=[])
+
+    await store.query(
+        "honcho.msg.test",
+        [0.1, 0.2, 0.3, 0.4],
+        include_attributes=["id", "message_id"],
+    )
+
+    chain.select.assert_called_once_with(["id", "message_id"])
+
+
+@pytest.mark.asyncio
+async def test_query_converts_rows_to_results_with_score_and_metadata(
+    store: LanceDBVectorStore,
+) -> None:
+    rows: list[dict[str, Any]] = [
+        {
+            "id": "vec_1",
+            "_distance": 0.12,
+            "vector": [0.0, 0.0, 0.0, 0.0],
+            "message_id": "msg_1",
+            "session_name": "sess_a",
+        },
+        {
+            "id": "vec_2",
+            "_distance": 0.34,
+            "message_id": "msg_2",
+        },
+    ]
+    _patch_table(store, rows=rows)
+
+    results = await store.query("honcho.msg.test", [0.1, 0.2, 0.3, 0.4])
+
+    assert [r.id for r in results] == ["vec_1", "vec_2"]
+    assert [r.score for r in results] == [0.12, 0.34]
+    # id, vector, _distance must not leak into metadata
+    assert results[0].metadata == {
+        "message_id": "msg_1",
+        "session_name": "sess_a",
+    }
+    assert results[1].metadata == {"message_id": "msg_2"}
+
+
+@pytest.mark.asyncio
+async def test_query_filters_by_max_distance(store: LanceDBVectorStore) -> None:
+    rows: list[dict[str, Any]] = [
+        {"id": "vec_close", "_distance": 0.05, "message_id": "msg_1"},
+        {"id": "vec_far", "_distance": 0.9, "message_id": "msg_2"},
+    ]
+    _patch_table(store, rows=rows)
+
+    results = await store.query(
+        "honcho.msg.test",
+        [0.1, 0.2, 0.3, 0.4],
+        max_distance=0.5,
+    )
+
+    assert [r.id for r in results] == ["vec_close"]

tests/vector_store/test_turbopuffer.py

@@ -79,3 +79,41 @@ async def test_upsert_many_succeeds_without_raising(
 
     assert result is None
     namespace_mock.write.assert_awaited_once()
+
+
+@pytest.mark.asyncio
+async def test_query_passes_requested_include_attributes(
+    store: TurbopufferVectorStore,
+) -> None:
+    namespace_mock = MagicMock()
+    namespace_mock.query = AsyncMock(return_value=MagicMock(rows=[]))
+    store._get_namespace = MagicMock(return_value=namespace_mock)  # pyright: ignore[reportPrivateUsage]
+
+    await store.query(
+        "honcho.msg.test",
+        [0.1, 0.2, 0.3, 0.4],
+        include_attributes=["message_id"],
+    )
+
+    namespace_mock.query.assert_awaited_once()
+    assert namespace_mock.query.await_args.kwargs["include_attributes"] == [
+        "message_id"
+    ]
+
+
+@pytest.mark.asyncio
+async def test_query_can_skip_attributes(
+    store: TurbopufferVectorStore,
+) -> None:
+    namespace_mock = MagicMock()
+    namespace_mock.query = AsyncMock(return_value=MagicMock(rows=[]))
+    store._get_namespace = MagicMock(return_value=namespace_mock)  # pyright: ignore[reportPrivateUsage]
+
+    await store.query(
+        "honcho.doc.test",
+        [0.1, 0.2, 0.3, 0.4],
+        include_attributes=False,
+    )
+
+    namespace_mock.query.assert_awaited_once()
+    assert namespace_mock.query.await_args.kwargs["include_attributes"] is False