Address review feedback and add KnowledgeBaseTimeoutError

Author: kashyapdayal

README.md

@@ -96,49 +96,6 @@ we recommend using [python-dotenv](https://pypi.org/project/python-dotenv/)
 to add `DIGITALOCEAN_ACCESS_TOKEN="My Access Token"`, `GRADIENT_MODEL_ACCESS_KEY="My Model Access Key"` to your `.env` file
 so that your keys are not stored in source control.
 
-## Knowledge Base Database Polling
-
-When creating a Knowledge Base, the database deployment can take several minutes. The `wait_for_database()` helper function simplifies polling for the database status:
-
-```python
-from gradient import Gradient
-from gradient.resources.knowledge_bases import KnowledgeBaseDatabaseError
-from gradient._exceptions import APITimeoutError
-
-client = Gradient()
-
-# Create a knowledge base
-kb_response = client.knowledge_bases.create(
-    name="My Knowledge Base",
-    region="nyc1",
-    embedding_model_uuid="your-embedding-model-uuid",
-)
-
-kb_uuid = kb_response.knowledge_base.uuid
-
-try:
-    # Wait for the database to be ready (default: 10 minute timeout, 5 second poll interval)
-    result = client.knowledge_bases.wait_for_database(kb_uuid)
-    print(f"Database status: {result.database_status}")  # "ONLINE"
-    
-    # Custom timeout and poll interval
-    result = client.knowledge_bases.wait_for_database(
-        kb_uuid,
-        timeout=900.0,       # 15 minutes
-        poll_interval=10.0   # Check every 10 seconds
-    )
-    
-except KnowledgeBaseDatabaseError as e:
-    # Database entered a failed state (DECOMMISSIONED or UNHEALTHY)
-    print(f"Database failed: {e}")
-    
-except APITimeoutError:
-    # Database did not become ready within the timeout period
-    print("Timeout: Database did not become ready in time")
-```
-
-The helper handles all state transitions and will raise appropriate exceptions for failed states or timeouts.
-
 ## Async usage
 
 Simply import `AsyncGradient` instead of `Gradient` and use `await` with each API call:

examples/wait_for_knowledge_base.py

@@ -0,0 +1,60 @@
+"""
+Example demonstrating how to use the wait_for_database helper function.
+
+This example shows how to:
+1. Create a knowledge base
+2. Wait for its database to be ready
+3. Handle errors and timeouts appropriately
+"""
+
+import os
+
+from gradient import Gradient
+from gradient.resources.knowledge_bases import KnowledgeBaseTimeoutError, KnowledgeBaseDatabaseError
+
+
+def main() -> None:
+    """Create a knowledge base and wait for its database to be ready."""
+    # Initialize the Gradient client
+    # Note: DIGITALOCEAN_ACCESS_TOKEN must be set in your environment
+    client = Gradient(
+        access_token=os.environ.get("DIGITALOCEAN_ACCESS_TOKEN"),
+    )
+
+    # Create a knowledge base
+    # Replace these values with your actual configuration
+    kb_response = client.knowledge_bases.create(
+        name="My Knowledge Base",
+        region="nyc1",  # Choose your preferred region
+        embedding_model_uuid="your-embedding-model-uuid",  # Use your embedding model UUID
+    )
+
+    kb_uuid = kb_response.knowledge_base.uuid
+    print(f"Created knowledge base: {kb_uuid}")
+
+    try:
+        # Wait for the database to be ready
+        # Default: 10 minute timeout, 5 second poll interval
+        print("Waiting for database to be ready...")
+        result = client.knowledge_bases.wait_for_database(kb_uuid)
+        print(f"Database status: {result.database_status}")  # "ONLINE"
+        print("Knowledge base is ready!")
+
+        # Alternative: Custom timeout and poll interval
+        # result = client.knowledge_bases.wait_for_database(
+        #     kb_uuid,
+        #     timeout=900.0,       # 15 minutes
+        #     poll_interval=10.0   # Check every 10 seconds
+        # )
+
+    except KnowledgeBaseDatabaseError as e:
+        # Database entered a failed state (DECOMMISSIONED or UNHEALTHY)
+        print(f"Database failed: {e}")
+
+    except KnowledgeBaseTimeoutError as e:
+        # Database did not become ready within the timeout period
+        print(f"Timeout: {e}")
+
+
+if __name__ == "__main__":
+    main()

pyproject.toml

@@ -253,5 +253,4 @@ known-first-party = ["gradient", "tests"]
 [tool.ruff.lint.per-file-ignores]
 "bin/**.py" = ["T201", "T203"]
 "scripts/**.py" = ["T201", "T203"]
-"tests/**.py" = ["T201", "T203", "ARG001"]
 "examples/**.py" = ["T201", "T203"]

src/gradient/resources/knowledge_bases/__init__.py

@@ -18,8 +18,9 @@
 )
 from .knowledge_bases import (
     KnowledgeBasesResource,
-    AsyncKnowledgeBasesResource,
+    KnowledgeBaseTimeoutError,
     KnowledgeBaseDatabaseError,
+    AsyncKnowledgeBasesResource,
     KnowledgeBasesResourceWithRawResponse,
     AsyncKnowledgeBasesResourceWithRawResponse,
     KnowledgeBasesResourceWithStreamingResponse,
@@ -42,6 +43,7 @@
     "KnowledgeBasesResource",
     "AsyncKnowledgeBasesResource",
     "KnowledgeBaseDatabaseError",
+    "KnowledgeBaseTimeoutError",
     "KnowledgeBasesResourceWithRawResponse",
     "AsyncKnowledgeBasesResourceWithRawResponse",
     "KnowledgeBasesResourceWithStreamingResponse",

src/gradient/resources/knowledge_bases/knowledge_bases.py

@@ -27,7 +27,6 @@
     DataSourcesResourceWithStreamingResponse,
     AsyncDataSourcesResourceWithStreamingResponse,
 )
