fix lint

Author: euphoria0-0

install.py

@@ -4,70 +4,52 @@
 
 
 def docker_tag_base():
-    return 'vdbbench'
+    return "vdbbench"
+
 
 def dockerfile_path_base():
-    return os.path.join('vectordb_bench/', '../Dockerfile')
+    return os.path.join("vectordb_bench/", "../Dockerfile")
+
 
 def docker_tag(track, algo):
-    return docker_tag_base() + '-' + track + '-' + algo
+    return docker_tag_base() + "-" + track + "-" + algo
 
 
 def build(tag, args, dockerfile):
-    print('Building %s...' % tag)
+    print("Building %s..." % tag)
     if args is not None and len(args) != 0:
         q = " ".join(["--build-arg " + x.replace(" ", "\\ ") for x in args])
     else:
         q = ""
 
     try:
-        command = 'docker build %s --rm -t %s -f' \
-                   % (q, tag)
-        command += ' %s .' % dockerfile
+        command = "docker build %s --rm -t %s -f" % (q, tag)
+        command += " %s ." % dockerfile
         print(command)
         subprocess.check_call(command, shell=True)
-        return {tag: 'success'}
+        return {tag: "success"}
     except subprocess.CalledProcessError:
-        return {tag: 'fail'}
+        return {tag: "fail"}
+
 
 def build_multiprocess(args):
     return build(*args)
 
 
 if __name__ == "__main__":
-    parser = argparse.ArgumentParser(
-        formatter_class=argparse.ArgumentDefaultsHelpFormatter)
-    parser.add_argument(
-        "--proc",
-        default=1,
-        type=int,
-        help="the number of process to build docker images")
+    parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
+    parser.add_argument("--proc", default=1, type=int, help="the number of process to build docker images")
+    parser.add_argument("--track", choices=["none"], default="none")
+    parser.add_argument("--algorithm", metavar="NAME", help="build only the named algorithm image", default=None)
     parser.add_argument(
-        '--track',
-        choices=['none'],
-        default='none'
+        "--dockerfile", metavar="PATH", help="build only the image from a Dockerfile path", default=None
     )
-    parser.add_argument(
-        '--algorithm',
-        metavar='NAME',
-        help='build only the named algorithm image',
-        default=None)
-    parser.add_argument(
-        '--dockerfile',
-        metavar='PATH',
-        help='build only the image from a Dockerfile path',
-        default=None)
-    parser.add_argument(
-        '--build-arg',
-        help='pass given args to all docker builds',
-        nargs="+")
+    parser.add_argument("--build-arg", help="pass given args to all docker builds", nargs="+")
     args = parser.parse_args()
 
-    print('Building base image...')
-
-    subprocess.check_call(
-        'docker build \
-        --rm -t %s -f %s .' % (docker_tag_base(), dockerfile_path_base()), shell=True)
+    print("Building base image...")
 
-    print('Building end.')
+    subprocess.check_call("docker build \
+        --rm -t %s -f %s ." % (docker_tag_base(), dockerfile_path_base()), shell=True)
 
+    print("Building end.")

tests/conftest.py

@@ -1,4 +1,5 @@
 import sys
 
 from os.path import dirname, abspath
+
 sys.path.append(dirname(dirname(abspath(__file__))))

tests/test_bench_runner.py

@@ -2,11 +2,16 @@
 import logging
 from vectordb_bench.interface import BenchMarkRunner
 from vectordb_bench.models import (
-    DB, IndexType, CaseType, TaskConfig, CaseConfig,
+    DB,
+    IndexType,
+    CaseType,
+    TaskConfig,
+    CaseConfig,
 )
 
 log = logging.getLogger(__name__)
 
+
 class TestBenchRunner:
     def test_get_results(self):
         runner = BenchMarkRunner()
@@ -17,7 +22,7 @@ def test_get_results(self):
     def test_performance_case_whole(self):
         runner = BenchMarkRunner()
 
