vectordb-bench 0.0.19__py3-none-any.whl → 0.0.21__py3-none-any.whl

vectordb_bench/backend/clients/__init__.py

@@ -1,12 +1,12 @@
 from enum import Enum
-from typing import Type
+
 from .api import (
-    VectorDB,
-    DBConfig,
     DBCaseConfig,
+    DBConfig,
     EmptyDBCaseConfig,
     IndexType,
     MetricType,
+    VectorDB,
 )
 
 
@@ -41,200 +41,265 @@ class DB(Enum):
     Test = "test"
     AliyunOpenSearch = "AliyunOpenSearch"
 
-
     @property
-    def init_cls(self) -> Type[VectorDB]:
+    def init_cls(self) -> type[VectorDB]:  # noqa: PLR0911, PLR0912, C901
         """Import while in use"""
         if self == DB.Milvus:
             from .milvus.milvus import Milvus
+
             return Milvus
 
         if self == DB.ZillizCloud:
             from .zilliz_cloud.zilliz_cloud import ZillizCloud
+
             return ZillizCloud
 
         if self == DB.Pinecone:
             from .pinecone.pinecone import Pinecone
+
             return Pinecone
 
         if self == DB.ElasticCloud:
             from .elastic_cloud.elastic_cloud import ElasticCloud
+
             return ElasticCloud
 
         if self == DB.QdrantCloud:
             from .qdrant_cloud.qdrant_cloud import QdrantCloud
+
             return QdrantCloud
 
         if self == DB.WeaviateCloud:
             from .weaviate_cloud.weaviate_cloud import WeaviateCloud
+
             return WeaviateCloud
 
         if self == DB.PgVector:
             from .pgvector.pgvector import PgVector
+
             return PgVector
 
         if self == DB.PgVectoRS:
             from .pgvecto_rs.pgvecto_rs import PgVectoRS
+
             return PgVectoRS
-        
+
         if self == DB.PgVectorScale:
             from .pgvectorscale.pgvectorscale import PgVectorScale
+
             return PgVectorScale
 
         if self == DB.PgDiskANN:
             from .pgdiskann.pgdiskann import PgDiskANN
+
             return PgDiskANN
 
         if self == DB.Redis:
             from .redis.redis import Redis
+
             return Redis
-        
+
         if self == DB.MemoryDB:
             from .memorydb.memorydb import MemoryDB
+
             return MemoryDB
 
         if self == DB.Chroma:
             from .chroma.chroma import ChromaClient
+
             return ChromaClient
 
         if self == DB.AWSOpenSearch:
             from .aws_opensearch.aws_opensearch import AWSOpenSearch
+
             return AWSOpenSearch
-        
+
         if self == DB.AlloyDB:
             from .alloydb.alloydb import AlloyDB
+
             return AlloyDB
 
         if self == DB.AliyunElasticsearch:
             from .aliyun_elasticsearch.aliyun_elasticsearch import AliyunElasticsearch
+
             return AliyunElasticsearch
 
         if self == DB.AliyunOpenSearch:
             from .aliyun_opensearch.aliyun_opensearch import AliyunOpenSearch
+
             return AliyunOpenSearch
 
+        if self == DB.Test:
+            from .test.test import Test
+
+            return Test
+
+        msg = f"Unknown DB: {self.name}"
+        raise ValueError(msg)
+
     @property
-    def config_cls(self) -> Type[DBConfig]:
+    def config_cls(self) -> type[DBConfig]:  # noqa: PLR0911, PLR0912, C901
         """Import while in use"""
         if self == DB.Milvus:
             from .milvus.config import MilvusConfig
+
             return MilvusConfig
 
         if self == DB.ZillizCloud:
             from .zilliz_cloud.config import ZillizCloudConfig
+
             return ZillizCloudConfig
 
         if self == DB.Pinecone:
             from .pinecone.config import PineconeConfig
+
             return PineconeConfig
 
         if self == DB.ElasticCloud:
             from .elastic_cloud.config import ElasticCloudConfig
+
             return ElasticCloudConfig
 
         if self == DB.QdrantCloud:
             from .qdrant_cloud.config import QdrantConfig
+
             return QdrantConfig
 
         if self == DB.WeaviateCloud:
             from .weaviate_cloud.config import WeaviateConfig
+
             return WeaviateConfig
 
         if self == DB.PgVector:
             from .pgvector.config import PgVectorConfig
+
             return PgVectorConfig
 
         if self == DB.PgVectoRS:
             from .pgvecto_rs.config import PgVectoRSConfig
+
             return PgVectoRSConfig
 
         if self == DB.PgVectorScale:
             from .pgvectorscale.config import PgVectorScaleConfig
+
             return PgVectorScaleConfig
 
         if self == DB.PgDiskANN:
             from .pgdiskann.config import PgDiskANNConfig
+
             return PgDiskANNConfig
 
         if self == DB.Redis:
             from .redis.config import RedisConfig
+
             return RedisConfig
-        
+
         if self == DB.MemoryDB:
             from .memorydb.config import MemoryDBConfig
+
             return MemoryDBConfig
 
         if self == DB.Chroma:
             from .chroma.config import ChromaConfig