-from ..._exceptions import APITimeoutError
 from .indexing_jobs import (
     IndexingJobsResource,
     AsyncIndexingJobsResource,
@@ -43,7 +42,12 @@
 from ...types.knowledge_base_update_response import KnowledgeBaseUpdateResponse
 from ...types.knowledge_base_retrieve_response import KnowledgeBaseRetrieveResponse
 
-__all__ = ["KnowledgeBasesResource", "AsyncKnowledgeBasesResource", "KnowledgeBaseDatabaseError"]
+__all__ = [
+    "KnowledgeBasesResource",
+    "AsyncKnowledgeBasesResource",
+    "KnowledgeBaseDatabaseError",
+    "KnowledgeBaseTimeoutError",
+]
 
 
 class KnowledgeBaseDatabaseError(Exception):
@@ -52,6 +56,12 @@ class KnowledgeBaseDatabaseError(Exception):
     pass
 
 
+class KnowledgeBaseTimeoutError(Exception):
+    """Raised when waiting for a knowledge base database times out."""
+
+    pass
+
+
 class KnowledgeBasesResource(SyncAPIResource):
     @cached_property
     def data_sources(self) -> DataSourcesResource:
@@ -377,7 +387,7 @@ def wait_for_database(
         Raises:
           KnowledgeBaseDatabaseError: If the database enters a failed state (DECOMMISSIONED, UNHEALTHY)
 
-          APITimeoutError: If the timeout is exceeded before the database becomes ONLINE
+          KnowledgeBaseTimeoutError: If the timeout is exceeded before the database becomes ONLINE
         """
         if not uuid:
             raise ValueError(f"Expected a non-empty value for `uuid` but received {uuid!r}")
@@ -388,11 +398,9 @@ def wait_for_database(
         while True:
             elapsed = time.time() - start_time
             if elapsed >= timeout:
-                raise APITimeoutError(
-                    request=httpx.Request(
-                        method="GET",
-                        url=f"https://api.digitalocean.com/v2/gen-ai/knowledge_bases/{uuid}",
-                    )
+                raise KnowledgeBaseTimeoutError(
+                    f"Timeout waiting for knowledge base database to become ready. "
+                    f"Database did not reach ONLINE status within {timeout} seconds."
                 )
 
             response = self.retrieve(
@@ -408,9 +416,7 @@ def wait_for_database(
                 return response
 
             if status in failed_states:
-                raise KnowledgeBaseDatabaseError(
-                    f"Knowledge base database entered failed state: {status}"
-                )
+                raise KnowledgeBaseDatabaseError(f"Knowledge base database entered failed state: {status}")
 
             # Sleep before next poll, but don't exceed timeout
             remaining_time = timeout - elapsed
@@ -744,7 +750,7 @@ async def wait_for_database(
         Raises:
           KnowledgeBaseDatabaseError: If the database enters a failed state (DECOMMISSIONED, UNHEALTHY)
 
-          APITimeoutError: If the timeout is exceeded before the database becomes ONLINE
+          KnowledgeBaseTimeoutError: If the timeout is exceeded before the database becomes ONLINE
         """
         if not uuid:
             raise ValueError(f"Expected a non-empty value for `uuid` but received {uuid!r}")
@@ -755,11 +761,9 @@ async def wait_for_database(
         while True:
             elapsed = time.time() - start_time
             if elapsed >= timeout:
-                raise APITimeoutError(
-                    request=httpx.Request(
-                        method="GET",
-                        url=f"https://api.digitalocean.com/v2/gen-ai/knowledge_bases/{uuid}",
-                    )
+                raise KnowledgeBaseTimeoutError(
+                    f"Timeout waiting for knowledge base database to become ready. "
+                    f"Database did not reach ONLINE status within {timeout} seconds."
                 )
 
             response = await self.retrieve(
@@ -775,9 +779,7 @@ async def wait_for_database(
                 return response
 
             if status in failed_states:
-                raise KnowledgeBaseDatabaseError(
-                    f"Knowledge base database entered failed state: {status}"
-                )
+                raise KnowledgeBaseDatabaseError(f"Knowledge base database entered failed state: {status}")
 
             # Sleep before next poll, but don't exceed timeout
             remaining_time = timeout - elapsed