feat: add distance scores to dense search results

.github/workflows/rust.yml

@@ -20,7 +20,7 @@ jobs:
       - name: Set Up Rust
         uses: actions-rs/toolchain@v1
         with:
-          toolchain: stable
+          toolchain: 1.98.1
           components: clippy, rustfmt
           override: true
 

client/python/USAGE.md

@@ -161,11 +161,11 @@ search(
     vector: DenseVector,
     similarity: Similarity,
     limit: int,
-) -> list[str]
+) -> list[tuple[str, float]]
 ```
 
 Returns
-- List of `point_id` strings
+- List of `(point_id, score)` tuples, where score is the distance (lower is closer)
 
 Raises
 - `TypeError` if `vector` is not a `DenseVector`
@@ -182,7 +182,7 @@ batch_search(
     queries,
     similarity: Similarity | None = None,
     limit: int | None = None,
-) -> list[list[str]]
+) -> list[list[tuple[str, float]]]
 ```
 
 Returns 

client/python/tests/test_async_client.py

@@ -173,7 +173,8 @@ async def run():
             result_point_ids=[
                 Mock(id=Mock(value="p1")),
                 Mock(id=Mock(value="p2")),
-            ]
+            ],
+            scores=[0.1, 0.2],
         )
 
         results = await client.search(
@@ -182,7 +183,7 @@ async def run():
             limit=2,
         )
 
-        assert results == ["p1", "p2"]
+        assert results == [("p1", 0.1), ("p2", 0.2)]
 
     asyncio.run(run())
 
@@ -192,20 +193,21 @@ async def run():
         mock_connection.call.return_value = Mock(
             result_point_ids=[
                 Mock(id=Mock(value="p1")),
-            ]
+            ],
+            scores=[0.1],
         )
 
         q = SearchQuery(DenseVector([1, 2, 3]), Similarity.COSINE, 2)
         results = await client.search(query=q)
 
-        assert results == ["p1"]
+        assert results == [("p1", 0.1)]
 
     asyncio.run(run())
 
 
 def test_async_search_accepts_ef(client, mock_connection):
     async def run():
-        mock_connection.call.return_value = Mock(result_point_ids=[])
+        mock_connection.call.return_value = Mock(result_point_ids=[], scores=[])
 
         await client.search(
             vector=DenseVector([1, 2, 3]),
@@ -239,16 +241,16 @@ def test_async_batch_search_full_tuple(client, mock_connection):
     async def run():
         mock_connection.call.return_value = Mock(
             results=[
-                Mock(result_point_ids=[Mock(id=Mock(value="p1"))]),
-                Mock(result_point_ids=[Mock(id=Mock(value="p2"))]),
+                Mock(result_point_ids=[Mock(id=Mock(value="p1"))], scores=[0.5]),
+                Mock(result_point_ids=[Mock(id=Mock(value="p2"))], scores=[0.5]),
             ]
         )
         queries = [
             (DenseVector([1, 2, 3]), Similarity.COSINE, 2),
             (DenseVector([4, 5, 6]), Similarity.EUCLIDEAN, 1),
         ]
         result = await client.batch_search(queries=queries)
-        assert result == [["p1"], ["p2"]]
+        assert result == [[("p1", 0.5)], [("p2", 0.5)]]
 
     asyncio.run(run())
 
@@ -257,14 +259,14 @@ def test_async_batch_search_searchquery_objects(client, mock_connection):
     async def run():
         mock_connection.call.return_value = Mock(
             results=[
-                Mock(result_point_ids=[Mock(id=Mock(value="p1"))]),
+                Mock(result_point_ids=[Mock(id=Mock(value="p1"))], scores=[0.5]),
             ]
         )
         queries = [
             SearchQuery(DenseVector([1, 2, 3]), Similarity.COSINE, 2),
         ]
         result = await client.batch_search(queries=queries)
-        assert result == [["p1"]]
+        assert result == [[("p1", 0.5)]]
 
     asyncio.run(run())
 
@@ -273,7 +275,7 @@ def test_async_batch_search_vectors_with_global_params(client, mock_connection):
     async def run():
         mock_connection.call.return_value = Mock(
             results=[
-                Mock(result_point_ids=[Mock(id=Mock(value="p1"))]),
+                Mock(result_point_ids=[Mock(id=Mock(value="p1"))], scores=[0.5]),
             ]
         )
         queries = [DenseVector([1, 2, 3])]
@@ -282,7 +284,7 @@ async def run():
             similarity=Similarity.MANHATTAN,
             limit=2,
         )
