vectordb-bench 0.0.29__py3-none-any.whl → 1.0.0__py3-none-any.whl

vectordb_bench/backend/clients/qdrant_local/cli.py

@@ -0,0 +1,60 @@
+from typing import Annotated, Unpack
+
+import click
+from pydantic import SecretStr
+
+from vectordb_bench.backend.clients import DB
+from vectordb_bench.cli.cli import (
+    CommonTypedDict,
+    cli,
+    click_parameter_decorators_from_typed_dict,
+    run,
+)
+
+DBTYPE = DB.QdrantLocal
+
+
+class QdrantLocalTypedDict(CommonTypedDict):
+    url: Annotated[
+        str,
+        click.option("--url", type=str, help="Qdrant url", required=True),
+    ]
+    on_disk: Annotated[
+        bool,
+        click.option("--on-disk", type=bool, default=False, help="Store the vectors and the HNSW index on disk"),
+    ]
+    m: Annotated[
+        int,
+        click.option("--m", type=int, default=16, help="HNSW index parameter m, set 0 to disable the index"),
+    ]
+    ef_construct: Annotated[
+        int,
+        click.option("--ef-construct", type=int, default=200, help="HNSW index parameter ef_construct"),
+    ]
+    hnsw_ef: Annotated[
+        int,
+        click.option(
+            "--hnsw-ef",
+            type=int,
+            default=0,
+            help="HNSW index parameter hnsw_ef, set 0 to use ef_construct for search",
+        ),
+    ]
+
+
+@cli.command()
+@click_parameter_decorators_from_typed_dict(QdrantLocalTypedDict)
+def QdrantLocal(**parameters: Unpack[QdrantLocalTypedDict]):
+    from .config import QdrantLocalConfig, QdrantLocalIndexConfig
+
+    run(
+        db=DBTYPE,
+        db_config=QdrantLocalConfig(url=SecretStr(parameters["url"])),
+        db_case_config=QdrantLocalIndexConfig(
+            on_disk=parameters["on_disk"],
+            m=parameters["m"],
+            ef_construct=parameters["ef_construct"],
+            hnsw_ef=parameters["hnsw_ef"],
+        ),
+        **parameters,
+    )

vectordb_bench/backend/clients/qdrant_local/config.py

@@ -0,0 +1,47 @@
+from pydantic import BaseModel, SecretStr
+
+from ..api import DBCaseConfig, DBConfig, MetricType
+
+
+class QdrantLocalConfig(DBConfig):
+    url: SecretStr
+
+    def to_dict(self) -> dict:
+        return {
+            "url": self.url.get_secret_value(),
+        }
+
+
+class QdrantLocalIndexConfig(BaseModel, DBCaseConfig):
+    metric_type: MetricType | None = None
+    m: int
+    ef_construct: int
+    hnsw_ef: int | None = 0
+    on_disk: bool | None = False
+
+    def parse_metric(self) -> str:
+        if self.metric_type == MetricType.L2:
+            return "Euclid"
+
+        if self.metric_type == MetricType.IP:
+            return "Dot"
+
+        return "Cosine"
+
+    def index_param(self) -> dict:
+        return {
+            "distance": self.parse_metric(),
+            "m": self.m,
+            "ef_construct": self.ef_construct,
+            "on_disk": self.on_disk,
+        }
+
+    def search_param(self) -> dict:
+        search_params = {
+            "exact": False,  # Force to use ANNs
+        }
+
+        if self.hnsw_ef != 0:
+            search_params["hnsw_ef"] = self.hnsw_ef
+
+        return search_params

vectordb_bench/backend/clients/qdrant_local/qdrant_local.py

