From eed294340c70cea27cf41ed8a04806c0b482a2a1 Mon Sep 17 00:00:00 2001 From: areeb1501 <64038736+areeb1501@users.noreply.github.com> Date: Wed, 8 Oct 2025 00:54:56 +0400 Subject: [PATCH 01/14] feat: add wait_for_completion method to IndexingJobs resource with sync/async support --- examples/knowledge_base_indexing_wait.py | 196 ++++++++++++++++++ .../knowledge_bases/indexing_jobs.py | 196 ++++++++++++++++++ 2 files changed, 392 insertions(+) create mode 100644 examples/knowledge_base_indexing_wait.py diff --git a/examples/knowledge_base_indexing_wait.py b/examples/knowledge_base_indexing_wait.py new file mode 100644 index 00000000..8242750b --- /dev/null +++ b/examples/knowledge_base_indexing_wait.py @@ -0,0 +1,196 @@ +#!/usr/bin/env python3 +""" +Example: Waiting for Knowledge Base Indexing Job Completion + +This example demonstrates how to use the wait_for_completion() method +to automatically wait for a knowledge base indexing job to finish, +without needing to write manual polling loops. +""" + +import os +from gradient import Gradient + + +def main(): + # Initialize the Gradient client + client = Gradient() + + # Example 1: Basic usage - wait for indexing job to complete + print("Example 1: Basic usage") + print("-" * 50) + + # Create an indexing job (replace with your actual knowledge base UUID) + knowledge_base_uuid = os.getenv("KNOWLEDGE_BASE_UUID", "your-kb-uuid-here") + + print(f"Creating indexing job for knowledge base: {knowledge_base_uuid}") + indexing_job = client.knowledge_bases.indexing_jobs.create( + knowledge_base_uuid=knowledge_base_uuid, + ) + + job_uuid = indexing_job.job.uuid if indexing_job.job else None + if not job_uuid: + print("Error: Could not create indexing job") + return + + print(f"Indexing job created with UUID: {job_uuid}") + print("Waiting for indexing job to complete...") + + try: + # Wait for the job to complete (polls every 5 seconds by default) + completed_job = client.knowledge_bases.indexing_jobs.wait_for_completion( + job_uuid + ) + + print("\n✅ Indexing job completed successfully!") + if completed_job.job: + print(f"Phase: {completed_job.job.phase}") + print(f"Total items indexed: {completed_job.job.total_items_indexed}") + print(f"Total items failed: {completed_job.job.total_items_failed}") + print(f"Total datasources: {completed_job.job.total_datasources}") + print(f"Completed datasources: {completed_job.job.completed_datasources}") + + except TimeoutError as e: + print(f"\n⏱️ Timeout: {e}") + except RuntimeError as e: + print(f"\n❌ Error: {e}") + except Exception as e: + print(f"\n❌ Unexpected error: {e}") + + +def example_with_custom_polling(): + """Example with custom polling interval and timeout""" + print("\n\nExample 2: Custom polling interval and timeout") + print("-" * 50) + + client = Gradient() + knowledge_base_uuid = os.getenv("KNOWLEDGE_BASE_UUID", "your-kb-uuid-here") + + print(f"Creating indexing job for knowledge base: {knowledge_base_uuid}") + indexing_job = client.knowledge_bases.indexing_jobs.create( + knowledge_base_uuid=knowledge_base_uuid, + ) + + job_uuid = indexing_job.job.uuid if indexing_job.job else None + if not job_uuid: + print("Error: Could not create indexing job") + return + + print(f"Indexing job created with UUID: {job_uuid}") + print("Waiting for indexing job to complete (polling every 10 seconds, 5 minute timeout)...") + + try: + # Wait with custom poll interval (10 seconds) and timeout (5 minutes = 300 seconds) + completed_job = client.knowledge_bases.indexing_jobs.wait_for_completion( + job_uuid, + poll_interval=10, # Poll every 10 seconds + timeout=300, # Timeout after 5 minutes + ) + + print("\n✅ Indexing job completed successfully!") + if completed_job.job: + print(f"Phase: {completed_job.job.phase}") + + except TimeoutError: + print("\n⏱️ Job did not complete within 5 minutes") + # You can still check the current status + current_status = client.knowledge_bases.indexing_jobs.retrieve(job_uuid) + if current_status.job: + print(f"Current phase: {current_status.job.phase}") + print(f"Completed datasources: {current_status.job.completed_datasources}/{current_status.job.total_datasources}") + except RuntimeError as e: + print(f"\n❌ Job failed: {e}") + + +def example_manual_polling(): + """Example of the old manual polling approach (for comparison)""" + print("\n\nExample 3: Manual polling (old approach)") + print("-" * 50) + + client = Gradient() + knowledge_base_uuid = os.getenv("KNOWLEDGE_BASE_UUID", "your-kb-uuid-here") + + indexing_job = client.knowledge_bases.indexing_jobs.create( + knowledge_base_uuid=knowledge_base_uuid, + ) + + job_uuid = indexing_job.job.uuid if indexing_job.job else None + if not job_uuid: + print("Error: Could not create indexing job") + return + + print(f"Indexing job created with UUID: {job_uuid}") + print("Manual polling (old approach)...") + + import time + + while True: + indexing_job = client.knowledge_bases.indexing_jobs.retrieve(job_uuid) + + if indexing_job.job and indexing_job.job.phase: + phase = indexing_job.job.phase + print(f"Current phase: {phase}") + + if phase in ["BATCH_JOB_PHASE_UNKNOWN", "BATCH_JOB_PHASE_PENDING", "BATCH_JOB_PHASE_RUNNING"]: + time.sleep(5) + continue + elif phase == "BATCH_JOB_PHASE_SUCCEEDED": + print("✅ Job completed successfully!") + break + else: + print(f"❌ Job ended with phase: {phase}") + break + + +async def example_async(): + """Example using async/await""" + print("\n\nExample 4: Async usage") + print("-" * 50) + + from gradient import AsyncGradient + + client = AsyncGradient() + knowledge_base_uuid = os.getenv("KNOWLEDGE_BASE_UUID", "your-kb-uuid-here") + + print(f"Creating indexing job for knowledge base: {knowledge_base_uuid}") + indexing_job = await client.knowledge_bases.indexing_jobs.create( + knowledge_base_uuid=knowledge_base_uuid, + ) + + job_uuid = indexing_job.job.uuid if indexing_job.job else None + if not job_uuid: + print("Error: Could not create indexing job") + return + + print(f"Indexing job created with UUID: {job_uuid}") + print("Waiting for indexing job to complete (async)...") + + try: + completed_job = await client.knowledge_bases.indexing_jobs.wait_for_completion( + job_uuid, + poll_interval=5, + timeout=600, # 10 minute timeout + ) + + print("\n✅ Indexing job completed successfully!") + if completed_job.job: + print(f"Phase: {completed_job.job.phase}") + + except TimeoutError as e: + print(f"\n⏱️ Timeout: {e}") + except RuntimeError as e: + print(f"\n❌ Error: {e}") + finally: + await client.close() + + +if __name__ == "__main__": + # Run the basic example + main() + + # Uncomment to run other examples: + # example_with_custom_polling() + # example_manual_polling() + + # For async example, you would need to run: + # import asyncio + # asyncio.run(example_async()) diff --git a/src/gradient/resources/knowledge_bases/indexing_jobs.py b/src/gradient/resources/knowledge_bases/indexing_jobs.py index 95898c2a..0ece3cfb 100644 --- a/src/gradient/resources/knowledge_bases/indexing_jobs.py +++ b/src/gradient/resources/knowledge_bases/indexing_jobs.py @@ -2,6 +2,8 @@ from __future__ import annotations +import time +import asyncio import httpx from ..._types import Body, Omit, Query, Headers, NotGiven, SequenceNotStr, omit, not_given @@ -259,6 +261,97 @@ def update_cancel( cast_to=IndexingJobUpdateCancelResponse, ) + def wait_for_completion( + self, + uuid: str, + *, + poll_interval: int = 5, + timeout: int | None = None, + # Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs. + # The extra values given here take precedence over values defined on the client or passed to this method. + extra_headers: Headers | None = None, + extra_query: Query | None = None, + extra_body: Body | None = None, + request_timeout: float | httpx.Timeout | None | NotGiven = not_given, + ) -> IndexingJobRetrieveResponse: + """ + Wait for an indexing job to complete by polling its status. + + This method polls the indexing job status at regular intervals until it reaches + a terminal state (succeeded, failed, error, or cancelled). It raises an exception + if the job fails or times out. + + Args: + uuid: The UUID of the indexing job to wait for. + + poll_interval: Time in seconds between status checks (default: 5 seconds). + + timeout: Maximum time in seconds to wait for completion. If None, waits indefinitely. + + extra_headers: Send extra headers + + extra_query: Add additional query parameters to the request + + extra_body: Add additional JSON properties to the request + + request_timeout: Override the client-level default timeout for this request, in seconds + + Returns: + The final IndexingJobRetrieveResponse when the job completes successfully. + + Raises: + TimeoutError: If the job doesn't complete within the specified timeout. + RuntimeError: 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}") + + start_time = time.time() + + while True: + response = self.retrieve( + uuid, + extra_headers=extra_headers, + extra_query=extra_query, + extra_body=extra_body, + timeout=request_timeout, + ) + + # Check if job is in a terminal state + if response.job and response.job.phase: + phase = response.job.phase + + # Success state + if phase == "BATCH_JOB_PHASE_SUCCEEDED": + return response + + # Failure states + if phase == "BATCH_JOB_PHASE_FAILED": + raise RuntimeError( + f"Indexing job {uuid} failed. " + f"Total items indexed: {response.job.total_items_indexed}, " + f"Total items failed: {response.job.total_items_failed}" + ) + + if phase == "BATCH_JOB_PHASE_ERROR": + raise RuntimeError(f"Indexing job {uuid} encountered an error") + + if phase == "BATCH_JOB_PHASE_CANCELLED": + raise RuntimeError(f"Indexing job {uuid} was cancelled") + + # Still in progress (UNKNOWN, PENDING, or RUNNING) + # Check timeout + if timeout is not None: + elapsed = time.time() - start_time + if elapsed >= timeout: + raise TimeoutError( + f"Indexing job {uuid} did not complete within {timeout} seconds. " + f"Current phase: {phase}" + ) + + # Wait before next poll + time.sleep(poll_interval) + class AsyncIndexingJobsResource(AsyncAPIResource): @cached_property @@ -490,6 +583,97 @@ async def update_cancel( cast_to=IndexingJobUpdateCancelResponse, ) + async def wait_for_completion( + self, + uuid: str, + *, + poll_interval: int = 5, + timeout: int | None = None, + # Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs. + # The extra values given here take precedence over values defined on the client or passed to this method. + extra_headers: Headers | None = None, + extra_query: Query | None = None, + extra_body: Body | None = None, + request_timeout: float | httpx.Timeout | None | NotGiven = not_given, + ) -> IndexingJobRetrieveResponse: + """ + Wait for an indexing job to complete by polling its status. + + This method polls the indexing job status at regular intervals until it reaches + a terminal state (succeeded, failed, error, or cancelled). It raises an exception + if the job fails or times out. + + Args: + uuid: The UUID of the indexing job to wait for. + + poll_interval: Time in seconds between status checks (default: 5 seconds). + + timeout: Maximum time in seconds to wait for completion. If None, waits indefinitely. + + extra_headers: Send extra headers + + extra_query: Add additional query parameters to the request + + extra_body: Add additional JSON properties to the request + + request_timeout: Override the client-level default timeout for this request, in seconds + + Returns: + The final IndexingJobRetrieveResponse when the job completes successfully. + + Raises: + TimeoutError: If the job doesn't complete within the specified timeout. + RuntimeError: 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}") + + start_time = time.time() + + while True: + response = await self.retrieve( + uuid, + extra_headers=extra_headers, + extra_query=extra_query, + extra_body=extra_body, + timeout=request_timeout, + ) + + # Check if job is in a terminal state + if response.job and response.job.phase: + phase = response.job.phase + + # Success state + if phase == "BATCH_JOB_PHASE_SUCCEEDED": + return response + + # Failure states + if phase == "BATCH_JOB_PHASE_FAILED": + raise RuntimeError( + f"Indexing job {uuid} failed. " + f"Total items indexed: {response.job.total_items_indexed}, " + f"Total items failed: {response.job.total_items_failed}" + ) + + if phase == "BATCH_JOB_PHASE_ERROR": + raise RuntimeError(f"Indexing job {uuid} encountered an error") + + if phase == "BATCH_JOB_PHASE_CANCELLED": + raise RuntimeError(f"Indexing job {uuid} was cancelled") + + # Still in progress (UNKNOWN, PENDING, or RUNNING) + # Check timeout + if timeout is not None: + elapsed = time.time() - start_time + if elapsed >= timeout: + raise TimeoutError( + f"Indexing job {uuid} did not complete within {timeout} seconds. " + f"Current phase: {phase}" + ) + + # Wait before next poll + await asyncio.sleep(poll_interval) + class IndexingJobsResourceWithRawResponse: def __init__(self, indexing_jobs: IndexingJobsResource) -> None: @@ -510,6 +694,9 @@ def __init__(self, indexing_jobs: IndexingJobsResource) -> None: self.update_cancel = to_raw_response_wrapper( indexing_jobs.update_cancel, ) + self.wait_for_completion = to_raw_response_wrapper( + indexing_jobs.wait_for_completion, + ) class AsyncIndexingJobsResourceWithRawResponse: @@ -531,6 +718,9 @@ def __init__(self, indexing_jobs: AsyncIndexingJobsResource) -> None: self.update_cancel = async_to_raw_response_wrapper( indexing_jobs.update_cancel, ) + self.wait_for_completion = async_to_raw_response_wrapper( + indexing_jobs.wait_for_completion, + ) class IndexingJobsResourceWithStreamingResponse: @@ -552,6 +742,9 @@ def __init__(self, indexing_jobs: IndexingJobsResource) -> None: self.update_cancel = to_streamed_response_wrapper( indexing_jobs.update_cancel, ) + self.wait_for_completion = to_streamed_response_wrapper( + indexing_jobs.wait_for_completion, + ) class AsyncIndexingJobsResourceWithStreamingResponse: @@ -573,3 +766,6 @@ def __init__(self, indexing_jobs: AsyncIndexingJobsResource) -> None: self.update_cancel = async_to_streamed_response_wrapper( indexing_jobs.update_cancel, ) + self.wait_for_completion = async_to_streamed_response_wrapper( + indexing_jobs.wait_for_completion, + ) From 9ef2314224434b98a55ce8447f6e9d68f5f64a10 Mon Sep 17 00:00:00 2001 From: areeb1501 <64038736+areeb1501@users.noreply.github.com> Date: Tue, 21 Oct 2025 02:12:46 +0400 Subject: [PATCH 02/14] Raise a named error to allow clients to handle failed jobs --- src/gradient/__init__.py | 2 + src/gradient/_exceptions.py | 13 ++++++ .../knowledge_bases/indexing_jobs.py | 41 ++++++++++++++----- 3 files changed, 46 insertions(+), 10 deletions(-) diff --git a/src/gradient/__init__.py b/src/gradient/__init__.py index a67cd2a7..0d9cb1e3 100644 --- a/src/gradient/__init__.py +++ b/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", diff --git a/src/gradient/_exceptions.py b/src/gradient/_exceptions.py index 5db08573..b5d835b4 100644 --- a/src/gradient/_exceptions.py +++ b/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 diff --git a/src/gradient/resources/knowledge_bases/indexing_jobs.py b/src/gradient/resources/knowledge_bases/indexing_jobs.py index 0ece3cfb..9d1a2f4e 100644 --- a/src/gradient/resources/knowledge_bases/indexing_jobs.py +++ b/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 From c9a4c6cc552f7dddae7aa2c6ea99899fa1f01a5d Mon Sep 17 00:00:00 2001 From: areeb1501 <64038736+areeb1501@users.noreply.github.com> Date: Tue, 21 Oct 2025 02:17:41 +0400 Subject: [PATCH 03/14] Handle polling timeouts directrly --- src/gradient/__init__.py | 2 + src/gradient/_exceptions.py | 19 +++++++- .../knowledge_bases/indexing_jobs.py | 48 +++++++++++-------- 3 files changed, 46 insertions(+), 23 deletions(-) diff --git a/src/gradient/__init__.py b/src/gradient/__init__.py index 0d9cb1e3..ed8bc3ae 100644 --- a/src/gradient/__init__.py +++ b/src/gradient/__init__.py @@ -36,6 +36,7 @@ UnprocessableEntityError, APIResponseValidationError, IndexingJobError, + IndexingJobTimeoutError, ) from ._base_client import DefaultHttpxClient, DefaultAioHttpClient, DefaultAsyncHttpxClient from ._utils._logs import setup_logging as _setup_logging @@ -67,6 +68,7 @@ "RateLimitError", "InternalServerError", "IndexingJobError", + "IndexingJobTimeoutError", "Timeout", "RequestOptions", "Client", diff --git a/src/gradient/_exceptions.py b/src/gradient/_exceptions.py index b5d835b4..984778fc 100644 --- a/src/gradient/_exceptions.py +++ b/src/gradient/_exceptions.py @@ -16,6 +16,7 @@ "RateLimitError", "InternalServerError", "IndexingJobError", + "IndexingJobTimeoutError", ] @@ -111,11 +112,25 @@ class InternalServerError(APIStatusError): 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 + + +class IndexingJobTimeoutError(GradientError): + """Raised when polling for an indexing job times out.""" + + uuid: str + phase: str + timeout: int + + def __init__(self, message: str, *, uuid: str, phase: str, timeout: int) -> None: + super().__init__(message) + self.uuid = uuid + self.phase = phase + self.timeout = timeout diff --git a/src/gradient/resources/knowledge_bases/indexing_jobs.py b/src/gradient/resources/knowledge_bases/indexing_jobs.py index 9d1a2f4e..c0c0508d 100644 --- a/src/gradient/resources/knowledge_bases/indexing_jobs.py +++ b/src/gradient/resources/knowledge_bases/indexing_jobs.py @@ -7,7 +7,7 @@ import httpx from ..._types import Body, Omit, Query, Headers, NotGiven, SequenceNotStr, omit, not_given -from ..._exceptions import IndexingJobError +from ..._exceptions import IndexingJobError, IndexingJobTimeoutError from ..._utils import maybe_transform, async_maybe_transform from ..._compat import cached_property from ..._resource import SyncAPIResource, AsyncAPIResource @@ -301,14 +301,14 @@ def wait_for_completion( The final IndexingJobRetrieveResponse when the job completes successfully. Raises: - TimeoutError: If the job doesn't complete within the specified timeout. + IndexingJobTimeoutError: If the job doesn't complete within the specified timeout. 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}") start_time = time.time() - + while True: response = self.retrieve( uuid, @@ -321,11 +321,11 @@ def wait_for_completion( # Check if job is in a terminal state if response.job and response.job.phase: phase = response.job.phase - + # Success state if phase == "BATCH_JOB_PHASE_SUCCEEDED": return response - + # Failure states if phase == "BATCH_JOB_PHASE_FAILED": raise IndexingJobError( @@ -335,31 +335,34 @@ def wait_for_completion( uuid=uuid, phase=phase, ) - + if phase == "BATCH_JOB_PHASE_ERROR": raise IndexingJobError( f"Indexing job {uuid} encountered an error", uuid=uuid, phase=phase, ) - + if phase == "BATCH_JOB_PHASE_CANCELLED": raise IndexingJobError( f"Indexing job {uuid} was cancelled", uuid=uuid, phase=phase, ) - + # Still in progress (UNKNOWN, PENDING, or RUNNING) # Check timeout if timeout is not None: elapsed = time.time() - start_time if elapsed >= timeout: - raise TimeoutError( + raise IndexingJobTimeoutError( f"Indexing job {uuid} did not complete within {timeout} seconds. " - f"Current phase: {phase}" + f"Current phase: {phase}", + uuid=uuid, + phase=phase, + timeout=timeout, ) - + # Wait before next poll time.sleep(poll_interval) @@ -633,14 +636,14 @@ async def wait_for_completion( The final IndexingJobRetrieveResponse when the job completes successfully. Raises: - TimeoutError: If the job doesn't complete within the specified timeout. + IndexingJobTimeoutError: If the job doesn't complete within the specified timeout. 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}") start_time = time.time() - + while True: response = await self.retrieve( uuid, @@ -653,11 +656,11 @@ async def wait_for_completion( # Check if job is in a terminal state if response.job and response.job.phase: phase = response.job.phase - + # Success state if phase == "BATCH_JOB_PHASE_SUCCEEDED": return response - + # Failure states if phase == "BATCH_JOB_PHASE_FAILED": raise IndexingJobError( @@ -667,31 +670,34 @@ async def wait_for_completion( uuid=uuid, phase=phase, ) - + if phase == "BATCH_JOB_PHASE_ERROR": raise IndexingJobError( f"Indexing job {uuid} encountered an error", uuid=uuid, phase=phase, ) - + if phase == "BATCH_JOB_PHASE_CANCELLED": raise IndexingJobError( f"Indexing job {uuid} was cancelled", uuid=uuid, phase=phase, ) - + # Still in progress (UNKNOWN, PENDING, or RUNNING) # Check timeout if timeout is not None: elapsed = time.time() - start_time if elapsed >= timeout: - raise TimeoutError( + raise IndexingJobTimeoutError( f"Indexing job {uuid} did not complete within {timeout} seconds. " - f"Current phase: {phase}" + f"Current phase: {phase}", + uuid=uuid, + phase=phase, + timeout=timeout, ) - + # Wait before next poll await asyncio.sleep(poll_interval) From e7a7e675a66a827f868ec7deb30575f78ce26c47 Mon Sep 17 00:00:00 2001 From: areeb1501 <64038736+areeb1501@users.noreply.github.com> Date: Tue, 21 Oct 2025 02:18:55 +0400 Subject: [PATCH 04/14] change type of poll interval from int to float --- src/gradient/_exceptions.py | 4 ++-- src/gradient/resources/knowledge_bases/indexing_jobs.py | 8 ++++---- 2 files changed, 6 insertions(+), 6 deletions(-) diff --git a/src/gradient/_exceptions.py b/src/gradient/_exceptions.py index 984778fc..8f415fbe 100644 --- a/src/gradient/_exceptions.py +++ b/src/gradient/_exceptions.py @@ -127,9 +127,9 @@ class IndexingJobTimeoutError(GradientError): uuid: str phase: str - timeout: int + timeout: float - def __init__(self, message: str, *, uuid: str, phase: str, timeout: int) -> None: + def __init__(self, message: str, *, uuid: str, phase: str, timeout: float) -> None: super().__init__(message) self.uuid = uuid self.phase = phase diff --git a/src/gradient/resources/knowledge_bases/indexing_jobs.py b/src/gradient/resources/knowledge_bases/indexing_jobs.py index c0c0508d..5d170781 100644 --- a/src/gradient/resources/knowledge_bases/indexing_jobs.py +++ b/src/gradient/resources/knowledge_bases/indexing_jobs.py @@ -266,8 +266,8 @@ def wait_for_completion( self, uuid: str, *, - poll_interval: int = 5, - timeout: int | None = None, + poll_interval: float = 5, + timeout: float | None = None, # Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs. # The extra values given here take precedence over values defined on the client or passed to this method. extra_headers: Headers | None = None, @@ -601,8 +601,8 @@ async def wait_for_completion( self, uuid: str, *, - poll_interval: int = 5, - timeout: int | None = None, + poll_interval: float = 5, + timeout: float | None = None, # Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs. # The extra values given here take precedence over values defined on the client or passed to this method. extra_headers: Headers | None = None, From d6a55bc623fd45aba91065b107634fed6ec4880b Mon Sep 17 00:00:00 2001 From: areeb1501 <64038736+areeb1501@users.noreply.github.com> Date: Tue, 21 Oct 2025 02:21:44 +0400 Subject: [PATCH 05/14] Resolve merge conflict --- src/gradient/__init__.py | 4 ++++ src/gradient/_exceptions.py | 18 ++++++++++++++++++ 2 files changed, 22 insertions(+) diff --git a/src/gradient/__init__.py b/src/gradient/__init__.py index ed8bc3ae..24bffd17 100644 --- a/src/gradient/__init__.py +++ b/src/gradient/__init__.py @@ -37,6 +37,8 @@ APIResponseValidationError, IndexingJobError, IndexingJobTimeoutError, + AgentDeploymentError, + AgentDeploymentTimeoutError, ) from ._base_client import DefaultHttpxClient, DefaultAioHttpClient, DefaultAsyncHttpxClient from ._utils._logs import setup_logging as _setup_logging @@ -69,6 +71,8 @@ "InternalServerError", "IndexingJobError", "IndexingJobTimeoutError", + "AgentDeploymentError", + "AgentDeploymentTimeoutError", "Timeout", "RequestOptions", "Client", diff --git a/src/gradient/_exceptions.py b/src/gradient/_exceptions.py index 8f415fbe..f0c6671d 100644 --- a/src/gradient/_exceptions.py +++ b/src/gradient/_exceptions.py @@ -17,6 +17,8 @@ "InternalServerError", "IndexingJobError", "IndexingJobTimeoutError", + "AgentDeploymentError", + "AgentDeploymentTimeoutError", ] @@ -134,3 +136,19 @@ def __init__(self, message: str, *, uuid: str, phase: str, timeout: float) -> No self.uuid = uuid self.phase = phase self.timeout = timeout + + +class AgentDeploymentError(GradientError): + """Raised when an agent deployment fails.""" + + def __init__(self, message: str, status: str) -> None: + super().__init__(message) + self.status = status + + +class AgentDeploymentTimeoutError(GradientError): + """Raised when waiting for an agent deployment times out.""" + + def __init__(self, message: str, agent_id: str) -> None: + super().__init__(message) + self.agent_id = agent_id From 0b9edf20e64006c1655008ca5ca0fa4618dcc5a9 Mon Sep 17 00:00:00 2001 From: areeb1501 <64038736+areeb1501@users.noreply.github.com> Date: Tue, 21 Oct 2025 02:28:21 +0400 Subject: [PATCH 06/14] Added test cases --- .../knowledge_bases/test_indexing_jobs.py | 238 ++++++++++++++++++ 1 file changed, 238 insertions(+) diff --git a/tests/api_resources/knowledge_bases/test_indexing_jobs.py b/tests/api_resources/knowledge_bases/test_indexing_jobs.py index 3dffaa69..c69e216d 100644 --- a/tests/api_resources/knowledge_bases/test_indexing_jobs.py +++ b/tests/api_resources/knowledge_bases/test_indexing_jobs.py @@ -5,9 +5,11 @@ import os from typing import Any, cast +import httpx import pytest from gradient import Gradient, AsyncGradient +from gradient import IndexingJobError, IndexingJobTimeoutError from tests.utils import assert_matches_type from gradient.types.knowledge_bases import ( IndexingJobListResponse, @@ -232,6 +234,124 @@ def test_path_params_update_cancel(self, client: Gradient) -> None: path_uuid="", ) + @parametrize + def test_wait_for_completion_raises_indexing_job_error_on_failed(self, client: Gradient, respx_mock: Any) -> None: + """Test that IndexingJobError is raised when job phase is FAILED""" + job_uuid = "test-job-uuid" + respx_mock.get(f"{base_url}/v2/gen-ai/indexing_jobs/{job_uuid}").mock( + return_value=httpx.Response( + 200, + json={ + "job": { + "uuid": job_uuid, + "phase": "BATCH_JOB_PHASE_FAILED", + "total_items_indexed": 10, + "total_items_failed": 5, + } + }, + ) + ) + + with pytest.raises(IndexingJobError) as exc_info: + client.knowledge_bases.indexing_jobs.wait_for_completion(job_uuid) + + assert exc_info.value.uuid == job_uuid + assert exc_info.value.phase == "BATCH_JOB_PHASE_FAILED" + assert "failed" in str(exc_info.value).lower() + + @parametrize + def test_wait_for_completion_raises_indexing_job_error_on_error(self, client: Gradient, respx_mock: Any) -> None: + """Test that IndexingJobError is raised when job phase is ERROR""" + job_uuid = "test-job-uuid" + respx_mock.get(f"{base_url}/v2/gen-ai/indexing_jobs/{job_uuid}").mock( + return_value=httpx.Response( + 200, + json={ + "job": { + "uuid": job_uuid, + "phase": "BATCH_JOB_PHASE_ERROR", + } + }, + ) + ) + + with pytest.raises(IndexingJobError) as exc_info: + client.knowledge_bases.indexing_jobs.wait_for_completion(job_uuid) + + assert exc_info.value.uuid == job_uuid + assert exc_info.value.phase == "BATCH_JOB_PHASE_ERROR" + assert "error" in str(exc_info.value).lower() + + @parametrize + def test_wait_for_completion_raises_indexing_job_error_on_cancelled(self, client: Gradient, respx_mock: Any) -> None: + """Test that IndexingJobError is raised when job phase is CANCELLED""" + job_uuid = "test-job-uuid" + respx_mock.get(f"{base_url}/v2/gen-ai/indexing_jobs/{job_uuid}").mock( + return_value=httpx.Response( + 200, + json={ + "job": { + "uuid": job_uuid, + "phase": "BATCH_JOB_PHASE_CANCELLED", + } + }, + ) + ) + + with pytest.raises(IndexingJobError) as exc_info: + client.knowledge_bases.indexing_jobs.wait_for_completion(job_uuid) + + assert exc_info.value.uuid == job_uuid + assert exc_info.value.phase == "BATCH_JOB_PHASE_CANCELLED" + assert "cancelled" in str(exc_info.value).lower() + + @parametrize + def test_wait_for_completion_raises_timeout_error(self, client: Gradient, respx_mock: Any) -> None: + """Test that IndexingJobTimeoutError is raised on timeout""" + job_uuid = "test-job-uuid" + respx_mock.get(f"{base_url}/v2/gen-ai/indexing_jobs/{job_uuid}").mock( + return_value=httpx.Response( + 200, + json={ + "job": { + "uuid": job_uuid, + "phase": "BATCH_JOB_PHASE_RUNNING", + } + }, + ) + ) + + with pytest.raises(IndexingJobTimeoutError) as exc_info: + client.knowledge_bases.indexing_jobs.wait_for_completion( + job_uuid, poll_interval=0.1, timeout=0.2 + ) + + assert exc_info.value.uuid == job_uuid + assert exc_info.value.phase == "BATCH_JOB_PHASE_RUNNING" + assert exc_info.value.timeout == 0.2 + + @parametrize + def test_wait_for_completion_succeeds(self, client: Gradient, respx_mock: Any) -> None: + """Test that wait_for_completion returns successfully when job succeeds""" + job_uuid = "test-job-uuid" + respx_mock.get(f"{base_url}/v2/gen-ai/indexing_jobs/{job_uuid}").mock( + return_value=httpx.Response( + 200, + json={ + "job": { + "uuid": job_uuid, + "phase": "BATCH_JOB_PHASE_SUCCEEDED", + "total_items_indexed": 100, + "total_items_failed": 0, + } + }, + ) + ) + + result = client.knowledge_bases.indexing_jobs.wait_for_completion(job_uuid) + assert_matches_type(IndexingJobRetrieveResponse, result, path=["response"]) + assert result.job.phase == "BATCH_JOB_PHASE_SUCCEEDED" + class TestAsyncIndexingJobs: parametrize = pytest.mark.parametrize( @@ -446,3 +566,121 @@ async def test_path_params_update_cancel(self, async_client: AsyncGradient) -> N await async_client.knowledge_bases.indexing_jobs.with_raw_response.update_cancel( path_uuid="", ) + + @parametrize + async def test_wait_for_completion_raises_indexing_job_error_on_failed(self, async_client: AsyncGradient, respx_mock: Any) -> None: + """Test that IndexingJobError is raised when job phase is FAILED""" + job_uuid = "test-job-uuid" + respx_mock.get(f"{base_url}/v2/gen-ai/indexing_jobs/{job_uuid}").mock( + return_value=httpx.Response( + 200, + json={ + "job": { + "uuid": job_uuid, + "phase": "BATCH_JOB_PHASE_FAILED", + "total_items_indexed": 10, + "total_items_failed": 5, + } + }, + ) + ) + + with pytest.raises(IndexingJobError) as exc_info: + await async_client.knowledge_bases.indexing_jobs.wait_for_completion(job_uuid) + + assert exc_info.value.uuid == job_uuid + assert exc_info.value.phase == "BATCH_JOB_PHASE_FAILED" + assert "failed" in str(exc_info.value).lower() + + @parametrize + async def test_wait_for_completion_raises_indexing_job_error_on_error(self, async_client: AsyncGradient, respx_mock: Any) -> None: + """Test that IndexingJobError is raised when job phase is ERROR""" + job_uuid = "test-job-uuid" + respx_mock.get(f"{base_url}/v2/gen-ai/indexing_jobs/{job_uuid}").mock( + return_value=httpx.Response( + 200, + json={ + "job": { + "uuid": job_uuid, + "phase": "BATCH_JOB_PHASE_ERROR", + } + }, + ) + ) + + with pytest.raises(IndexingJobError) as exc_info: + await async_client.knowledge_bases.indexing_jobs.wait_for_completion(job_uuid) + + assert exc_info.value.uuid == job_uuid + assert exc_info.value.phase == "BATCH_JOB_PHASE_ERROR" + assert "error" in str(exc_info.value).lower() + + @parametrize + async def test_wait_for_completion_raises_indexing_job_error_on_cancelled(self, async_client: AsyncGradient, respx_mock: Any) -> None: + """Test that IndexingJobError is raised when job phase is CANCELLED""" + job_uuid = "test-job-uuid" + respx_mock.get(f"{base_url}/v2/gen-ai/indexing_jobs/{job_uuid}").mock( + return_value=httpx.Response( + 200, + json={ + "job": { + "uuid": job_uuid, + "phase": "BATCH_JOB_PHASE_CANCELLED", + } + }, + ) + ) + + with pytest.raises(IndexingJobError) as exc_info: + await async_client.knowledge_bases.indexing_jobs.wait_for_completion(job_uuid) + + assert exc_info.value.uuid == job_uuid + assert exc_info.value.phase == "BATCH_JOB_PHASE_CANCELLED" + assert "cancelled" in str(exc_info.value).lower() + + @parametrize + async def test_wait_for_completion_raises_timeout_error(self, async_client: AsyncGradient, respx_mock: Any) -> None: + """Test that IndexingJobTimeoutError is raised on timeout""" + job_uuid = "test-job-uuid" + respx_mock.get(f"{base_url}/v2/gen-ai/indexing_jobs/{job_uuid}").mock( + return_value=httpx.Response( + 200, + json={ + "job": { + "uuid": job_uuid, + "phase": "BATCH_JOB_PHASE_RUNNING", + } + }, + ) + ) + + with pytest.raises(IndexingJobTimeoutError) as exc_info: + await async_client.knowledge_bases.indexing_jobs.wait_for_completion( + job_uuid, poll_interval=0.1, timeout=0.2 + ) + + assert exc_info.value.uuid == job_uuid + assert exc_info.value.phase == "BATCH_JOB_PHASE_RUNNING" + assert exc_info.value.timeout == 0.2 + + @parametrize + async def test_wait_for_completion_succeeds(self, async_client: AsyncGradient, respx_mock: Any) -> None: + """Test that wait_for_completion returns successfully when job succeeds""" + job_uuid = "test-job-uuid" + respx_mock.get(f"{base_url}/v2/gen-ai/indexing_jobs/{job_uuid}").mock( + return_value=httpx.Response( + 200, + json={ + "job": { + "uuid": job_uuid, + "phase": "BATCH_JOB_PHASE_SUCCEEDED", + "total_items_indexed": 100, + "total_items_failed": 0, + } + }, + ) + ) + + result = await async_client.knowledge_bases.indexing_jobs.wait_for_completion(job_uuid) + assert_matches_type(IndexingJobRetrieveResponse, result, path=["response"]) + assert result.job.phase == "BATCH_JOB_PHASE_SUCCEEDED" From 192cf3e787ccb4b0198e0df2f7d3d7eaa4e10c4a Mon Sep 17 00:00:00 2001 From: areeb1501 <64038736+areeb1501@users.noreply.github.com> Date: Wed, 22 Oct 2025 21:29:06 +0400 Subject: [PATCH 07/14] Adding newly created error types --- examples/knowledge_base_indexing_wait.py | 15 ++++++++------- 1 file changed, 8 insertions(+), 7 deletions(-) diff --git a/examples/knowledge_base_indexing_wait.py b/examples/knowledge_base_indexing_wait.py index 8242750b..9fa618db 100644 --- a/examples/knowledge_base_indexing_wait.py +++ b/examples/knowledge_base_indexing_wait.py @@ -9,6 +9,7 @@ import os from gradient import Gradient +from gradient import IndexingJobError, IndexingJobTimeoutError def main(): @@ -49,9 +50,9 @@ def main(): print(f"Total datasources: {completed_job.job.total_datasources}") print(f"Completed datasources: {completed_job.job.completed_datasources}") - except TimeoutError as e: + except IndexingJobTimeoutError as e: print(f"\n⏱️ Timeout: {e}") - except RuntimeError as e: + except IndexingJobError as e: print(f"\n❌ Error: {e}") except Exception as e: print(f"\n❌ Unexpected error: {e}") @@ -90,14 +91,14 @@ def example_with_custom_polling(): if completed_job.job: print(f"Phase: {completed_job.job.phase}") - except TimeoutError: + except IndexingJobTimeoutError: print("\n⏱️ Job did not complete within 5 minutes") # You can still check the current status current_status = client.knowledge_bases.indexing_jobs.retrieve(job_uuid) if current_status.job: print(f"Current phase: {current_status.job.phase}") print(f"Completed datasources: {current_status.job.completed_datasources}/{current_status.job.total_datasources}") - except RuntimeError as e: + except IndexingJobError as e: print(f"\n❌ Job failed: {e}") @@ -146,7 +147,7 @@ async def example_async(): print("\n\nExample 4: Async usage") print("-" * 50) - from gradient import AsyncGradient + from gradient import AsyncGradient, IndexingJobError, IndexingJobTimeoutError client = AsyncGradient() knowledge_base_uuid = os.getenv("KNOWLEDGE_BASE_UUID", "your-kb-uuid-here") @@ -175,9 +176,9 @@ async def example_async(): if completed_job.job: print(f"Phase: {completed_job.job.phase}") - except TimeoutError as e: + except IndexingJobTimeoutError as e: print(f"\n⏱️ Timeout: {e}") - except RuntimeError as e: + except IndexingJobError as e: print(f"\n❌ Error: {e}") finally: await client.close() From 6eca8ccda5019c5e128fa97677d9121230f7324b Mon Sep 17 00:00:00 2001 From: areeb1501 <64038736+areeb1501@users.noreply.github.com> Date: Wed, 22 Oct 2025 21:31:12 +0400 Subject: [PATCH 08/14] Removing redundant imports --- examples/knowledge_base_indexing_wait.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/examples/knowledge_base_indexing_wait.py b/examples/knowledge_base_indexing_wait.py index 9fa618db..036177b2 100644 --- a/examples/knowledge_base_indexing_wait.py +++ b/examples/knowledge_base_indexing_wait.py @@ -147,7 +147,7 @@ async def example_async(): print("\n\nExample 4: Async usage") print("-" * 50) - from gradient import AsyncGradient, IndexingJobError, IndexingJobTimeoutError + from gradient import AsyncGradient client = AsyncGradient() knowledge_base_uuid = os.getenv("KNOWLEDGE_BASE_UUID", "your-kb-uuid-here") From 01dac39dea059dbcc32d9ca871b50ed20a02b796 Mon Sep 17 00:00:00 2001 From: areeb1501 <64038736+areeb1501@users.noreply.github.com> Date: Wed, 22 Oct 2025 21:42:49 +0400 Subject: [PATCH 09/14] combine imports --- examples/knowledge_base_indexing_wait.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/examples/knowledge_base_indexing_wait.py b/examples/knowledge_base_indexing_wait.py index 036177b2..3ee762bb 100644 --- a/examples/knowledge_base_indexing_wait.py +++ b/examples/knowledge_base_indexing_wait.py @@ -8,8 +8,8 @@ """ import os -from gradient import Gradient -from gradient import IndexingJobError, IndexingJobTimeoutError + +from gradient import Gradient, IndexingJobError, IndexingJobTimeoutError def main(): From dedbad9af9d080fd9bfa1d62d66d6462932b0429 Mon Sep 17 00:00:00 2001 From: areeb1501 <64038736+areeb1501@users.noreply.github.com> Date: Wed, 22 Oct 2025 21:44:57 +0400 Subject: [PATCH 10/14] Update test cases --- .../knowledge_bases/test_indexing_jobs.py | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/tests/api_resources/knowledge_bases/test_indexing_jobs.py b/tests/api_resources/knowledge_bases/test_indexing_jobs.py index c69e216d..b55ea5b0 100644 --- a/tests/api_resources/knowledge_bases/test_indexing_jobs.py +++ b/tests/api_resources/knowledge_bases/test_indexing_jobs.py @@ -245,8 +245,8 @@ def test_wait_for_completion_raises_indexing_job_error_on_failed(self, client: G "job": { "uuid": job_uuid, "phase": "BATCH_JOB_PHASE_FAILED", - "total_items_indexed": 10, - "total_items_failed": 5, + "total_items_indexed": "10", + "total_items_failed": "5", } }, ) @@ -341,8 +341,8 @@ def test_wait_for_completion_succeeds(self, client: Gradient, respx_mock: Any) - "job": { "uuid": job_uuid, "phase": "BATCH_JOB_PHASE_SUCCEEDED", - "total_items_indexed": 100, - "total_items_failed": 0, + "total_items_indexed": "100", + "total_items_failed": "0", } }, ) @@ -578,8 +578,8 @@ async def test_wait_for_completion_raises_indexing_job_error_on_failed(self, asy "job": { "uuid": job_uuid, "phase": "BATCH_JOB_PHASE_FAILED", - "total_items_indexed": 10, - "total_items_failed": 5, + "total_items_indexed": "10", + "total_items_failed": "5", } }, ) @@ -674,8 +674,8 @@ async def test_wait_for_completion_succeeds(self, async_client: AsyncGradient, r "job": { "uuid": job_uuid, "phase": "BATCH_JOB_PHASE_SUCCEEDED", - "total_items_indexed": 100, - "total_items_failed": 0, + "total_items_indexed": "100", + "total_items_failed": "0", } }, ) From 8bda2013245931d3ea03c881c718cc4bd6d20735 Mon Sep 17 00:00:00 2001 From: areeb1501 <64038736+areeb1501@users.noreply.github.com> Date: Thu, 23 Oct 2025 01:56:13 +0400 Subject: [PATCH 11/14] Fix linting --- src/gradient/__init__.py | 6 +++--- src/gradient/resources/knowledge_bases/indexing_jobs.py | 3 ++- tests/api_resources/knowledge_bases/test_indexing_jobs.py | 3 +-- 3 files changed, 6 insertions(+), 6 deletions(-) diff --git a/src/gradient/__init__.py b/src/gradient/__init__.py index 24bffd17..864f1484 100644 --- a/src/gradient/__init__.py +++ b/src/gradient/__init__.py @@ -29,15 +29,15 @@ RateLimitError, APITimeoutError, BadRequestError, + IndexingJobError, APIConnectionError, AuthenticationError, InternalServerError, + AgentDeploymentError, PermissionDeniedError, + IndexingJobTimeoutError, UnprocessableEntityError, APIResponseValidationError, - IndexingJobError, - IndexingJobTimeoutError, - AgentDeploymentError, AgentDeploymentTimeoutError, ) from ._base_client import DefaultHttpxClient, DefaultAioHttpClient, DefaultAsyncHttpxClient diff --git a/src/gradient/resources/knowledge_bases/indexing_jobs.py b/src/gradient/resources/knowledge_bases/indexing_jobs.py index 5d170781..647ab308 100644 --- a/src/gradient/resources/knowledge_bases/indexing_jobs.py +++ b/src/gradient/resources/knowledge_bases/indexing_jobs.py @@ -4,10 +4,10 @@ import time import asyncio + import httpx from ..._types import Body, Omit, Query, Headers, NotGiven, SequenceNotStr, omit, not_given -from ..._exceptions import IndexingJobError, IndexingJobTimeoutError from ..._utils import maybe_transform, async_maybe_transform from ..._compat import cached_property from ..._resource import SyncAPIResource, AsyncAPIResource @@ -17,6 +17,7 @@ async_to_raw_response_wrapper, async_to_streamed_response_wrapper, ) +from ..._exceptions import IndexingJobError, IndexingJobTimeoutError from ..._base_client import make_request_options from ...types.knowledge_bases import ( indexing_job_list_params, diff --git a/tests/api_resources/knowledge_bases/test_indexing_jobs.py b/tests/api_resources/knowledge_bases/test_indexing_jobs.py index b55ea5b0..216f09e0 100644 --- a/tests/api_resources/knowledge_bases/test_indexing_jobs.py +++ b/tests/api_resources/knowledge_bases/test_indexing_jobs.py @@ -8,8 +8,7 @@ import httpx import pytest -from gradient import Gradient, AsyncGradient -from gradient import IndexingJobError, IndexingJobTimeoutError +from gradient import Gradient, AsyncGradient, IndexingJobError, IndexingJobTimeoutError from tests.utils import assert_matches_type from gradient.types.knowledge_bases import ( IndexingJobListResponse, From e37f640eac4a864771a9a9e947703a0949f63d75 Mon Sep 17 00:00:00 2001 From: areeb1501 <64038736+areeb1501@users.noreply.github.com> Date: Thu, 23 Oct 2025 01:58:37 +0400 Subject: [PATCH 12/14] Linting Errors --- examples/agent_wait_until_ready.py | 28 +++---- examples/knowledge_base_indexing_wait.py | 74 +++++++++---------- .../knowledge_bases/test_indexing_jobs.py | 20 +++-- tests/api_resources/test_agents.py | 46 ++++++------ 4 files changed, 88 insertions(+), 80 deletions(-) diff --git a/examples/agent_wait_until_ready.py b/examples/agent_wait_until_ready.py index 3ea7b4a3..df8c8cc6 100644 --- a/examples/agent_wait_until_ready.py +++ b/examples/agent_wait_until_ready.py @@ -24,7 +24,7 @@ if agent_id: print(f"Agent created with ID: {agent_id}") print("Waiting for agent to be ready...") - + try: # Wait for the agent to be deployed and ready # This will poll the agent status every 5 seconds (default) @@ -32,24 +32,24 @@ ready_agent = client.agents.wait_until_ready( agent_id, poll_interval=5.0, # Check every 5 seconds - timeout=300.0, # Wait up to 5 minutes + timeout=300.0, # Wait up to 5 minutes ) - + if ready_agent.agent and ready_agent.agent.deployment: print(f"Agent is ready! Status: {ready_agent.agent.deployment.status}") print(f"Agent URL: {ready_agent.agent.url}") - + # Now you can use the agent # ... - + except AgentDeploymentError as e: print(f"Agent deployment failed: {e}") print(f"Failed status: {e.status}") - + except AgentDeploymentTimeoutError as e: print(f"Agent deployment timed out: {e}") print(f"Agent ID: {e.agent_id}") - + except Exception as e: print(f"Unexpected error: {e}") @@ -60,7 +60,7 @@ async def main() -> None: async_client = AsyncGradient() - + # Create a new agent agent_response = await async_client.agents.create( name="My Async Agent", @@ -68,13 +68,13 @@ async def main() -> None: model_uuid="", region="nyc1", ) - + agent_id = agent_response.agent.uuid if agent_response.agent else None - + if agent_id: print(f"Agent created with ID: {agent_id}") print("Waiting for agent to be ready...") - + try: # Wait for the agent to be deployed and ready (async) ready_agent = await async_client.agents.wait_until_ready( @@ -82,15 +82,15 @@ async def main() -> None: poll_interval=5.0, timeout=300.0, ) - + if ready_agent.agent and ready_agent.agent.deployment: print(f"Agent is ready! Status: {ready_agent.agent.deployment.status}") print(f"Agent URL: {ready_agent.agent.url}") - + except AgentDeploymentError as e: print(f"Agent deployment failed: {e}") print(f"Failed status: {e.status}") - + except AgentDeploymentTimeoutError as e: print(f"Agent deployment timed out: {e}") print(f"Agent ID: {e.agent_id}") diff --git a/examples/knowledge_base_indexing_wait.py b/examples/knowledge_base_indexing_wait.py index 3ee762bb..2917e31a 100644 --- a/examples/knowledge_base_indexing_wait.py +++ b/examples/knowledge_base_indexing_wait.py @@ -19,29 +19,27 @@ def main(): # Example 1: Basic usage - wait for indexing job to complete print("Example 1: Basic usage") print("-" * 50) - + # Create an indexing job (replace with your actual knowledge base UUID) knowledge_base_uuid = os.getenv("KNOWLEDGE_BASE_UUID", "your-kb-uuid-here") - + print(f"Creating indexing job for knowledge base: {knowledge_base_uuid}") indexing_job = client.knowledge_bases.indexing_jobs.create( knowledge_base_uuid=knowledge_base_uuid, ) - + job_uuid = indexing_job.job.uuid if indexing_job.job else None if not job_uuid: print("Error: Could not create indexing job") return - + print(f"Indexing job created with UUID: {job_uuid}") print("Waiting for indexing job to complete...") - + try: # Wait for the job to complete (polls every 5 seconds by default) - completed_job = client.knowledge_bases.indexing_jobs.wait_for_completion( - job_uuid - ) - + completed_job = client.knowledge_bases.indexing_jobs.wait_for_completion(job_uuid) + print("\n✅ Indexing job completed successfully!") if completed_job.job: print(f"Phase: {completed_job.job.phase}") @@ -49,7 +47,7 @@ def main(): print(f"Total items failed: {completed_job.job.total_items_failed}") print(f"Total datasources: {completed_job.job.total_datasources}") print(f"Completed datasources: {completed_job.job.completed_datasources}") - + except IndexingJobTimeoutError as e: print(f"\n⏱️ Timeout: {e}") except IndexingJobError as e: @@ -62,42 +60,44 @@ def example_with_custom_polling(): """Example with custom polling interval and timeout""" print("\n\nExample 2: Custom polling interval and timeout") print("-" * 50) - + client = Gradient() knowledge_base_uuid = os.getenv("KNOWLEDGE_BASE_UUID", "your-kb-uuid-here") - + print(f"Creating indexing job for knowledge base: {knowledge_base_uuid}") indexing_job = client.knowledge_bases.indexing_jobs.create( knowledge_base_uuid=knowledge_base_uuid, ) - + job_uuid = indexing_job.job.uuid if indexing_job.job else None if not job_uuid: print("Error: Could not create indexing job") return - + print(f"Indexing job created with UUID: {job_uuid}") print("Waiting for indexing job to complete (polling every 10 seconds, 5 minute timeout)...") - + try: # Wait with custom poll interval (10 seconds) and timeout (5 minutes = 300 seconds) completed_job = client.knowledge_bases.indexing_jobs.wait_for_completion( job_uuid, poll_interval=10, # Poll every 10 seconds - timeout=300, # Timeout after 5 minutes + timeout=300, # Timeout after 5 minutes ) - + print("\n✅ Indexing job completed successfully!") if completed_job.job: print(f"Phase: {completed_job.job.phase}") - + except IndexingJobTimeoutError: print("\n⏱️ Job did not complete within 5 minutes") # You can still check the current status current_status = client.knowledge_bases.indexing_jobs.retrieve(job_uuid) if current_status.job: print(f"Current phase: {current_status.job.phase}") - print(f"Completed datasources: {current_status.job.completed_datasources}/{current_status.job.total_datasources}") + print( + f"Completed datasources: {current_status.job.completed_datasources}/{current_status.job.total_datasources}" + ) except IndexingJobError as e: print(f"\n❌ Job failed: {e}") @@ -106,31 +106,31 @@ def example_manual_polling(): """Example of the old manual polling approach (for comparison)""" print("\n\nExample 3: Manual polling (old approach)") print("-" * 50) - + client = Gradient() knowledge_base_uuid = os.getenv("KNOWLEDGE_BASE_UUID", "your-kb-uuid-here") - + indexing_job = client.knowledge_bases.indexing_jobs.create( knowledge_base_uuid=knowledge_base_uuid, ) - + job_uuid = indexing_job.job.uuid if indexing_job.job else None if not job_uuid: print("Error: Could not create indexing job") return - + print(f"Indexing job created with UUID: {job_uuid}") print("Manual polling (old approach)...") - + import time - + while True: indexing_job = client.knowledge_bases.indexing_jobs.retrieve(job_uuid) - + if indexing_job.job and indexing_job.job.phase: phase = indexing_job.job.phase print(f"Current phase: {phase}") - + if phase in ["BATCH_JOB_PHASE_UNKNOWN", "BATCH_JOB_PHASE_PENDING", "BATCH_JOB_PHASE_RUNNING"]: time.sleep(5) continue @@ -146,36 +146,36 @@ async def example_async(): """Example using async/await""" print("\n\nExample 4: Async usage") print("-" * 50) - + from gradient import AsyncGradient - + client = AsyncGradient() knowledge_base_uuid = os.getenv("KNOWLEDGE_BASE_UUID", "your-kb-uuid-here") - + print(f"Creating indexing job for knowledge base: {knowledge_base_uuid}") indexing_job = await client.knowledge_bases.indexing_jobs.create( knowledge_base_uuid=knowledge_base_uuid, ) - + job_uuid = indexing_job.job.uuid if indexing_job.job else None if not job_uuid: print("Error: Could not create indexing job") return - + print(f"Indexing job created with UUID: {job_uuid}") print("Waiting for indexing job to complete (async)...") - + try: completed_job = await client.knowledge_bases.indexing_jobs.wait_for_completion( job_uuid, poll_interval=5, timeout=600, # 10 minute timeout ) - + print("\n✅ Indexing job completed successfully!") if completed_job.job: print(f"Phase: {completed_job.job.phase}") - + except IndexingJobTimeoutError as e: print(f"\n⏱️ Timeout: {e}") except IndexingJobError as e: @@ -187,11 +187,11 @@ async def example_async(): if __name__ == "__main__": # Run the basic example main() - + # Uncomment to run other examples: # example_with_custom_polling() # example_manual_polling() - + # For async example, you would need to run: # import asyncio # asyncio.run(example_async()) diff --git a/tests/api_resources/knowledge_bases/test_indexing_jobs.py b/tests/api_resources/knowledge_bases/test_indexing_jobs.py index 216f09e0..53e5ee1b 100644 --- a/tests/api_resources/knowledge_bases/test_indexing_jobs.py +++ b/tests/api_resources/knowledge_bases/test_indexing_jobs.py @@ -282,7 +282,9 @@ def test_wait_for_completion_raises_indexing_job_error_on_error(self, client: Gr assert "error" in str(exc_info.value).lower() @parametrize - def test_wait_for_completion_raises_indexing_job_error_on_cancelled(self, client: Gradient, respx_mock: Any) -> None: + def test_wait_for_completion_raises_indexing_job_error_on_cancelled( + self, client: Gradient, respx_mock: Any + ) -> None: """Test that IndexingJobError is raised when job phase is CANCELLED""" job_uuid = "test-job-uuid" respx_mock.get(f"{base_url}/v2/gen-ai/indexing_jobs/{job_uuid}").mock( @@ -321,9 +323,7 @@ def test_wait_for_completion_raises_timeout_error(self, client: Gradient, respx_ ) with pytest.raises(IndexingJobTimeoutError) as exc_info: - client.knowledge_bases.indexing_jobs.wait_for_completion( - job_uuid, poll_interval=0.1, timeout=0.2 - ) + client.knowledge_bases.indexing_jobs.wait_for_completion(job_uuid, poll_interval=0.1, timeout=0.2) assert exc_info.value.uuid == job_uuid assert exc_info.value.phase == "BATCH_JOB_PHASE_RUNNING" @@ -567,7 +567,9 @@ async def test_path_params_update_cancel(self, async_client: AsyncGradient) -> N ) @parametrize - async def test_wait_for_completion_raises_indexing_job_error_on_failed(self, async_client: AsyncGradient, respx_mock: Any) -> None: + async def test_wait_for_completion_raises_indexing_job_error_on_failed( + self, async_client: AsyncGradient, respx_mock: Any + ) -> None: """Test that IndexingJobError is raised when job phase is FAILED""" job_uuid = "test-job-uuid" respx_mock.get(f"{base_url}/v2/gen-ai/indexing_jobs/{job_uuid}").mock( @@ -592,7 +594,9 @@ async def test_wait_for_completion_raises_indexing_job_error_on_failed(self, asy assert "failed" in str(exc_info.value).lower() @parametrize - async def test_wait_for_completion_raises_indexing_job_error_on_error(self, async_client: AsyncGradient, respx_mock: Any) -> None: + async def test_wait_for_completion_raises_indexing_job_error_on_error( + self, async_client: AsyncGradient, respx_mock: Any + ) -> None: """Test that IndexingJobError is raised when job phase is ERROR""" job_uuid = "test-job-uuid" respx_mock.get(f"{base_url}/v2/gen-ai/indexing_jobs/{job_uuid}").mock( @@ -615,7 +619,9 @@ async def test_wait_for_completion_raises_indexing_job_error_on_error(self, asyn assert "error" in str(exc_info.value).lower() @parametrize - async def test_wait_for_completion_raises_indexing_job_error_on_cancelled(self, async_client: AsyncGradient, respx_mock: Any) -> None: + async def test_wait_for_completion_raises_indexing_job_error_on_cancelled( + self, async_client: AsyncGradient, respx_mock: Any + ) -> None: """Test that IndexingJobError is raised when job phase is CANCELLED""" job_uuid = "test-job-uuid" respx_mock.get(f"{base_url}/v2/gen-ai/indexing_jobs/{job_uuid}").mock( diff --git a/tests/api_resources/test_agents.py b/tests/api_resources/test_agents.py index 5777c3ea..1ba3e093 100644 --- a/tests/api_resources/test_agents.py +++ b/tests/api_resources/test_agents.py @@ -368,9 +368,10 @@ def test_path_params_update_status(self, client: Gradient) -> None: def test_method_wait_until_ready(self, client: Gradient, respx_mock: Any) -> None: """Test successful wait_until_ready when agent becomes ready.""" agent_uuid = "test-agent-id" - + # Create side effect that returns different responses call_count = [0] + def get_response(_: httpx.Request) -> httpx.Response: call_count[0] += 1 if call_count[0] == 1: @@ -395,9 +396,9 @@ def get_response(_: httpx.Request) -> httpx.Response: } }, ) - + respx_mock.get(f"/v2/gen-ai/agents/{agent_uuid}").mock(side_effect=get_response) - + agent = client.agents.wait_until_ready(agent_uuid, poll_interval=0.1, timeout=10.0) assert_matches_type(AgentRetrieveResponse, agent, path=["response"]) assert agent.agent is not None @@ -408,9 +409,9 @@ def get_response(_: httpx.Request) -> httpx.Response: def test_wait_until_ready_timeout(self, client: Gradient, respx_mock: Any) -> None: """Test that wait_until_ready raises timeout error.""" from gradient._exceptions import AgentDeploymentTimeoutError - + agent_uuid = "test-agent-id" - + # Mock always returns deploying respx_mock.get(f"/v2/gen-ai/agents/{agent_uuid}").mock( return_value=httpx.Response( @@ -423,10 +424,10 @@ def test_wait_until_ready_timeout(self, client: Gradient, respx_mock: Any) -> No }, ) ) - + with pytest.raises(AgentDeploymentTimeoutError) as exc_info: client.agents.wait_until_ready(agent_uuid, poll_interval=0.1, timeout=0.5) - + assert "did not reach STATUS_RUNNING within" in str(exc_info.value) assert exc_info.value.agent_id == agent_uuid @@ -434,9 +435,9 @@ def test_wait_until_ready_timeout(self, client: Gradient, respx_mock: Any) -> No def test_wait_until_ready_deployment_failed(self, client: Gradient, respx_mock: Any) -> None: """Test that wait_until_ready raises error on deployment failure.""" from gradient._exceptions import AgentDeploymentError - + agent_uuid = "test-agent-id" - + # Mock returns failed status respx_mock.get(f"/v2/gen-ai/agents/{agent_uuid}").mock( return_value=httpx.Response( @@ -449,10 +450,10 @@ def test_wait_until_ready_deployment_failed(self, client: Gradient, respx_mock: }, ) ) - + with pytest.raises(AgentDeploymentError) as exc_info: client.agents.wait_until_ready(agent_uuid, poll_interval=0.1, timeout=10.0) - + assert "deployment failed with status: STATUS_FAILED" in str(exc_info.value) assert exc_info.value.status == "STATUS_FAILED" @@ -810,9 +811,10 @@ async def test_path_params_update_status(self, async_client: AsyncGradient) -> N async def test_method_wait_until_ready(self, async_client: AsyncGradient, respx_mock: Any) -> None: """Test successful async wait_until_ready when agent becomes ready.""" agent_uuid = "test-agent-id" - + # Create side effect that returns different responses call_count = [0] + def get_response(_: httpx.Request) -> httpx.Response: call_count[0] += 1 if call_count[0] == 1: @@ -837,9 +839,9 @@ def get_response(_: httpx.Request) -> httpx.Response: } }, ) - + respx_mock.get(f"/v2/gen-ai/agents/{agent_uuid}").mock(side_effect=get_response) - + agent = await async_client.agents.wait_until_ready(agent_uuid, poll_interval=0.1, timeout=10.0) assert_matches_type(AgentRetrieveResponse, agent, path=["response"]) assert agent.agent is not None @@ -850,9 +852,9 @@ def get_response(_: httpx.Request) -> httpx.Response: async def test_wait_until_ready_timeout(self, async_client: AsyncGradient, respx_mock: Any) -> None: """Test that async wait_until_ready raises timeout error.""" from gradient._exceptions import AgentDeploymentTimeoutError - + agent_uuid = "test-agent-id" - + # Mock always returns deploying respx_mock.get(f"/v2/gen-ai/agents/{agent_uuid}").mock( return_value=httpx.Response( @@ -865,10 +867,10 @@ async def test_wait_until_ready_timeout(self, async_client: AsyncGradient, respx }, ) ) - + with pytest.raises(AgentDeploymentTimeoutError) as exc_info: await async_client.agents.wait_until_ready(agent_uuid, poll_interval=0.1, timeout=0.5) - + assert "did not reach STATUS_RUNNING within" in str(exc_info.value) assert exc_info.value.agent_id == agent_uuid @@ -876,9 +878,9 @@ async def test_wait_until_ready_timeout(self, async_client: AsyncGradient, respx async def test_wait_until_ready_deployment_failed(self, async_client: AsyncGradient, respx_mock: Any) -> None: """Test that async wait_until_ready raises error on deployment failure.""" from gradient._exceptions import AgentDeploymentError - + agent_uuid = "test-agent-id" - + # Mock returns failed status respx_mock.get(f"/v2/gen-ai/agents/{agent_uuid}").mock( return_value=httpx.Response( @@ -891,9 +893,9 @@ async def test_wait_until_ready_deployment_failed(self, async_client: AsyncGradi }, ) ) - + with pytest.raises(AgentDeploymentError) as exc_info: await async_client.agents.wait_until_ready(agent_uuid, poll_interval=0.1, timeout=10.0) - + assert "deployment failed with status: STATUS_FAILED" in str(exc_info.value) assert exc_info.value.status == "STATUS_FAILED" From a413919729bcdd929991e5015609aa9fba313bb0 Mon Sep 17 00:00:00 2001 From: areeb1501 <64038736+areeb1501@users.noreply.github.com> Date: Thu, 23 Oct 2025 20:28:52 +0400 Subject: [PATCH 13/14] fixing linting errors --- tests/api_resources/knowledge_bases/test_indexing_jobs.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/tests/api_resources/knowledge_bases/test_indexing_jobs.py b/tests/api_resources/knowledge_bases/test_indexing_jobs.py index 53e5ee1b..88c551c8 100644 --- a/tests/api_resources/knowledge_bases/test_indexing_jobs.py +++ b/tests/api_resources/knowledge_bases/test_indexing_jobs.py @@ -349,6 +349,7 @@ def test_wait_for_completion_succeeds(self, client: Gradient, respx_mock: Any) - result = client.knowledge_bases.indexing_jobs.wait_for_completion(job_uuid) assert_matches_type(IndexingJobRetrieveResponse, result, path=["response"]) + assert result.job is not None assert result.job.phase == "BATCH_JOB_PHASE_SUCCEEDED" @@ -688,4 +689,5 @@ async def test_wait_for_completion_succeeds(self, async_client: AsyncGradient, r result = await async_client.knowledge_bases.indexing_jobs.wait_for_completion(job_uuid) assert_matches_type(IndexingJobRetrieveResponse, result, path=["response"]) + assert result.job is not None assert result.job.phase == "BATCH_JOB_PHASE_SUCCEEDED" From 5df83227eaf3f6c67668489cae680d17f33b03c9 Mon Sep 17 00:00:00 2001 From: areeb1501 <64038736+areeb1501@users.noreply.github.com> Date: Fri, 24 Oct 2025 20:53:10 +0400 Subject: [PATCH 14/14] Added return types where applicable --- examples/knowledge_base_indexing_wait.py | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/examples/knowledge_base_indexing_wait.py b/examples/knowledge_base_indexing_wait.py index 2917e31a..1171fea3 100644 --- a/examples/knowledge_base_indexing_wait.py +++ b/examples/knowledge_base_indexing_wait.py @@ -12,7 +12,7 @@ from gradient import Gradient, IndexingJobError, IndexingJobTimeoutError -def main(): +def main() -> None: # Initialize the Gradient client client = Gradient() @@ -56,7 +56,7 @@ def main(): print(f"\n❌ Unexpected error: {e}") -def example_with_custom_polling(): +def example_with_custom_polling() -> None: """Example with custom polling interval and timeout""" print("\n\nExample 2: Custom polling interval and timeout") print("-" * 50) @@ -102,7 +102,7 @@ def example_with_custom_polling(): print(f"\n❌ Job failed: {e}") -def example_manual_polling(): +def example_manual_polling() -> None: """Example of the old manual polling approach (for comparison)""" print("\n\nExample 3: Manual polling (old approach)") print("-" * 50) @@ -125,10 +125,10 @@ def example_manual_polling(): import time while True: - indexing_job = client.knowledge_bases.indexing_jobs.retrieve(job_uuid) + indexing_job_status = client.knowledge_bases.indexing_jobs.retrieve(job_uuid) - if indexing_job.job and indexing_job.job.phase: - phase = indexing_job.job.phase + if indexing_job_status.job and indexing_job_status.job.phase: + phase = indexing_job_status.job.phase print(f"Current phase: {phase}") if phase in ["BATCH_JOB_PHASE_UNKNOWN", "BATCH_JOB_PHASE_PENDING", "BATCH_JOB_PHASE_RUNNING"]: @@ -142,7 +142,7 @@ def example_manual_polling(): break -async def example_async(): +async def example_async() -> None: """Example using async/await""" print("\n\nExample 4: Async usage") print("-" * 50)