Add batch support to the async client

Author: peopleig

client/python/USAGE.md

@@ -41,10 +41,10 @@ Example available in:
 
 ### Async Client Support
 
-For async applications, use `AsyncVortexDB`. It mirrors the synchronous client API and uses `grpc.aio` under the hood.
+For async applications, use `AsyncVortexDB`. It mirrors the synchronous client API and uses `grpc.aio` under the hood, including full support for `batch_insert` and `batch_search`.
 
-Example available in:
-```examples/async_usage.py```
+Examples available in:
+```examples/async_usage.py``` & ```examples/async_batch_usage.py```
 
 ```python
 async with AsyncVortexDB(
@@ -56,11 +56,12 @@ async with AsyncVortexDB(
         payload=Payload.text("hello async vortex"),
     )
 ```
+
 ### Batch Insertion and Search Support  
 
-The client now supports batch insertion and batch search queries.  
+Both `VortexDB` and `AsyncVortexDB` support batch insertion and batch search queries.  
 Methods of usage and examples available in:  
-```examples/batch_insert_usage.py``` & ```examples/search_query_usage.py```  
+```examples/batch_insert_usage.py``` & ```examples/search_query_usage.py``` & ```examples/async_batch_usage.py```  
 
 ---
 
@@ -76,8 +77,10 @@ Async client class for I/O-heavy applications. It has the same constructor and m
 
 ```
 await db.insert(...)
+await db.batch_insert(...)
 await db.get(...)
 await db.search(...)
+await db.batch_search(...)
 await db.delete(...)
 await db.close()
 ```
@@ -371,4 +374,4 @@ python -m grpc_tools.protoc \
 
 After running this:  
 - `vector_db_pb2_grpc.py` and `vector_db_pb2.py` will be updated
-- No other client code should need changes 
+- No other client code should need changes

client/python/examples/async_batch_usage.py

@@ -0,0 +1,86 @@
+import asyncio
+
+from vortexdb import AsyncVortexDB
+from vortexdb import DenseVector, Payload, Similarity, SearchQuery, to_dense_vectors
+
+
+async def main():
+    async with AsyncVortexDB(
+        grpc_url="localhost:50051",
+        api_key="my-secret-password",
+    ) as db:
+
+        raw_vectors = [
+            [0.1, 0.2, 0.3],
+            [0.4, 0.5, 0.6],
+            [0.7, 0.8, 0.9],
+        ]
+        vectors = to_dense_vectors(raw_vectors)
+
+        p1 = Payload.text("hello world")
+        p2 = Payload.image("/img/a.png")
+        p3 = Payload.text("foo bar")
+
+        items = [
+            (vectors[0], p1),
+            (vectors[1], p2),
+            (vectors[2], p3),
+        ]
+
+        # Batch Insert
+        point_ids = await db.batch_insert(items=items)
+        print("Inserted ids:\n", point_ids)
+
+        q = SearchQuery(
+            vector=vectors[0],
+            similarity=Similarity.COSINE,
+            limit=3,
+        )
+        res = await db.search(query=q)
+        print("\nSingle SearchQuery:\n", res)
+
+        # List of SearchQuery
+        queries = [
+            SearchQuery(vectors[0], Similarity.HAMMING, 3),
+            SearchQuery(vectors[1], Similarity.EUCLIDEAN, 2),
+            q,
+        ]
+        res = await db.batch_search(queries=queries)
+        print("\nBatch SearchQuery:\n", res)
+
+        # List of vectors with global Similarity and Limit
+        res = await db.batch_search(
+            queries=vectors,
+            similarity=Similarity.COSINE,
+            limit=3,
+        )
+        print("\nList of DenseVectors:\n", res)
+
+        # List of tuple (DenseVector, Similarity) with global Limit
+        queries = [
+            (vectors[0], Similarity.COSINE),
+            (vectors[1], Similarity.MANHATTAN),
+        ]
+        res = await db.batch_search(
+            queries=queries,
+            limit=3,
+        )
+        print("\nList of (DenseVector, Similarity):\n", res)
+
+        # List of tuple (DenseVector, Limit) with global Similarity
+        queries = [
+            (vectors[0], 2),
+            (vectors[1], 4),
+        ]
+        res = await db.batch_search(
+            queries=queries,
+            similarity=Similarity.COSINE,
+        )
+        print("\nList of (DenseVector, Limit):\n", res)
+
+        for pid in point_ids:
+            await db.delete(point_id=pid)
+
+
+if __name__ == "__main__":
+    asyncio.run(main())

client/python/tests/test_async_client.py

@@ -5,7 +5,7 @@
 
 from vortexdb.async_client import AsyncVortexDB
 from vortexdb.async_connection import AsyncGRPCConnection
-from vortexdb.models import ContentType, DenseVector, Payload, Point, Similarity
+from vortexdb.models import ContentType, DenseVector, Payload, Point, Similarity, SearchQuery
 
 
 @pytest.fixture
@@ -58,6 +58,59 @@ async def run():
     asyncio.run(run())
 
 
