Async batch support (#85)

* Update Python client to support batch insert and search

* Add batch support to the async client

* Add master test script for all usage example files

Run linter on client/python

* Add test for api parity between sync and async classes

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(
@@ -57,6 +57,12 @@ async with AsyncVortexDB(
     )
 ```
 
+### Batch Insertion and Search Support  
+
+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/async_batch_usage.py```  
+
 ---
 
 ## Client API
@@ -71,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()
 ```
@@ -115,6 +123,22 @@ Raises
 
 ---
 
+#### **Batch Insert**
+
+Insert multiple vectors with payloads in a single request  
+```
+batch_insert(*, items: list[tuple[DenseVector, Payload]]) -> list[str]
+```
+
+Returns
+- List of `point_id` (UUID string)
+
+Raises
+- `TypeError` if input structure is invalid
+- gRPC-mapped errors (see Error Handling)
+
+---
+
 #### **Get**
 
 Fetch a point by its ID
@@ -149,6 +173,32 @@ Raises
 
 ---
 
+#### **Batch Search**
+
+Search for nearest neighbours for multiple queries in a single request
+```
+batch_search(
+    *,
+    queries,
+    similarity: Similarity | None = None,
+    limit: int | None = None,
+) -> list[list[str]]
+```
+
+Returns 
+- `TypeError` for invalid query formats
+- `ValueError` if required parameters are missing
+
+Supported Input Formats:  
+The `queries` parameter is flexible and supports multiple formats:
+- List of `SearchQuery` objects
+- List of `(DenseVector, Similarity, Limit)` tuples
+- List of `(DenseVector, Similarity)` tuples with a global `Limit`
+- List of `(DenseVector, Limit)` tuples with a global `Similarity`
+- List of `DenseVector` with global `Similarity` and `Limit`
+
+---
+
 #### **Delete**
 
 Delete a point by its ID
@@ -214,6 +264,19 @@ All fields are directly accessible:
 
 ---
 
+### `SearchQuery`
+
+```
+SearchQuery(
+    vector: DenseVector,
+    similarity: Similarity,
+    limit: int,
+)
+```
+Structured representation of a search request  
+
+---
+
 ### `Similarity`
 
 Enum representing distance functions: 
@@ -311,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/all.py

@@ -0,0 +1,28 @@
+# This file is like a master test. Runs all the examples
+# Not exactly the purpose of the examples dir,
+# but helps in checking if any code updates haven't broken the API
+
+from pathlib import Path
+import subprocess
+import pytest
+
+# Didn't know I could do this with pytest, so cool
+# Just run: pytest ./all.py -v
+
+EXAMPLES_DIR = Path(__file__).parent
+example_files = sorted(EXAMPLES_DIR.glob("*_usage.py"))
+
+
+@pytest.mark.parametrize(
+    "script_path",
+    example_files,
+    ids=lambda p: p.stem,
+)
+def test(script_path):
+    """Run all example scripts to check if they crash or not"""
+    result = subprocess.run(
+        ["python3", str(script_path)], capture_output=True, text=True
+    )
+    assert result.returncode == 0, (
+        f"Script {script_path} failed with stderr:\n{result.stderr}"
+    )

client/python/examples/async_batch_usage.py

@@ -0,0 +1,85 @@
+import asyncio
+
+from vortexdb import AsyncVortexDB
+from vortexdb import 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/examples/async_usage.py

@@ -6,7 +6,7 @@
 async def main():
     async with AsyncVortexDB(
         grpc_url="localhost:50051",
-        api_key="your-api-key",
+        api_key="my-secret-password",
     ) as db:
         point_id = await db.insert(
             vector=DenseVector([0.1, 0.2, 0.3]),

client/python/examples/basic_usage.py

@@ -1,5 +1,6 @@
 from vortexdb import VortexDB
-from vortexdb import DenseVector, Payload, Similarity   # from vortexdb.models
+from vortexdb import DenseVector, Payload, Similarity  # from vortexdb.models
+
 
 def main():
     # Initialize client
@@ -32,5 +33,6 @@ def main():
     # Close connection
     db.close()
 
+
 if __name__ == "__main__":
     main()

client/python/examples/batch_insert_usage.py

@@ -0,0 +1,39 @@
+from vortexdb import VortexDB
+from vortexdb import Payload, to_dense_vectors
+
+
+def main():
+    db = VortexDB(
+        grpc_url="localhost:50051",
+        api_key="my-secret-password",
+    )
+
+    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 = db.batch_insert(items=items)
+    print("Inserted ids:\n", point_ids)
+
+    for pid in point_ids:
+        db.delete(point_id=pid)
+
+    db.close()
+
+
+if __name__ == "__main__":
+    main()

client/python/examples/context_manager_usage.py

@@ -1,11 +1,11 @@
 from vortexdb import VortexDB, DenseVector, Payload, Similarity
 
+
 def main():
     with VortexDB(
         grpc_url="localhost:50051",
         api_key="my-secret-password",
     ) as db:
-
         # Insert a vector
         point_id = db.insert(
             vector=DenseVector([0.1, 0.2, 0.3]),
@@ -30,5 +30,6 @@ def main():
     # At this point, the gRPC channel is closed automatically
     print("Connection closed")
 
+
 if __name__ == "__main__":
     main()

client/python/examples/search_query_usage.py

@@ -0,0 +1,69 @@
+from vortexdb import VortexDB
+from vortexdb import Similarity, SearchQuery, to_dense_vectors
+
+
+def main():
+    db = VortexDB(
+        grpc_url="localhost:50051",
+        api_key="my-secret-password",
+    )
+
+    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)
+
+    q = SearchQuery(
+        vector=vectors[0],
+        similarity=Similarity.COSINE,
+        limit=3,
+    )
+    res = db.search(query=q)
+    print("Single SearchQuery:\n", res)
+
+    # List of SearchQuery
+    queries = [
+        SearchQuery(vectors[0], Similarity.HAMMING, 3),
+        SearchQuery(vectors[1], Similarity.EUCLIDEAN, 2),
+        q,
+    ]
+    res = db.batch_search(queries=queries)
+    print("\nBatch SearchQuery:\n", res)
+
+    # List of vectors with global Similarity and Limit
+    res = 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 = 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 = db.batch_search(
+        queries=queries,
+        similarity=Similarity.COSINE,
+    )
+    print("\nList of (DenseVector, Limit):\n", res)
+
+    db.close()
+
+
+if __name__ == "__main__":
+    main()