diff --git a/docs/api.md b/docs/api.md index 7dd0cef..0eda6e6 100644 --- a/docs/api.md +++ b/docs/api.md @@ -152,6 +152,7 @@ docs.query( measure="cosine_distance", # distance measure to use include_value=False, # should distance measure values be returned? include_metadata=False, # should record metadata be returned? + include_vector=False, # should vectors be returned? ) ``` diff --git a/src/tests/test_collection.py b/src/tests/test_collection.py index b7cde59..0653329 100644 --- a/src/tests/test_collection.py +++ b/src/tests/test_collection.py @@ -244,6 +244,18 @@ def test_query(client: vecs.Client) -> None: assert res[0][0] == "vec5" assert res[0][1] == query_meta + # include_vector + res = bar.query( + data=query_vec, + limit=top_k, + filters=None, + measure="cosine_distance", + include_vector=True, + ) + assert len(res[0]) == 2 + assert res[0][0] == "vec5" + assert all(res[0][1] == query_vec) + # test for different numbers of probes assert len(bar.query(data=query_vec, limit=top_k, probes=10)) == top_k diff --git a/src/vecs/collection.py b/src/vecs/collection.py index e00ebce..eab8f93 100644 --- a/src/vecs/collection.py +++ b/src/vecs/collection.py @@ -456,6 +456,7 @@ def query( measure: Union[IndexMeasure, str] = IndexMeasure.cosine_distance, include_value: bool = False, include_metadata: bool = False, + include_vector: bool = False, *, probes: Optional[int] = None, ef_search: Optional[int] = None, @@ -537,6 +538,9 @@ def query( if include_value: cols.append(distance_clause) + if include_vector: + cols.append(self.table.c.vec) + if include_metadata: cols.append(self.table.c.metadata)