+# Batch Insert
+
+def test_async_batch_insert_success(client, mock_connection):
+    async def run():
+        response = Mock()
+        response.ids = [
+            Mock(id=Mock(value="p1")),
+            Mock(id=Mock(value="p2")),
+        ]
+        mock_connection.call.return_value = response
+
+        items = [
+            (DenseVector([1, 2, 3]), Payload.text("a")),
+            (DenseVector([4, 5, 6]), Payload.text("b")),
+        ]
+        result = await client.batch_insert(items=items)
+        assert result == ["p1", "p2"]
+
+    asyncio.run(run())
+
+
+def test_async_batch_insert_invalid_items_type(client):
+    async def run():
+        with pytest.raises(TypeError):
+            await client.batch_insert(items="not-a-list")
+
+    asyncio.run(run())
+
+
+def test_async_batch_insert_invalid_tuple_structure(client):
+    async def run():
+        items = [
+            (DenseVector([1, 2, 3]),),  # only one element
+        ]
+        with pytest.raises(TypeError):
+            await client.batch_insert(items=items)
+
+    asyncio.run(run())
+
+
+def test_async_batch_insert_invalid_vector(client):
+    async def run():
+        items = [
+            ([1, 2, 3], Payload.text("a")),  # not DenseVector
+        ]
+        with pytest.raises(TypeError):
+            await client.batch_insert(items=items)
+
+    asyncio.run(run())
+
+
+# Get
+
 def test_async_get_point_success(client, mock_connection):
     async def run():
         proto_point = Mock()
@@ -88,6 +141,8 @@ async def run():
     asyncio.run(run())
 
 
+# Delete
+
 def test_async_delete_success(client, mock_connection):
     async def run():
         mock_connection.call.return_value = None
@@ -99,6 +154,8 @@ async def run():
     asyncio.run(run())
 
 
+# Search
+
 def test_async_search_success(client, mock_connection):
     async def run():
         mock_connection.call.return_value = Mock(
@@ -119,6 +176,39 @@ async def run():
     asyncio.run(run())
 
 
+def test_async_search_with_query_object(client, mock_connection):
+    async def run():
+        mock_connection.call.return_value = Mock(
+            result_point_ids=[
+                Mock(id=Mock(value="p1")),
+            ]
+        )
+
+        q = SearchQuery(DenseVector([1, 2, 3]), Similarity.COSINE, 2)
+        results = await client.search(query=q)
+
+        assert results == ["p1"]
+
+    asyncio.run(run())
+
+
+def test_async_search_accepts_ef(client, mock_connection):
+    async def run():
+        mock_connection.call.return_value = Mock(result_point_ids=[])
+
+        await client.search(
+            vector=DenseVector([1, 2, 3]),
+            similarity=Similarity.COSINE,
+            limit=2,
+            ef=128,
+        )
+
+        request = mock_connection.call.call_args.args[1]
+        assert request.ef == 128
+
+    asyncio.run(run())
+
+
 def test_async_search_invalid_vector(client):
     async def run():
         with pytest.raises(TypeError):
@@ -131,6 +221,128 @@ async def run():
     asyncio.run(run())
 
 
+# Batch Search
+
+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"))]),
+            ]
+        )
+        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"]]
+
+    asyncio.run(run())
+
+
+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"))]),
+            ]
+        )
+        queries = [
+            SearchQuery(DenseVector([1, 2, 3]), Similarity.COSINE, 2),
+        ]
+        result = await client.batch_search(queries=queries)
+        assert result == [["p1"]]
+
+    asyncio.run(run())
+
+
+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"))]),
+            ]
+        )
+        queries = [DenseVector([1, 2, 3])]
+        result = await client.batch_search(
+            queries=queries,
+            similarity=Similarity.MANHATTAN,
+            limit=2,
+        )
+        assert result == [["p1"]]
+
+    asyncio.run(run())
+
+
+def test_async_batch_search_vector_similarity_with_global_limit(client, mock_connection):
+    async def run():
+        mock_connection.call.return_value = Mock(
+            results=[
+                Mock(result_point_ids=[Mock(id=Mock(value="p1"))]),
+            ]
+        )
+        queries = [
+            (DenseVector([1, 2, 3]), Similarity.COSINE),
+        ]
+        result = await client.batch_search(
+            queries=queries,
+            limit=2,
+        )
+        assert result == [["p1"]]
+
+    asyncio.run(run())
+
+
+def test_async_batch_search_accepts_ef(client, mock_connection):
+    async def run():
+        mock_connection.call.return_value = Mock(results=[])
+
+        await client.batch_search(
+            queries=[
+                (DenseVector([1, 2, 3]), Similarity.COSINE, 2),
+                (DenseVector([4, 5, 6]), Similarity.COSINE, 1),
+            ],
+            ef=256,
+        )
+
+        request = mock_connection.call.call_args.args[1]
+        assert [query.ef for query in request.queries] == [256, 256]
+
+    asyncio.run(run())
+
+
+def test_async_batch_search_missing_globals_for_vector(client):
+    async def run():
+        queries = [DenseVector([1, 2, 3])]
+        with pytest.raises(ValueError):
+            await client.batch_search(queries=queries)
+
+    asyncio.run(run())
+
+
+def test_async_batch_search_missing_limit(client):
+    async def run():
+        queries = [
+            (DenseVector([1, 2, 3]), Similarity.COSINE),
+        ]
+        with pytest.raises(ValueError):
+            await client.batch_search(queries=queries)
+
+    asyncio.run(run())
+
+
+def test_async_batch_search_invalid_format(client):
+    async def run():
+        queries = ["invalid"]
+        with pytest.raises(TypeError):
+            await client.batch_search(queries=queries)
+
+    asyncio.run(run())
+
+
+# Close / Context Manager
+
 def test_async_close_closes_connection(client, mock_connection):
     async def run():
         await client.close()
@@ -155,4 +367,4 @@ async def run():
 
         conn.close.assert_awaited_once()
 
-    asyncio.run(run())
+    asyncio.run(run())