Raise a named error to allow clients to handle failed jobs

Author: areeb1501

src/gradient/__init__.py

@@ -35,6 +35,7 @@
     PermissionDeniedError,
     UnprocessableEntityError,
     APIResponseValidationError,
+    IndexingJobError,
 )
 from ._base_client import DefaultHttpxClient, DefaultAioHttpClient, DefaultAsyncHttpxClient
 from ._utils._logs import setup_logging as _setup_logging
@@ -65,6 +66,7 @@
     "UnprocessableEntityError",
     "RateLimitError",
     "InternalServerError",
+    "IndexingJobError",
     "Timeout",
     "RequestOptions",
     "Client",

src/gradient/_exceptions.py

@@ -15,6 +15,7 @@
     "UnprocessableEntityError",
     "RateLimitError",
     "InternalServerError",
+    "IndexingJobError",
 ]
 
 
@@ -106,3 +107,15 @@ class RateLimitError(APIStatusError):
 
 class InternalServerError(APIStatusError):
     pass
+
+
+class IndexingJobError(GradientError):
+    """Raised when an indexing job fails, encounters an error, or is cancelled."""
+    
+    uuid: str
+    phase: str
+    
+    def __init__(self, message: str, *, uuid: str, phase: str) -> None:
+        super().__init__(message)
+        self.uuid = uuid
+        self.phase = phase

src/gradient/resources/knowledge_bases/indexing_jobs.py

@@ -7,6 +7,7 @@
 import httpx
 
 from ..._types import Body, Omit, Query, Headers, NotGiven, SequenceNotStr, omit, not_given
+from ..._exceptions import IndexingJobError
 from ..._utils import maybe_transform, async_maybe_transform
 from ..._compat import cached_property
 from ..._resource import SyncAPIResource, AsyncAPIResource
@@ -301,7 +302,7 @@ def wait_for_completion(
 
         Raises:
           TimeoutError: If the job doesn't complete within the specified timeout.
-          RuntimeError: If the job fails, errors, or is cancelled.
+          IndexingJobError: If the job fails, errors, or is cancelled.
         """
         if not uuid:
             raise ValueError(f"Expected a non-empty value for `uuid` but received {uuid!r}")
@@ -327,17 +328,27 @@ def wait_for_completion(
 
                 # Failure states
                 if phase == "BATCH_JOB_PHASE_FAILED":
-                    raise RuntimeError(
+                    raise IndexingJobError(
                         f"Indexing job {uuid} failed. "
                         f"Total items indexed: {response.job.total_items_indexed}, "
-                        f"Total items failed: {response.job.total_items_failed}"
+                        f"Total items failed: {response.job.total_items_failed}",
+                        uuid=uuid,
+                        phase=phase,
                     )
 
                 if phase == "BATCH_JOB_PHASE_ERROR":
-                    raise RuntimeError(f"Indexing job {uuid} encountered an error")
+                    raise IndexingJobError(
+                        f"Indexing job {uuid} encountered an error",
+                        uuid=uuid,
+                        phase=phase,
+                    )
 
                 if phase == "BATCH_JOB_PHASE_CANCELLED":
-                    raise RuntimeError(f"Indexing job {uuid} was cancelled")
+                    raise IndexingJobError(
+                        f"Indexing job {uuid} was cancelled",
+                        uuid=uuid,
+                        phase=phase,
+                    )
 
                 # Still in progress (UNKNOWN, PENDING, or RUNNING)
                 # Check timeout
@@ -623,7 +634,7 @@ async def wait_for_completion(
 
         Raises:
           TimeoutError: If the job doesn't complete within the specified timeout.
-          RuntimeError: If the job fails, errors, or is cancelled.
+          IndexingJobError: If the job fails, errors, or is cancelled.
         """
         if not uuid:
             raise ValueError(f"Expected a non-empty value for `uuid` but received {uuid!r}")
@@ -649,17 +660,27 @@ async def wait_for_completion(
 
                 # Failure states
                 if phase == "BATCH_JOB_PHASE_FAILED":
-                    raise RuntimeError(
+                    raise IndexingJobError(
                         f"Indexing job {uuid} failed. "
                         f"Total items indexed: {response.job.total_items_indexed}, "
-                        f"Total items failed: {response.job.total_items_failed}"
+                        f"Total items failed: {response.job.total_items_failed}",
+                        uuid=uuid,
+                        phase=phase,
                     )
 
                 if phase == "BATCH_JOB_PHASE_ERROR":
-                    raise RuntimeError(f"Indexing job {uuid} encountered an error")
+                    raise IndexingJobError(
+                        f"Indexing job {uuid} encountered an error",
+                        uuid=uuid,
+                        phase=phase,
+                    )
 
                 if phase == "BATCH_JOB_PHASE_CANCELLED":
-                    raise RuntimeError(f"Indexing job {uuid} was cancelled")
+                    raise IndexingJobError(
+                        f"Indexing job {uuid} was cancelled",
+                        uuid=uuid,
+                        phase=phase,
+                    )
 
                 # Still in progress (UNKNOWN, PENDING, or RUNNING)
                 # Check timeout