-        assert result == [["p1"]]
+        assert result == [[("p1", 0.5)]]
 
     asyncio.run(run())
 
@@ -293,7 +295,7 @@ def test_async_batch_search_vector_similarity_with_global_limit(
     async def run():
         mock_connection.call.return_value = Mock(
             results=[
-                Mock(result_point_ids=[Mock(id=Mock(value="p1"))]),
+                Mock(result_point_ids=[Mock(id=Mock(value="p1"))], scores=[0.5]),
             ]
         )
         queries = [
@@ -303,7 +305,7 @@ async def run():
             queries=queries,
             limit=2,
         )
-        assert result == [["p1"]]
+        assert result == [[("p1", 0.5)]]
 
     asyncio.run(run())
 

client/python/tests/test_client.py

@@ -139,7 +139,8 @@ def test_search_success(client, mock_connection):
         result_point_ids=[
             Mock(id=Mock(value="p1")),
             Mock(id=Mock(value="p2")),
-        ]
+        ],
+        scores=[0.1, 0.2],
     )
 
     results = client.search(
@@ -148,11 +149,11 @@ def test_search_success(client, mock_connection):
         limit=2,
     )
 
-    assert results == ["p1", "p2"]
+    assert results == [("p1", 0.1), ("p2", 0.2)]
 
 
 def test_search_accepts_ef(client, mock_connection):
-    mock_connection.call.return_value = Mock(result_point_ids=[])
+    mock_connection.call.return_value = Mock(result_point_ids=[], scores=[])
 
     client.search(
         vector=DenseVector([1, 2, 3]),
@@ -195,22 +196,22 @@ def test_batch_search_accepts_ef(client, mock_connection):
 def test_batch_search_full_tuple(client, mock_connection):
     mock_connection.call.return_value = Mock(
         results=[
-            Mock(result_point_ids=[Mock(id=Mock(value="p1"))]),
-            Mock(result_point_ids=[Mock(id=Mock(value="p2"))]),
+            Mock(result_point_ids=[Mock(id=Mock(value="p1"))], scores=[0.5]),
+            Mock(result_point_ids=[Mock(id=Mock(value="p2"))], scores=[0.5]),
         ]
     )
     queries = [
         (DenseVector([1, 2, 3]), Similarity.COSINE, 2),
         (DenseVector([4, 5, 6]), Similarity.EUCLIDEAN, 1),
     ]
     result = client.batch_search(queries=queries)
-    assert result == [["p1"], ["p2"]]
+    assert result == [[("p1", 0.5)], [("p2", 0.5)]]
 
 
 def test_batch_search_vectors_with_global_params(client, mock_connection):
     mock_connection.call.return_value = Mock(
         results=[
-            Mock(result_point_ids=[Mock(id=Mock(value="p1"))]),
+            Mock(result_point_ids=[Mock(id=Mock(value="p1"))], scores=[0.5]),
         ]
     )
     queries = [DenseVector([1, 2, 3])]
@@ -219,13 +220,13 @@ def test_batch_search_vectors_with_global_params(client, mock_connection):
         similarity=Similarity.MANHATTAN,
         limit=2,
     )
