indexing impl

Author: jihan-yin

nucleus/__init__.py

@@ -94,6 +94,7 @@
     JOB_CREATION_TIME_KEY,
     IMAGE_KEY,
     IMAGE_URL_KEY,
+    INDEX_CONTINUOUS_ENABLE_KEY,
     ITEM_METADATA_SCHEMA_KEY,
     ITEMS_KEY,
     KEEP_HISTORY_KEY,
@@ -1205,6 +1206,37 @@ def delete_custom_index(self, dataset_id: str):
             requests_command=requests.delete,
         )
 
+    def toggle_continuous_indexing(self, dataset_id: str, enable: bool = True):
+        """
+        Toggles continuous indexing for a given dataset, which will automatically generate embeddings whenever
+        new images are uploaded. This endpoint is currently locked for most users. Please contact the nucleus team if
+        you wish to use this endpoint.
+
+        :param
+        dataset_id: id of dataset that continuous indexing is being toggled for
+        enable: boolean, sets whether we are enabling or disabling continuous indexing. The default behavior is to enable.
+        """
+        return self.make_request(
+            {INDEX_CONTINUOUS_ENABLE_KEY: enable},
+            f"indexing/{dataset_id}/toggleContinuous",
+            requests_command=requests.post,
+        )
+
+    def kickoff_image_indexing(self, dataset_id: str, enable: bool = True):
+        """
+        Starts indexing generating embeddings for images without embeddings in a given dataset. These embeddings will
+        be used for autotag and similarity search. This endpoint is currently locked for most users.
+        Please contact the nucleus team if you wish to use this endpoint.
+
+        :param
+        dataset_id: id of dataset for generating embeddings on.
+        """
+        return self.make_request(
+            {},
+            f"indexing/{dataset_id}/kickoffImage",
+            requests_command=requests.post,
+        )
+
     def make_request(
         self, payload: dict, route: str, requests_command=requests.post
     ) -> dict:

nucleus/constants.py

@@ -37,6 +37,7 @@
 IMAGE_KEY = "image"
 IMAGE_URL_KEY = "image_url"
 INDEX_KEY = "index"
+INDEX_CONTINUOUS_ENABLE_KEY = "enable"
 ITEMS_KEY = "items"
 ITEM_ID_KEY = "item_id"
 ITEM_KEY = "item"

nucleus/dataset.py

@@ -376,6 +376,12 @@ def create_custom_index(self, embeddings_urls: list, embedding_dim: int):
     def delete_custom_index(self):
         return self._client.delete_custom_index(self.id)
 
+    def toggle_continuous_indexing(self, enable: bool = True):
+        return self._client.toggle_continuous_indexing(self.id, enable)
+
+    def kickoff_image_indexing(self):
+        return self._client.kickoff_image_indexing(self.id)
+
     def check_index_status(self, job_id: str):
         return self._client.check_index_status(job_id)