Add AWSOpenSearch cli

Signed-off-by: Cai Yudong <[email protected]>

vectordb_bench/backend/clients/aws_opensearch/cli.py

@@ -0,0 +1,48 @@
+from typing import Annotated, TypedDict, Unpack
+
+import click
+from pydantic import SecretStr
+
+from ....cli.cli import (
+    CommonTypedDict,
+    HNSWFlavor2,
+    cli,
+    click_parameter_decorators_from_typed_dict,
+    run,
+)
+from .. import DB
+
+
+class AWSOpenSearchTypedDict(TypedDict):
+    host: Annotated[
+        str, click.option("--host", type=str, help="Db host", required=True)
+    ]
+    password: Annotated[str, click.option("--password", type=str, help="Db password")]
+    port: Annotated[int, click.option("--port", type=int, default=443, help="Db Port")]
+    timeout: Annotated[int, click.option("--timeout", type=int, default=600, help="timeout")]
+
+
+class AWSOpenSearchHNSWTypedDict(CommonTypedDict, AWSOpenSearchTypedDict, HNSWFlavor2):
+    ...
+
+
+@cli.command()
+@click_parameter_decorators_from_typed_dict(AWSOpenSearchHNSWTypedDict)
+def AWSOpenSearch(**parameters: Unpack[AWSOpenSearchHNSWTypedDict]):
+    from .config import AWSOpenSearchConfig, AWSOpenSearchIndexConfig
+    run(
+        db=DB.AWSOpenSearch,
+        db_config=AWSOpenSearchConfig(
+            hosts=[{'host': parameters["host"], 'port': parameters["port"]}],
+            http_auth=("admin", SecretStr(parameters["password"])),
+            use_ssl=True,
+            http_compress=True,
+            verify_certs=True,
+            ssl_assert_hostname=False,
+            ssl_show_warn=False,
+            timeout=parameters["timeout"],
+        ),
+        db_case_config=AWSOpenSearchIndexConfig(
+        ),
+        **parameters,
+    )

vectordb_bench/cli/vectordbbench.py

@@ -4,6 +4,7 @@
 from ..backend.clients.weaviate_cloud.cli import Weaviate
 from ..backend.clients.zilliz_cloud.cli import ZillizAutoIndex
 from ..backend.clients.milvus.cli import MilvusAutoIndex
+from ..backend.clients.aws_opensearch.cli import AWSOpenSearch
 
 
 from .cli import cli
@@ -14,6 +15,7 @@
 cli.add_command(Test)
 cli.add_command(ZillizAutoIndex)
 cli.add_command(MilvusAutoIndex)
+cli.add_command(AWSOpenSearch)
 
 
 if __name__ == "__main__":