From 092b60520f0258ae29f0611e33a90476a591dc2e Mon Sep 17 00:00:00 2001 From: Rajat Ahuja Date: Thu, 14 May 2026 13:14:15 -0400 Subject: [PATCH] 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 | 1 + src/crud/message.py | 1 + src/utils/search.py | 1 + src/vector_store/__init__.py | 3 + src/vector_store/lancedb.py | 11 +- src/vector_store/turbopuffer.py | 5 +- tests/vector_store/test_lancedb.py | 151 +++++++++++++++++++++++++ tests/vector_store/test_turbopuffer.py | 38 +++++++ 8 files changed, 209 insertions(+), 2 deletions(-) create mode 100644 tests/vector_store/test_lancedb.py diff --git a/src/crud/document.py b/src/crud/document.py index 97a57e36..fd068d6a 100644 --- a/src/crud/document.py +++ b/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: diff --git a/src/crud/message.py b/src/crud/message.py index c79c3fc4..da81fdb3 100644 --- a/src/crud/message.py +++ b/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: diff --git a/src/utils/search.py b/src/utils/search.py index 59c8f137..9c44bb1b 100644 --- a/src/utils/search.py +++ b/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: diff --git a/src/vector_store/__init__.py b/src/vector_store/__init__.py index 5a22abd9..406c524b 100644 --- a/src/vector_store/__init__.py +++ b/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) diff --git a/src/vector_store/lancedb.py b/src/vector_store/lancedb.py index f63b8cfd..640c9888 100644 --- a/src/vector_store/lancedb.py +++ b/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) diff --git a/src/vector_store/turbopuffer.py b/src/vector_store/turbopuffer.py index c84c8774..20af77dc 100644 --- a/src/vector_store/turbopuffer.py +++ b/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 diff --git a/tests/vector_store/test_lancedb.py b/tests/vector_store/test_lancedb.py new file mode 100644 index 00000000..0c502e4c --- /dev/null +++ b/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"] diff --git a/tests/vector_store/test_turbopuffer.py b/tests/vector_store/test_turbopuffer.py index 73bf3c20..ca8c4ff3 100644 --- a/tests/vector_store/test_turbopuffer.py +++ b/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