Feat/endee version

pyproject.toml

@@ -79,7 +79,7 @@ all = [
     "mysql-connector-python",
     "turbopuffer[fast]",
     'zvec',
-    "endee",
+    "endee==0.1.10", # compatible with pydantic<2
 ]
 
 qdrant          = [ "qdrant-client" ]
@@ -109,7 +109,7 @@ alisql          = [ "mysql-connector-python" ]
 doris           = [ "doris-vector-search" ]
 turbopuffer     = [ "turbopuffer" ]
 zvec            = [ "zvec" ]
-endee           = [ "endee" ]
+endee           = [ "endee==0.1.10" ]
 
 [project.urls]
 Repository  = "https://github.com/zilliztech/VectorDBBench"

vectordb_bench/backend/clients/api.py

@@ -1,13 +1,13 @@
 from abc import ABC, abstractmethod
 from contextlib import contextmanager
-from enum import Enum
+from enum import StrEnum
 
 from pydantic import BaseModel, SecretStr, validator
 
 from vectordb_bench.backend.filter import Filter, FilterOp
 
 
-class MetricType(str, Enum):
+class MetricType(StrEnum):
     L2 = "L2"
     COSINE = "COSINE"
     IP = "IP"
@@ -16,7 +16,7 @@ class MetricType(str, Enum):
     JACCARD = "JACCARD"
 
 
-class IndexType(str, Enum):
+class IndexType(StrEnum):
     HNSW = "HNSW"
     HNSW_SQ = "HNSW_SQ"
     HNSW_BQ = "HNSW_BQ"
@@ -47,7 +47,7 @@ class IndexType(str, Enum):
     NONE = "NONE"
 
 
-class SQType(str, Enum):
+class SQType(StrEnum):
     SQ6 = "SQ6"
     SQ8 = "SQ8"
     BF16 = "BF16"

vectordb_bench/backend/clients/elastic_cloud/config.py

@@ -1,4 +1,4 @@
-from enum import Enum
+from enum import StrEnum
 
 from pydantic import BaseModel, SecretStr
 
@@ -16,7 +16,7 @@ def to_dict(self) -> dict:
         }
 
 
-class ESElementType(str, Enum):
+class ESElementType(StrEnum):
     float = "float"  # 4 byte
     byte = "byte"  # 1 byte, -128 to 127
 

vectordb_bench/backend/clients/endee/cli.py

@@ -38,7 +38,7 @@ class EndeeTypedDict(CommonTypedDict):
         str,
         click.option(
             "--space-type",
-            type=click.Choice(["cosine", "l2", "dot_product"]),
+            type=click.Choice(["cosine", "l2", "ip"]),
             default="cosine",
             help="Distance metric",
             show_default=True,

vectordb_bench/backend/clients/endee/config.py

@@ -5,7 +5,7 @@
 
 class EndeeConfig(DBConfig):
     token: SecretStr | None = None
-    region: str | None = "as1"
+    region: str | None = ""
     base_url: str = "http://127.0.0.1:8080/api/v1"
     space_type: str = "cosine"
     precision: str = "int8d"

vectordb_bench/backend/clients/endee/endee.py

@@ -34,7 +34,7 @@ def __init__(
         **kwargs,
     ):
         self.token = db_config.get("token", "")
-        self.region = db_config.get("region", "as1")
+        self.region = db_config.get("region", "")
         self.base_url = db_config.get("base_url")
 
         self.collection_name = db_config.get("collection_name") or db_config.get("index_name")

vectordb_bench/backend/clients/tencent_elasticsearch/config.py

@@ -1,4 +1,4 @@
-from enum import Enum
+from enum import StrEnum
 
 from pydantic import BaseModel, SecretStr
 
@@ -20,7 +20,7 @@ def to_dict(self) -> dict:
         }
 
 
-class ESElementType(str, Enum):
+class ESElementType(StrEnum):
     float = "float"  # 4 byte
     byte = "byte"  # 1 byte, -128 to 127