+
             return ChromaConfig
 
         if self == DB.AWSOpenSearch:
             from .aws_opensearch.config import AWSOpenSearchConfig
+
             return AWSOpenSearchConfig
-        
+
         if self == DB.AlloyDB:
             from .alloydb.config import AlloyDBConfig
+
             return AlloyDBConfig
 
         if self == DB.AliyunElasticsearch:
             from .aliyun_elasticsearch.config import AliyunElasticsearchConfig
+
             return AliyunElasticsearchConfig
 
         if self == DB.AliyunOpenSearch:
             from .aliyun_opensearch.config import AliyunOpenSearchConfig
+
             return AliyunOpenSearchConfig
 
-    def case_config_cls(self, index_type: IndexType | None = None) -> Type[DBCaseConfig]:
+        if self == DB.Test:
+            from .test.config import TestConfig
+
+            return TestConfig
+
+        msg = f"Unknown DB: {self.name}"
+        raise ValueError(msg)
+
+    def case_config_cls(  # noqa: PLR0911
+        self,
+        index_type: IndexType | None = None,
+    ) -> type[DBCaseConfig]:
         if self == DB.Milvus:
             from .milvus.config import _milvus_case_config
+
             return _milvus_case_config.get(index_type)
 
         if self == DB.ZillizCloud:
             from .zilliz_cloud.config import AutoIndexConfig
+
             return AutoIndexConfig
 
         if self == DB.ElasticCloud:
             from .elastic_cloud.config import ElasticCloudIndexConfig
+
             return ElasticCloudIndexConfig
 
         if self == DB.QdrantCloud:
             from .qdrant_cloud.config import QdrantIndexConfig
+
             return QdrantIndexConfig
 
         if self == DB.WeaviateCloud:
             from .weaviate_cloud.config import WeaviateIndexConfig
+
             return WeaviateIndexConfig
 
         if self == DB.PgVector:
             from .pgvector.config import _pgvector_case_config
+
             return _pgvector_case_config.get(index_type)
 
         if self == DB.PgVectoRS:
             from .pgvecto_rs.config import _pgvecto_rs_case_config
+
             return _pgvecto_rs_case_config.get(index_type)
 
         if self == DB.AWSOpenSearch:
             from .aws_opensearch.config import AWSOpenSearchIndexConfig
+
             return AWSOpenSearchIndexConfig
 
         if self == DB.PgVectorScale:
             from .pgvectorscale.config import _pgvectorscale_case_config
+
             return _pgvectorscale_case_config.get(index_type)
 
         if self == DB.PgDiskANN:
             from .pgdiskann.config import _pgdiskann_case_config
+
             return _pgdiskann_case_config.get(index_type)
-        
+
         if self == DB.AlloyDB:
             from .alloydb.config import _alloydb_case_config
+
             return _alloydb_case_config.get(index_type)
 
         if self == DB.AliyunElasticsearch:
             from .elastic_cloud.config import ElasticCloudIndexConfig
+
             return ElasticCloudIndexConfig
 
         if self == DB.AliyunOpenSearch:
             from .aliyun_opensearch.config import AliyunOpenSearchIndexConfig
+
             return AliyunOpenSearchIndexConfig
 
         # DB.Pinecone, DB.Chroma, DB.Redis
@@ -242,5 +307,11 @@ class DB(Enum):
 
 
 __all__ = [
-    "DB", "VectorDB", "DBConfig", "DBCaseConfig", "IndexType", "MetricType", "EmptyDBCaseConfig",
+    "DB",
+    "DBCaseConfig",
+    "DBConfig",
+    "EmptyDBCaseConfig",
+    "IndexType",
+    "MetricType",
+    "VectorDB",
 ]

vectordb_bench/backend/clients/aliyun_elasticsearch/aliyun_elasticsearch.py

@@ -1,5 +1,5 @@
-from ..elastic_cloud.elastic_cloud import ElasticCloud
 from ..elastic_cloud.config import ElasticCloudIndexConfig
+from ..elastic_cloud.elastic_cloud import ElasticCloud
 
 
 class AliyunElasticsearch(ElasticCloud):
@@ -24,4 +24,3 @@ class AliyunElasticsearch(ElasticCloud):
             drop_old=drop_old,
             **kwargs,
         )
-

vectordb_bench/backend/clients/aliyun_elasticsearch/config.py

@@ -1,7 +1,6 @@
-from enum import Enum
-from pydantic import SecretStr, BaseModel
+from pydantic import BaseModel, SecretStr
 
-from ..api import DBConfig, DBCaseConfig, MetricType, IndexType
+from ..api import DBConfig
 
 
 class AliyunElasticsearchConfig(DBConfig, BaseModel):
@@ -14,6 +13,6 @@ class AliyunElasticsearchConfig(DBConfig, BaseModel):
 
     def to_dict(self) -> dict:
         return {
-            "hosts": [{'scheme': self.scheme, 'host': self.host, 'port': self.port}],
+            "hosts": [{"scheme": self.scheme, "host": self.host, "port": self.port}],
             "basic_auth": (self.user, self.password.get_secret_value()),
         }