-        task_config=TaskConfig(
+        task_config = TaskConfig(
             db=DB.Milvus,
             db_config=DB.Milvus.config(),
             db_case_config=DB.Milvus.case_config_cls(index=IndexType.Flat)(),
@@ -32,7 +37,7 @@ def test_performance_case_whole(self):
     def test_performance_case_clean(self):
         runner = BenchMarkRunner()
 
-        task_config=TaskConfig(
+        task_config = TaskConfig(
             db=DB.Milvus,
             db_config=DB.Milvus.config(),
             db_case_config=DB.Milvus.case_config_cls(index=IndexType.Flat)(),
@@ -44,17 +49,18 @@ def test_performance_case_clean(self):
         runner.stop_running()
 
     def test_performance_case_no_error(self):
-        task_config=TaskConfig(
+        task_config = TaskConfig(
             db=DB.ZillizCloud,
             db_config=DB.ZillizCloud.config(uri="xxx", user="abc", password="1234"),
             db_case_config=DB.ZillizCloud.case_config_cls()(),
             case_config=CaseConfig(case_id=CaseType.PerformanceSZero),
         )
 
         t = task_config.copy()
-        d = t.json(exclude={'db_config': {'password', 'api_key'}})
+        d = t.json(exclude={"db_config": {"password", "api_key"}})
         log.info(f"{d}")
 
         import ujson
+
         loads = ujson.loads(d)
         log.info(f"{loads}")

tests/test_chroma.py

@@ -6,7 +6,6 @@
 import numpy as np
 import chromadb
 
-
 log = logging.getLogger(__name__)
 
 """ Tests for Chroma, assumes Chroma is running on localhost:8000, 
@@ -19,13 +18,11 @@
     3. default port is 8000, default host is localhost"""
 
 
-
-dict = {} #Assumes chroma is acception connections on localhost:8000
-dict['name'] = "chroma"
-dict['host'] = "localhost"
-dict['port'] = 8000
-dict['password'] = "chroma"
-
+dict = {}  # Assumes chroma is acception connections on localhost:8000
+dict["name"] = "chroma"
+dict["host"] = "localhost"
+dict["port"] = 8000
+dict["password"] = "chroma"
 
 
 class TestChroma:
@@ -34,7 +31,6 @@ def test_insert_and_search(self):
 
         dbcls = DB.Chroma.init_cls
         dbConfig = DB.Chroma.config_cls
-        
 
         dim = 16
         chrma = dbcls(
@@ -49,50 +45,37 @@ def test_insert_and_search(self):
         filter_value = 0.9
         embeddings = [[np.random.random() for _ in range(dim)] for _ in range(count)]
 
-
         # insert
         with chrma.init():
-            #chrma.client.delete_collection("example2")
-            assert (chrma.client.heartbeat() is not None), "chroma client is not connected"
+            # chrma.client.delete_collection("example2")
+            assert chrma.client.heartbeat() is not None, "chroma client is not connected"
             res = chrma.insert_embeddings(embeddings=embeddings, metadata=range(count))
             # bulk_insert return
-            assert (
-                res[0] == count
-            ), f"the return count of bulk insert ({res}) is not equal to count ({count})"
+            assert res[0] == count, f"the return count of bulk insert ({res}) is not equal to count ({count})"
 
             # count entries in chroma database
             countRes = chrma.collection.count()
-            
-            assert (
-                countRes == count
-            ), f"the return count of redis client ({countRes}) is not equal to count ({count})"
+
+            assert countRes == count, f"the return count of redis client ({countRes}) is not equal to count ({count})"
 
         # search
         with chrma.init():
             test_id = np.random.randint(count)
-            #log.info(f"test_id: {test_id}")
+            # log.info(f"test_id: {test_id}")
             q = embeddings[test_id]
 
             res = chrma.search_embedding(query=q, k=100)
             print(res)
-            assert (
-                res[0] == int(test_id)
-            ), f"the most nearest neighbor ({res[0]}) id is not test_id ({int(test_id)}"
-            
+            assert res[0] == int(test_id), f"the most nearest neighbor ({res[0]}) id is not test_id ({int(test_id)}"
 
         # search with filters, assumes filter format {id: int, metadata: >=int}
         with chrma.init():
             filter_value = int(count * filter_value)
             test_id = np.random.randint(filter_value, count)
             q = embeddings[test_id]
 
-
-            res = chrma.search_embedding(
-                query=q, k=100, filters={"id": filter_value}
-            )
-            assert (
-                res[0] == int(test_id)
-            ), f"the most nearest neighbor ({res[0]}) id is not test_id ({test_id})"
+            res = chrma.search_embedding(query=q, k=100, filters={"id": filter_value})
+            assert res[0] == int(test_id), f"the most nearest neighbor ({res[0]}) id is not test_id ({test_id})"
             isFilter = True
             id_list = []
             for id in res:
@@ -101,5 +84,3 @@ def test_insert_and_search(self):
                     isFilter = False
                     break
             assert isFilter, f"Filter not working, id_list: {id_list}"
-
-            

tests/test_data_source.py

@@ -5,14 +5,12 @@
 
 log = logging.getLogger("vectordb_bench")
 
+
 class TestReader:
-    @pytest.mark.parametrize("type_case", [
-        (k, v) for k, v in type2case.items()
-    ])
+    @pytest.mark.parametrize("type_case", [(k, v) for k, v in type2case.items()])
     def test_type_cases(self, type_case):
         self.per_case_test(type_case)
 
-
     def per_case_test(self, type_case):
         t, ca_cls = type_case
         ca = ca_cls()

tests/test_dataset.py

@@ -4,9 +4,9 @@
 from pydantic import ValidationError
 from vectordb_bench.backend.data_source import DatasetSource
 
-
 log = logging.getLogger("vectordb_bench")
 
+
 class TestDataSet:
     def test_iter_dataset(self):
         for ds in Dataset:
@@ -29,6 +29,7 @@ def test_iter_cohere(self):
         cohere_10m.prepare()
 
         import time
+
         before = time.time()
         for i in cohere_10m:
             log.debug(i.head(1))
@@ -40,9 +41,11 @@ def test_iter_cohere(self):
     def test_iter_laion(self):
         laion_100m = Dataset.LAION.manager(100_000_000)
         from vectordb_bench.backend.data_source import DatasetSource
+
         laion_100m.prepare(source=DatasetSource.AliyunOSS)
 
         import time
+
         before = time.time()
         for i in laion_100m:
             log.debug(i.head(1))
@@ -74,4 +77,3 @@ def test_download_small(self):
             files=files,
             local_ds_root=openai_50k.data_dir,
         )
-

tests/test_elasticsearch_cloud.py

@@ -6,7 +6,6 @@
 )
 import numpy as np
 
-
 log = logging.getLogger(__name__)
 
 cloud_id = ""
@@ -41,9 +40,7 @@ def test_insert_and_search(self):
         with es.init():
             res = es.insert_embeddings(embeddings=embeddings, metadata=range(count))
             # bulk_insert return
-            assert (
-                res == count
-            ), f"the return count of bulk insert ({res}) is not equal to count ({count})"
+            assert res == count, f"the return count of bulk insert ({res}) is not equal to count ({count})"
 
             # indice_count return
             es.client.indices.refresh()
@@ -61,23 +58,17 @@ def test_insert_and_search(self):
 
             res = es.search_embedding(query=q, k=100)
             log.info(f"search_results_id: {res}")
-            assert (
-                res[0] == test_id
-            ), f"the most nearest neighbor ({res[0]}) id is not test_id ({test_id})"
+            assert res[0] == test_id, f"the most nearest neighbor ({res[0]}) id is not test_id ({test_id})"
 
         # search with filters
         with es.init():
             test_id = np.random.randint(count * filter_rate, count)
             log.info(f"test_id: {test_id}")
             q = embeddings[test_id]
 
-            res = es.search_embedding(
-                query=q, k=100, filters={"id": count * filter_rate}
-            )
+            res = es.search_embedding(query=q, k=100, filters={"id": count * filter_rate})
             log.info(f"search_results_id: {res}")
-            assert (
-                res[0] == test_id
-            ), f"the most nearest neighbor ({res[0]}) id is not test_id ({test_id})"
+            assert res[0] == test_id, f"the most nearest neighbor ({res[0]}) id is not test_id ({test_id})"
             isFilter = True
             for id in res:
                 if id < count * filter_rate:

tests/test_models.py

@@ -1,18 +1,10 @@
 import pytest
 import logging
-from vectordb_bench.models import (
-    TaskConfig, CaseConfig,
-    CaseResult, TestResult,
-    Metric, CaseType
-)
-from vectordb_bench.backend.clients import (
-    DB,
-    IndexType
-)
+from vectordb_bench.models import TaskConfig, CaseConfig, CaseResult, TestResult, Metric, CaseType
+from vectordb_bench.backend.clients import DB, IndexType
 
 from vectordb_bench import config
 
-
 log = logging.getLogger("vectordb_bench")
 
 
@@ -33,7 +25,7 @@ def test_test_result(self):
         test_result.flush()
 
         with pytest.raises(ValueError):
-            result = TestResult.read_file('nosuchfile.json')
+            result = TestResult.read_file("nosuchfile.json")
 
     def test_test_result_read_write(self):
         result_dir = config.RESULTS_LOCAL_DIR