@@ -0,0 +1,232 @@
+"""Wrapper around the Qdrant over VectorDB"""
+
+import logging
+import time
+from collections.abc import Iterable
+from contextlib import contextmanager
+
+from qdrant_client import QdrantClient
+from qdrant_client.http.models import (
+    Batch,
+    CollectionStatus,
+    FieldCondition,
+    Filter,
+    HnswConfigDiff,
+    OptimizersConfigDiff,
+    PayloadSchemaType,
+    Range,
+    SearchParams,
+    VectorParams,
+)
+
+from ..api import VectorDB
+from .config import QdrantLocalIndexConfig
+
+log = logging.getLogger(__name__)
+
+SECONDS_WAITING_FOR_INDEXING_API_CALL = 5
+QDRANT_BATCH_SIZE = 100
+
+
+def qdrant_collection_exists(client: QdrantClient, collection_name: str) -> bool:
+    collection_exists = True
+
+    try:
+        client.get_collection(collection_name)
+    except Exception:
+        collection_exists = False
+
+    return collection_exists
+
+
+class QdrantLocal(VectorDB):
+    def __init__(
+        self,
+        dim: int,
+        db_config: dict,
+        db_case_config: QdrantLocalIndexConfig,
+        collection_name: str = "QdrantLocalCollection",
+        drop_old: bool = False,
+        name: str = "QdrantLocal",
+        **kwargs,
+    ):
+        """Initialize wrapper around the qdrant."""
+        self.name = name
+        self.db_config = db_config
+        self.case_config = db_case_config
+        self.search_parameter = self.case_config.search_param()
+        self.collection_name = collection_name
+        self.client = None
+
+        self._primary_field = "pk"
+        self._vector_field = "vector"
+
+        client = QdrantClient(**self.db_config)
+
+        # Lets just print the parameters here for double check
+        log.info(f"Case config: {self.case_config.index_param()}")
+        log.info(f"Search parameter: {self.search_parameter}")
+
+        if drop_old and qdrant_collection_exists(client, self.collection_name):
+            log.info(f"{self.name} client drop_old collection: {self.collection_name}")
+            client.delete_collection(self.collection_name)
+
+        if not qdrant_collection_exists(client, self.collection_name):
+            log.info(f"{self.name} create collection: {self.collection_name}")
+            self._create_collection(dim, client)
+
+        client = None
+
+    @contextmanager
+    def init(self):
+        """
+        Examples:
+            >>> with self.init():
+            >>>     self.insert_embeddings()
+            >>>     self.search_embedding()
+        """
+        # create connection
+        self.client = QdrantClient(**self.db_config)
+        yield
+        self.client = None
+        del self.client
+
+    def _create_collection(self, dim: int, qdrant_client: QdrantClient):
+        log.info(f"Create collection: {self.collection_name}")
+        log.info(
+            f"Index parameters: m={self.case_config.index_param()['m']}, "
+            f"ef_construct={self.case_config.index_param()['ef_construct']}, "
+            f"on_disk={self.case_config.index_param()['on_disk']}"
+        )
+
+        # If the on_disk is true, we enable both on disk index and vectors.
+        try:
+            qdrant_client.create_collection(
+                collection_name=self.collection_name,
+                vectors_config=VectorParams(
+                    size=dim,
+                    distance=self.case_config.index_param()["distance"],
+                    on_disk=self.case_config.index_param()["on_disk"],
+                ),
+                hnsw_config=HnswConfigDiff(
+                    m=self.case_config.index_param()["m"],
+                    ef_construct=self.case_config.index_param()["ef_construct"],
+                    on_disk=self.case_config.index_param()["on_disk"],
+                ),
+            )
+
+            qdrant_client.create_payload_index(
+                collection_name=self.collection_name,
+                field_name=self._primary_field,
+                field_schema=PayloadSchemaType.INTEGER,
+            )
+
+        except Exception as e:
+            if "already exists!" in str(e):
+                return
+            log.warning(f"Failed to create collection: {self.collection_name} error: {e}")
+            raise e from None
+
+    def optimize(self, data_size: int | None = None):
+        assert self.client, "Please call self.init() before"
+        # wait for vectors to be fully indexed
+        try:
+            while True:
+                info = self.client.get_collection(self.collection_name)
+                time.sleep(SECONDS_WAITING_FOR_INDEXING_API_CALL)
+                if info.status != CollectionStatus.GREEN:
+                    continue
+                if info.status == CollectionStatus.GREEN:
+                    log.info(f"Finishing building index for collection: {self.collection_name}")
+                    msg = (
+                        f"Stored vectors: {info.vectors_count}, Indexed vectors: {info.indexed_vectors_count}, "
+                        f"Collection status: {info.indexed_vectors_count}"
+                    )
+                    log.info(msg)
+                    return
+
+        except Exception as e:
+            log.warning(f"QdrantCloud ready to search error: {e}")
+            raise e from None
+
+    def insert_embeddings(
+        self,
+        embeddings: Iterable[list[float]],
+        metadata: list[int],
+        **kwargs,
+    ) -> tuple[int, Exception]:
+        """Insert embeddings into the database.
+
+        Args:
+            embeddings(list[list[float]]): list of embeddings
+            metadata(list[int]): list of metadata
+            kwargs: other arguments
+
+        Returns:
+            tuple[int, Exception]: number of embeddings inserted and exception if any
+        """
+        assert self.client is not None
+        assert len(embeddings) == len(metadata)
+        insert_count = 0
+
+        # disable indexing for quick insertion
+        self.client.update_collection(
+            collection_name=self.collection_name,
+            optimizer_config=OptimizersConfigDiff(indexing_threshold=0),
+        )
+        try:
+            for offset in range(0, len(embeddings), QDRANT_BATCH_SIZE):
+                vectors = embeddings[offset : offset + QDRANT_BATCH_SIZE]
+                ids = metadata[offset : offset + QDRANT_BATCH_SIZE]
+                payloads = [{self._primary_field: v} for v in ids]
+                _ = self.client.upsert(
+                    collection_name=self.collection_name,
+                    wait=True,
+                    points=Batch(ids=ids, payloads=payloads, vectors=vectors),
+                )
+                insert_count += QDRANT_BATCH_SIZE
+            # enable indexing after insertion
+            self.client.update_collection(
+                collection_name=self.collection_name,
+                optimizer_config=OptimizersConfigDiff(indexing_threshold=100),
+            )
+
+        except Exception as e:
+            log.info(f"Failed to insert data, {e}")
+            return insert_count, e
+        else:
+            return insert_count, None
+
+    def search_embedding(
+        self,
+        query: list[float],
+        k: int = 100,
+        filters: dict | None = None,
+        timeout: int | None = None,
+    ) -> list[int]:
+        """Perform a search on a query embedding and return results with score.
+        Should call self.init() first.
+        """
+        assert self.client is not None
+
+        f = None
+        if filters:
+            f = Filter(
+                must=[
+                    FieldCondition(
+                        key=self._primary_field,
+                        range=Range(
+                            gt=filters.get("id"),
+                        ),
+                    ),
+                ],
+            )
+        res = self.client.query_points(
+            collection_name=self.collection_name,
+            query=query,
+            limit=k,
+            query_filter=f,
+            search_params=SearchParams(**self.search_parameter),
+        ).points
+
+        return [result.id for result in res]

