When adding a new client, you can enable command-line execution by defining a CLI module. This uses click and pydantic to map command-line arguments to the client's configuration.
Key components:
CommonTypedDict: Used to define the expected CLI arguments.@click.option: Used within the Annotated type hints to define flags (e.g., --uri, --password).run(): The core function that executes the benchmark with the provided db, db_config, and db_case_config.
Note: For databases with multiple index configurations (like pgvector or milvus), you should repeat this process for each index configuration type.
from typing import Annotated, Unpack
import click
import os
from pydantic import SecretStr
from vectordb_bench.cli.cli import (
CommonTypedDict,
cli,
click_parameter_decorators_from_typed_dict,
run,
)
from vectordb_bench.backend.clients import DB
class ZillizTypedDict(CommonTypedDict):
uri: Annotated[
str, click.option("--uri", type=str, help="uri connection string", required=True)
]
user_name: Annotated[
str, click.option("--user-name", type=str, help="Db username", required=True)
]
password: Annotated[
str,
click.option(
"--password",
type=str,
help="Zilliz password",
default=lambda: os.environ.get("ZILLIZ_PASSWORD", ""),
show_default="$ZILLIZ_PASSWORD",
),
]
level: Annotated[
str, click.option("--level", type=str, help="Zilliz index level", required=False),
]
@cli.command()
@click_parameter_decorators_from_typed_dict(ZillizTypedDict)
def ZillizAutoIndex(**parameters: Unpack[ZillizTypedDict]):
from .config import ZillizCloudConfig, AutoIndexConfig
run(
db=DB.ZillizCloud,
db_config=ZillizCloudConfig(
db_label=parameters["db_label"],
uri=SecretStr(parameters["uri"]),
user=parameters["user_name"],
password=SecretStr(parameters["password"]),
),
db_case_config=AutoIndexConfig(
params={parameters["level"]},
),
**parameters,
)
