add new endpoints

Author: sasha-scale

nucleus/__init__.py

@@ -119,6 +119,7 @@
     ANNOTATION_METADATA_SCHEMA_KEY,
     ITEM_METADATA_SCHEMA_KEY,
     FORCE_KEY,
+    EMBEDDINGS_URL_KEY,
 )
 from .model import Model
 from .errors import (
@@ -991,6 +992,27 @@ def delete_model(self, model_id: str) -> dict:
         )
         return response
 
+    def create_custom_index(self, dataset_id: str, embeddings_url: str):
+        return self._make_request(
+            {EMBEDDINGS_URL_KEY: embeddings_url},
+            f"indexing/{dataset_id}",
+            requests_command=requests.post,
+        )
+
+    def check_index_status(self, job_id: str):
+        return self._make_request(
+            {},
+            f"indexing/{job_id}",
+            requests_command=requests.get,
+        )
+
+    def delete_custom_index(self, dataset_id: str):
+        return self._make_request(
+            {},
+            f"indexing/{dataset_id}",
+            requests_command=requests.delete,
+        )
+
     def _make_grequest(
         self,
         payload: dict,

nucleus/constants.py

@@ -56,3 +56,4 @@
 MASK_URL_KEY = "mask_url"
 INDEX_KEY = "index"
 SEGMENTATIONS_KEY = "segmentations"
+EMBEDDINGS_URL_KEY = "embeddings_url"

nucleus/dataset.py

@@ -4,7 +4,6 @@
     Annotation,
     BoxAnnotation,
     PolygonAnnotation,
-    SegmentationAnnotation,
 )
 from .constants import (
     DATASET_NAME_KEY,
@@ -17,9 +16,6 @@
     ITEM_KEY,
     DEFAULT_ANNOTATION_UPDATE_MODE,
     ANNOTATIONS_KEY,
-    BOX_TYPE,
-    POLYGON_TYPE,
-    SEGMENTATION_TYPE,
     ANNOTATION_TYPES,
 )
 from .payload_constructor import construct_model_run_creation_payload
@@ -260,3 +256,12 @@ def _format_dataset_item_response(self, response: dict) -> dict:
             ITEM_KEY: DatasetItem.from_json(item),
             ANNOTATIONS_KEY: annotation_response,
         }
+
+    def create_custom_index(self, embeddings_url: str):
+        return self._client.create_custom_index(self.id, embeddings_url)
+
+    def delete_custom_index(self):
+        return self._client.delete_custom_index(self.id)
+
+    def check_index_status(self, job_id: str):
+        return self._client.check_index_status(self, job_id)