working (#283)

Author: jihan-yin

nucleus/__init__.py

@@ -742,14 +742,23 @@ def create_custom_index(
         )
 
     @deprecated("Prefer calling Dataset.delete_custom_index instead.")
-    def delete_custom_index(self, dataset_id: str):
+    def delete_custom_index(self, dataset_id: str, image: bool):
         # TODO: deprecate in favor of Dataset.delete_custom_index invocation
         return self.make_request(
-            {},
+            {"image": image},
             f"indexing/{dataset_id}",
             requests_command=requests.delete,
         )
 
+    @deprecated("Prefer calling Dataset.set_primary_index instead.")
+    def set_primary_index(self, dataset_id: str, image: bool, custom: bool):
+        # TODO: deprecate in favor of Dataset.set_primary_index invocation
+        return self.make_request(
+            {"image": image, "custom": custom},
+            f"indexing/{dataset_id}/setPrimary",
+            requests_command=requests.post,
+        )
+
     @deprecated("Prefer calling Dataset.set_continuous_indexing instead.")
     def set_continuous_indexing(self, dataset_id: str, enable: bool = True):
         # TODO: deprecate in favor of Dataset.set_continuous_indexing invocation

nucleus/dataset.py

@@ -944,7 +944,7 @@ def create_custom_index(
             self._client,
         )
 
-    def delete_custom_index(self):
+    def delete_custom_index(self, image: bool = True):
         """Deletes the custom index uploaded to the dataset.
 
         Returns:
@@ -956,7 +956,18 @@ def delete_custom_index(self):
                     "message": str
                 }
         """
-        return self._client.delete_custom_index(self.id)
+        return self._client.delete_custom_index(self.id, image)
+
+    def set_primary_index(self, image: bool = True, custom: bool = False):
+        """Sets the primary index used for Autotag and Similarity Search on this dataset.
+
+        Returns:
+
+            {
+                "success": bool,
+            }
+        """
+        return self._client.set_primary_index(self.id, image, custom)
 
     def set_continuous_indexing(self, enable: bool = True):
         """Toggle whether embeddings are automatically generated for new data.
@@ -985,9 +996,10 @@ def set_continuous_indexing(self, enable: bool = True):
             MESSAGE_KEY: preprocessed_response[MESSAGE_KEY],
         }
         if enable:
-            response[BACKFILL_JOB_KEY] = (
-                AsyncJob.from_json(preprocessed_response, self._client),
+            response[BACKFILL_JOB_KEY] = AsyncJob.from_json(
+                preprocessed_response, self._client
             )
+
         return response
 
     def create_image_index(self):

tests/test_indexing.py

@@ -2,6 +2,7 @@
 
 from nucleus import DatasetItem
 from nucleus.constants import (
+    BACKFILL_JOB_KEY,
     ERROR_PAYLOAD,
     JOB_ID_KEY,
     MESSAGE_KEY,
@@ -38,7 +39,31 @@ def dataset(CLIENT):
 
 
 @pytest.mark.integration
-def test_index_integration(dataset):
+def test_set_continuous_indexing(dataset):
+    resp = dataset.set_continuous_indexing(True)
+    job = resp[BACKFILL_JOB_KEY]
+    print(job)
+    assert job
+    assert job.job_id
+    assert job.job_last_known_status
+    assert job.job_type
+    assert job.job_creation_time
+
+    job_status_response = job.status()
+    assert STATUS_KEY in job_status_response
+    assert JOB_ID_KEY in job_status_response
+    assert MESSAGE_KEY in job_status_response
+
+
+@pytest.mark.integration
+def test_set_primary_index(dataset):
+    dataset.set_continuous_indexing()
+    resp = dataset.set_primary_index(image=True, custom=False)
+    assert resp["success"]
+
+
+@pytest.mark.integration
+def test_create_custom_index(dataset):
     signed_embeddings_url = TEST_INDEX_EMBEDDINGS_FILE
     job = dataset.create_custom_index([signed_embeddings_url], embedding_dim=3)
     assert job.job_id
@@ -52,6 +77,21 @@ def test_index_integration(dataset):
     assert MESSAGE_KEY in job_status_response
 
 
+@pytest.mark.integration
+def test_create_and_delete_custom_index(dataset):
+    # Creates image index
+    resp = dataset.set_continuous_indexing(True)
+
+    # Starts custom indexing job
+    signed_embeddings_url = TEST_INDEX_EMBEDDINGS_FILE
+    dataset.create_custom_index([signed_embeddings_url], embedding_dim=3)
+
+    resp = dataset.set_primary_index(image=True, custom=True)
+    assert resp["success"]
+
+    dataset.delete_custom_index(image=True)
+
+
 @pytest.mark.skip(reason="Times out consistently")
 def test_generate_image_index_integration(dataset):
     job = dataset.create_image_index()