add from json function

Author: Bihan Jiang

nucleus/__init__.py

@@ -95,6 +95,7 @@
     ITEM_METADATA_SCHEMA_KEY,
     ITEMS_KEY,
     KEEP_HISTORY_KEY,
+    MESSAGE_KEY,
     MODEL_RUN_ID_KEY,
     NAME_KEY,
     NUCLEUS_ENDPOINT,
@@ -216,7 +217,7 @@ def list_jobs(
         return [
             AsyncJob(
                 job_id=job[JOB_ID_KEY],
-                job_status=job[JOB_STATUS_KEY],
+                job_last_known_status=job[JOB_STATUS_KEY],
                 job_type=job[JOB_TYPE_KEY],
                 job_creation_time=job[JOB_CREATION_TIME_KEY],
                 client=self,
@@ -1162,7 +1163,7 @@ def create_custom_index(
         embeddings_urls: list of urls, each of which being a json mapping dataset_item_id -> embedding vector
         embedding_dim: the dimension of the embedding vectors, must be consistent for all embedding vectors in the index.
         """
-        return self.make_request(
+        response_objects = self.make_request(
             {
                 EMBEDDINGS_URL_KEY: embeddings_urls,
                 EMBEDDING_DIMENSION_KEY: embedding_dim,
@@ -1171,6 +1172,18 @@ def create_custom_index(
             requests_command=requests.post,
         )
 
+        job = AsyncJob(
+            job_id=response_objects[JOB_ID_KEY],
+            job_last_known_status=response_objects[JOB_STATUS_KEY],
+            job_type=response_objects[JOB_TYPE_KEY],
+            job_creation_time=response_objects[JOB_CREATION_TIME_KEY],
+            client=self,
+        )
+        dataset_id = response_objects[DATASET_ID_KEY]
+        message = response_objects[MESSAGE_KEY]
+
+        return dataset_id, job, message
+
     def check_index_status(self, job_id: str):
         return self.make_request(
             {},

nucleus/dataset.py

@@ -23,10 +23,6 @@
     DATASET_SLICES_KEY,
     DEFAULT_ANNOTATION_UPDATE_MODE,
     EXPORTED_ROWS,
-    JOB_ID_KEY,
-    JOB_STATUS_KEY,
-    JOB_TYPE_KEY,
-    JOB_CREATION_TIME_KEY,
     NAME_KEY,
     REFERENCE_IDS_KEY,
     REQUEST_ID_KEY,
@@ -184,14 +180,7 @@ def annotate(
                 payload={REQUEST_ID_KEY: request_id, UPDATE_KEY: update},
                 route=f"dataset/{self.id}/annotate?async=1",
             )
-
-            return AsyncJob(
-                response[JOB_ID_KEY],
-                response[JOB_STATUS_KEY],
-                response[JOB_TYPE_KEY],
-                response[JOB_CREATION_TIME_KEY],
-                self._client,
-            )
+            return AsyncJob.from_json(response, self._client)
 
         return self._client.annotate_dataset(
             self.id, annotations, update=update, batch_size=batch_size
@@ -250,13 +239,7 @@ def append(
                 payload={REQUEST_ID_KEY: request_id, UPDATE_KEY: update},
                 route=f"dataset/{self.id}/append?async=1",
             )
-            return AsyncJob(
-                response[JOB_ID_KEY],
-                response[JOB_STATUS_KEY],
-                response[JOB_TYPE_KEY],
-                response[JOB_CREATION_TIME_KEY],
-                self._client,
-            )
+            return AsyncJob.from_json(response, self._client)
 
         return self._client.populate_dataset(
             self.id,

nucleus/job.py

@@ -1,16 +1,21 @@
 from dataclasses import dataclass
 import time
 from typing import Dict, List
-
 import requests
+from nucleus.constants import (
+    JOB_CREATION_TIME_KEY,
+    JOB_ID_KEY,
+    JOB_STATUS_KEY,
+    JOB_TYPE_KEY,
+)
 
 JOB_POLLING_INTERVAL = 5
 
 
 @dataclass
 class AsyncJob:
     job_id: str
-    job_status: str
+    job_last_known_status: str
     job_type: str
     job_creation_time: str
     client: "NucleusClient"  # type: ignore # noqa: F821
@@ -45,6 +50,16 @@ def sleep_until_complete(self, verbose_std_out=True):
         if final_status["status"] == "Errored":
             raise JobError(final_status, self)
 
+    @classmethod
+    def from_json(cls, payload: dict, client):
+        return cls(
+            job_id=payload[JOB_ID_KEY],
+            job_last_known_status=payload[JOB_STATUS_KEY],
+            job_type=payload[JOB_TYPE_KEY],
+            job_creation_time=payload[JOB_CREATION_TIME_KEY],
+            client=client,
+        )
+
 
 class JobError(Exception):
     def __init__(self, job_status: Dict[str, str], job: AsyncJob):

nucleus/model_run.py

@@ -8,10 +8,6 @@
     ANNOTATIONS_KEY,
     BOX_TYPE,
     DEFAULT_ANNOTATION_UPDATE_MODE,
-    JOB_ID_KEY,
-    JOB_STATUS_KEY,
-    JOB_TYPE_KEY,
-    JOB_CREATION_TIME_KEY,
     POLYGON_TYPE,
     REQUEST_ID_KEY,
     SEGMENTATION_TYPE,
@@ -118,14 +114,7 @@ def predict(
                 payload={REQUEST_ID_KEY: request_id, UPDATE_KEY: update},
                 route=f"modelRun/{self.model_run_id}/predict?async=1",
             )
-
-            return AsyncJob(
-                response[JOB_ID_KEY],
-                response[JOB_STATUS_KEY],
-                response[JOB_TYPE_KEY],
-                response[JOB_CREATION_TIME_KEY],
-                self._client,
-            )
+            return AsyncJob.from_json(response, self._client)
         else:
             return self._client.predict(self.model_run_id, annotations, update)
 

nucleus/slice.py

@@ -8,10 +8,6 @@
 from nucleus.utils import convert_export_payload, format_dataset_item_response
 from nucleus.constants import (
     EXPORTED_ROWS,
-    JOB_ID_KEY,
-    JOB_STATUS_KEY,
-    JOB_TYPE_KEY,
-    JOB_CREATION_TIME_KEY,
 )
 
 
@@ -128,13 +124,7 @@ def send_to_labeling(self, project_id: str):
         response = self._client.make_request(
             {}, f"slice/{self.slice_id}/{project_id}/send_to_labeling"
         )
-        return AsyncJob(
-            response[JOB_ID_KEY],
-            response[JOB_STATUS_KEY],
-            response[JOB_TYPE_KEY],
-            response[JOB_CREATION_TIME_KEY],
-            self._client,
-        )
+        return AsyncJob.from_json(response, self._client)
 
 
 def check_annotations_are_in_slice(