fix saving engine names

qdrant · Mar 16, 2024 · 5f8ebee · 5f8ebee
1 parent f112b6b
commit 5f8ebee
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 2 deletions.
diff --git a/engine/base_client/client.py b/engine/base_client/client.py
@@ -20,6 +20,7 @@ class BaseClient:
     def __init__(
         self,
         name: str,  # name of the experiment
+        engine: str,  # name of the engine
         configurator: BaseConfigurator,
         uploader: BaseUploader,
         searchers: List[BaseSearcher],
@@ -28,6 +29,7 @@ def __init__(
         self.configurator = configurator
         self.uploader = uploader
         self.searchers = searchers
+        self.engine = engine
 
     def save_search_results(
         self, dataset_name: str, results: dict, search_id: int, search_params: dict
@@ -40,7 +42,15 @@ def save_search_results(
         result_path = RESULTS_DIR / experiments_file
         with open(result_path, "w") as out:
             out.write(
-                json.dumps({"params": search_params, "results": results}, indent=2)
+                json.dumps({
+                    "params": {
+                        "dataset": dataset_name,
+                        "experiment": self.name,
+                        "engine": self.engine,
+                        **search_params
+                    },
+                    "results": results
+                }, indent=2)
             )
         return result_path
 
@@ -53,7 +63,8 @@ def save_upload_results(
         with open(RESULTS_DIR / experiments_file, "w") as out:
             upload_stats = {
                 "params": {
-                    "engine": self.name,
+                    "experiment": self.name,
+                    "engine": self.engine,
                     "dataset": dataset_name,
                     **upload_params
                 },

diff --git a/engine/clients/client_factory.py b/engine/clients/client_factory.py
@@ -65,8 +65,10 @@
 class ClientFactory(ABC):
     def __init__(self, host):
         self.host = host
+        self.engine = None
 
     def _create_configurator(self, experiment) -> BaseConfigurator:
+        self.engine = experiment["engine"]
         engine_configurator_class = ENGINE_CONFIGURATORS[experiment["engine"]]
         engine_configurator = engine_configurator_class(
             self.host,
@@ -103,6 +105,7 @@ def _create_searchers(self, experiment) -> List[BaseSearcher]:
     def build_client(self, experiment):
         return BaseClient(
             name=experiment["name"],
+            engine=experiment["engine"],
             configurator=self._create_configurator(experiment),
             uploader=self._create_uploader(experiment),
             searchers=self._create_searchers(experiment),