-    assert result == [["p1"]]
+    assert result == [[("p1", 0.5)]]
 
 
 def test_batch_search_vector_similarity_with_global_limit(client, mock_connection):
     mock_connection.call.return_value = Mock(
         results=[
-            Mock(result_point_ids=[Mock(id=Mock(value="p1"))]),
+            Mock(result_point_ids=[Mock(id=Mock(value="p1"))], scores=[0.5]),
         ]
     )
     queries = [
@@ -235,20 +236,20 @@ def test_batch_search_vector_similarity_with_global_limit(client, mock_connectio
         queries=queries,
         limit=2,
     )
-    assert result == [["p1"]]
+    assert result == [[("p1", 0.5)]]
 
 
 def test_batch_search_searchquery_objects(client, mock_connection):
     mock_connection.call.return_value = Mock(
         results=[
-            Mock(result_point_ids=[Mock(id=Mock(value="p1"))]),
+            Mock(result_point_ids=[Mock(id=Mock(value="p1"))], scores=[0.5]),
         ]
     )
     queries = [
         SearchQuery(DenseVector([1, 2, 3]), Similarity.COSINE, 2),
     ]
     result = client.batch_search(queries=queries)
-    assert result == [["p1"]]
+    assert result == [[("p1", 0.5)]]
 
 
 def test_batch_search_missing_globals_for_vector(client):

client/python/vortexdb/async_client.py

@@ -1,4 +1,4 @@
-from typing import List
+from typing import List, Tuple
 
 from vortexdb import protoutils as proto
 from vortexdb.async_connection import AsyncGRPCConnection
@@ -98,10 +98,10 @@ async def search(
         limit: int | None = None,
         query: SearchQuery | None = None,
         ef: int | None = None,
-    ) -> List[str]:
+    ) -> List[Tuple[str, float]]:
         """
         Search for nearest neighbors.
-        Returns: List of point IDs
+        Returns: List of (point ID, score) tuples
         """
         if query is not None:
             if not isinstance(query, SearchQuery):
@@ -131,7 +131,10 @@ async def search(
             request,
         )
 
-        return [pid.id.value for pid in response.result_point_ids]
+        return [
+            (pid.id.value, score)
+            for pid, score in zip(response.result_point_ids, response.scores)
+        ]
 
     async def batch_search(
         self,
@@ -140,7 +143,7 @@ async def batch_search(
         similarity: Similarity | None = None,
         limit: int | None = None,
         ef: int | None = None,
-    ) -> List[List[str]]:
+    ) -> List[List[Tuple[str, float]]]:
         """
         Flexible batch search.
 
@@ -194,7 +197,10 @@ async def batch_search(
         request = proto.build_batch_search_request(queries=normalized, ef=ef)
         response = await self._conn.call(self._conn.stub.SearchPointsBatch, request)
         return [
-            [pid.id.value for pid in result.result_point_ids]
+            [
+                (pid.id.value, score)
+                for pid, score in zip(result.result_point_ids, result.scores)
+            ]
             for result in response.results
         ]
 

client/python/vortexdb/client.py

@@ -1,4 +1,4 @@
-from typing import List
+from typing import List, Tuple
 
 from vortexdb.connection import GRPCConnection
 from vortexdb.config import VortexDBConfig
@@ -101,10 +101,10 @@ def search(
         limit: int | None = None,
         query: SearchQuery | None = None,
         ef: int | None = None,
-    ) -> List[str]:
+    ) -> List[Tuple[str, float]]:
         """
         Search for nearest neighbors.
-        Returns: List of point IDs
+        Returns: List of (point ID, score) tuples
         """
         if query is not None:
             if not isinstance(query, SearchQuery):
@@ -129,7 +129,10 @@ def search(
             self._conn.stub.SearchPoints,
             request,
         )
-        return [pid.id.value for pid in response.result_point_ids]
+        return [
+            (pid.id.value, score)
+            for pid, score in zip(response.result_point_ids, response.scores)
+        ]
 
     def batch_search(
         self,
@@ -138,7 +141,7 @@ def batch_search(
         similarity: Similarity | None = None,
         limit: int | None = None,
         ef: int | None = None,
-    ) -> List[List[str]]:
+    ) -> List[List[Tuple[str, float]]]:
         """
         Flexible batch search.
 
@@ -192,7 +195,10 @@ def batch_search(
         request = proto.build_batch_search_request(queries=normalized, ef=ef)
         response = self._conn.call(self._conn.stub.SearchPointsBatch, request)
         return [
-            [pid.id.value for pid in result.result_point_ids]
+            [
+                (pid.id.value, score)
+                for pid, score in zip(result.result_point_ids, result.scores)
+            ]
             for result in response.results
         ]