vectordb_bench/backend/clients/weaviate_cloud/cli.py

@@ -15,12 +15,33 @@ from .. import DB
 class WeaviateTypedDict(CommonTypedDict):
     api_key: Annotated[
         str,
-        click.option("--api-key", type=str, help="Weaviate api key", required=True),
+        click.option("--api-key", type=str, help="Weaviate api key", required=False, default=""),
     ]
     url: Annotated[
         str,
         click.option("--url", type=str, help="Weaviate url", required=True),
     ]
+    no_auth: Annotated[
+        bool,
+        click.option(
+            "--no-auth",
+            is_flag=True,
+            help="Do not use api-key, set it to true if you are using a local setup. Default is False.",
+            default=False,
+        ),
+    ]
+    m: Annotated[
+        int,
+        click.option("--m", type=int, default=16, help="HNSW index parameter m."),
+    ]
+    ef_construct: Annotated[
+        int,
+        click.option("--ef-construction", type=int, default=256, help="HNSW index parameter ef_construction"),
+    ]
+    ef: Annotated[
+        int,
+        click.option("--ef", type=int, default=256, help="HNSW index parameter ef for search"),
+    ]
 
 
 @cli.command()
@@ -32,9 +53,14 @@ def Weaviate(**parameters: Unpack[WeaviateTypedDict]):
         db=DB.WeaviateCloud,
         db_config=WeaviateConfig(
             db_label=parameters["db_label"],
-            api_key=SecretStr(parameters["api_key"]),
+            api_key=SecretStr(parameters["api_key"]) if parameters["api_key"] != "" else SecretStr("-"),
             url=SecretStr(parameters["url"]),
+            no_auth=parameters["no_auth"],
+        ),
+        db_case_config=WeaviateIndexConfig(
+            efConstruction=parameters["ef_construction"],
+            maxConnections=parameters["m"],
+            ef=parameters["ef"],
         ),
-        db_case_config=WeaviateIndexConfig(ef=256, efConstruction=256, maxConnections=16),
         **parameters,
     )

vectordb_bench/backend/clients/weaviate_cloud/config.py

@@ -6,11 +6,13 @@ from ..api import DBCaseConfig, DBConfig, MetricType
 class WeaviateConfig(DBConfig):
     url: SecretStr
     api_key: SecretStr
+    no_auth: bool | None = False
 
     def to_dict(self) -> dict:
         return {
             "url": self.url.get_secret_value(),
             "auth_client_secret": self.api_key.get_secret_value(),
+            "no_auth": self.no_auth,
         }
 
 

vectordb_bench/backend/clients/weaviate_cloud/weaviate_cloud.py

@@ -38,6 +38,11 @@ class WeaviateCloud(VectorDB):
         self._vector_field = "vector"
         self._index_name = "vector_idx"
 
+        # If local setup is used, we
+        if db_config["no_auth"]:
+            del db_config["auth_client_secret"]
+        del db_config["no_auth"]
+
         from weaviate import Client
 
         client = Client(**db_config)