Add S3vectors cli support (#600)

* Add s3vectors cli support

* Last commit.

Co-authored-by: Enrique enrique.molina@urv.cat
Co-authored-by: Octavio octaviohoracio.iacoponelli@urv.cat

* Fix linting with ruff --fix and black.

Co-authored-by: Enrique enrique.molina@urv.cat
Co-authored-by: Octavio octaviohoracio.iacoponelli@urv.cat

Author: acanadil

vectordb_bench/backend/clients/s3_vectors/cli.py

@@ -0,0 +1,67 @@
+from typing import Annotated, TypedDict, Unpack
+
+import click
+from pydantic import SecretStr
+
+from ....cli.cli import (
+    CommonTypedDict,
+    cli,
+    click_parameter_decorators_from_typed_dict,
+    get_custom_case_config,
+    run,
+)
+from .. import DB
+from ..api import MetricType
+from .config import S3VectorsIndexConfig
+
+
+class S3VectorsTypedDict(TypedDict):
+    region_name: Annotated[
+        str, click.option("--region", type=str, help="AWS region for S3 bucket (eg. us-east-1)", default="us-east-1")
+    ]
+    access_key_id: Annotated[str, click.option("--access_key_id", type=str, help="AWS access key ID", required=True)]
+    secret_access_key: Annotated[
+        str, click.option("--secret_access_key", type=str, help="AWS secret access key", required=True)
+    ]
+
+    bucket: Annotated[str, click.option("--bucket", type=str, help="S3 bucket name", required=True)]
+    index: Annotated[str, click.option("--index", type=str, help="Unique vector index name", default="vdbbench-index")]
+
+    metric: Annotated[
+        str,
+        click.option(
+            "--metric",
+            type=str,
+            help="Distance metric for vector similarity (e.g., 'cosine', 'euclidean').",
+            default=None,
+        ),
+    ]
+
+
+class S3VectorsIndexTypedDict(CommonTypedDict, S3VectorsTypedDict): ...
+
+
+@cli.command()
+@click_parameter_decorators_from_typed_dict(S3VectorsIndexTypedDict)
+def S3Vectors(**parameters: Unpack[S3VectorsIndexTypedDict]):
+    from .config import S3VectorsConfig
+
+    parameters["custom_case"] = get_custom_case_config(parameters)
+    run(
+        db=DB.S3Vectors,
+        db_config=S3VectorsConfig(
+            region_name=parameters["region"],
+            access_key_id=SecretStr(parameters["access_key_id"]),
+            secret_access_key=SecretStr(parameters["secret_access_key"]),
+            bucket_name=parameters["bucket"],
+            index_name=parameters["index"] if parameters["index"] else "vdbbench-index",
+        ),
+        db_case_config=S3VectorsIndexConfig(
+            metric_type=(
+                MetricType.COSINE
+                if parameters["metric"] == "cosine"
+                else MetricType.L2 if parameters["metric"] == "l2" else None
+            )
+        ),
+        **parameters,
+    )

vectordb_bench/cli/vectordbbench.py

@@ -15,6 +15,7 @@
 from ..backend.clients.qdrant_cloud.cli import QdrantCloud
 from ..backend.clients.qdrant_local.cli import QdrantLocal
 from ..backend.clients.redis.cli import Redis
+from ..backend.clients.s3_vectors.cli import S3Vectors
 from ..backend.clients.test.cli import Test
 from ..backend.clients.tidb.cli import TiDB
 from ..backend.clients.vespa.cli import Vespa
@@ -48,6 +49,7 @@
 cli.add_command(QdrantCloud)
 cli.add_command(QdrantLocal)
 cli.add_command(BatchCli)
+cli.add_command(S3Vectors)
 
 
 if __name__ == "__main__":