feat: add wait_until_database_online helper for knowledge base databa…

examples/agent_wait_until_ready.py

@@ -24,32 +24,32 @@
 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)
         # and wait up to 5 minutes (default timeout=300 seconds)
         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,37 +60,37 @@
 
 async def main() -> None:
     async_client = AsyncGradient()
-    
+
     # Create a new agent
     agent_response = await async_client.agents.create(
         name="My Async Agent",
         instruction="You are a helpful assistant",
         model_uuid="<your-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(
                 agent_id,
                 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}")

src/gradient/_exceptions.py

@@ -15,6 +15,8 @@
     "UnprocessableEntityError",
     "RateLimitError",
     "InternalServerError",
+    "KnowledgeBaseDatabaseError",
+    "KnowledgeBaseDatabaseTimeoutError",
     "AgentDeploymentError",
     "AgentDeploymentTimeoutError",
 ]
@@ -124,3 +126,19 @@ class AgentDeploymentTimeoutError(GradientError):
     def __init__(self, message: str, agent_id: str) -> None:
         super().__init__(message)
         self.agent_id = agent_id
+
+
+class KnowledgeBaseDatabaseError(GradientError):
+    """Raised when a knowledge base database creation fails."""
+
+    def __init__(self, message: str, status: str) -> None:
+        super().__init__(message)
+        self.status = status
+
+
+class KnowledgeBaseDatabaseTimeoutError(GradientError):
+    """Raised when waiting for a knowledge base database creation times out."""
+
+    def __init__(self, message: str, knowledge_base_id: str) -> None:
+        super().__init__(message)
+        self.knowledge_base_id = knowledge_base_id

src/gradient/resources/knowledge_bases/knowledge_bases.py

@@ -2,6 +2,7 @@
 
 from __future__ import annotations
 
+import time
 from typing import Iterable
 
 import httpx
@@ -330,6 +331,56 @@ def delete(
             cast_to=KnowledgeBaseDeleteResponse,
         )
 
+    def wait_until_database_online(
+        self,
+        uuid: str,
+        *,
+        timeout: float = 300.0,
+        poll_interval: float = 5.0,
+        extra_headers: Headers | None = None,
+        extra_query: Query | None = None,
+        extra_body: Body | None = None,
+    ) -> KnowledgeBaseRetrieveResponse:
+        """Wait for a knowledge base's associated database to reach ONLINE.
+
+        This polls `retrieve` until `database_status` equals "ONLINE", or raises
+        on terminal failure or timeout.
+        """
+        from ..._exceptions import KnowledgeBaseDatabaseError, KnowledgeBaseDatabaseTimeoutError
+
+        if not uuid:
+            raise ValueError(f"Expected a non-empty value for `uuid` but received {uuid!r}")
+
+        start_time = time.time()
+
+        while True:
+            kb_response = self.retrieve(
+                uuid, extra_headers=extra_headers, extra_query=extra_query, extra_body=extra_body
+            )
+
+            status = kb_response.database_status if kb_response else None
+
+            # Success
+            if status == "ONLINE":
+                return kb_response
+
+            # Failure cases - treat some terminal statuses as failures
+            if status in ("DECOMMISSIONED", "UNHEALTHY"):
+                raise KnowledgeBaseDatabaseError(
+                    f"Knowledge base database creation failed with status: {status}", status=status
+                )
+
+            # Timeout
+            elapsed_time = time.time() - start_time
+            if elapsed_time >= timeout:
+                current_status = status or "UNKNOWN"
+                raise KnowledgeBaseDatabaseTimeoutError(
+                    f"Knowledge base database did not reach ONLINE within {timeout} seconds. Current status: {current_status}",
+                    knowledge_base_id=uuid,
+                )
+
+            time.sleep(poll_interval)
+
 
 class AsyncKnowledgeBasesResource(AsyncAPIResource):
     @cached_property
@@ -618,6 +669,51 @@ async def delete(
             cast_to=KnowledgeBaseDeleteResponse,
         )
 
+    async def wait_until_database_online(
+        self,
+        uuid: str,
+        *,
+        timeout: float = 300.0,
+        poll_interval: float = 5.0,
+        extra_headers: Headers | None = None,
+        extra_query: Query | None = None,
+        extra_body: Body | None = None,
+    ) -> KnowledgeBaseRetrieveResponse:
+        """Async version of `wait_until_database_online`."""
+        import asyncio
+
+        from ..._exceptions import KnowledgeBaseDatabaseError, KnowledgeBaseDatabaseTimeoutError
+
+        if not uuid:
+            raise ValueError(f"Expected a non-empty value for `uuid` but received {uuid!r}")
+
+        start_time = time.time()
+
+        while True:
+            kb_response = await self.retrieve(
+                uuid, extra_headers=extra_headers, extra_query=extra_query, extra_body=extra_body
+            )
+
+            status = kb_response.database_status if kb_response else None
+
+            if status == "ONLINE":
+                return kb_response
+
+            if status in ("DECOMMISSIONED", "UNHEALTHY"):
+                raise KnowledgeBaseDatabaseError(
+                    f"Knowledge base database creation failed with status: {status}", status=status
+                )
+
+            elapsed_time = time.time() - start_time
+            if elapsed_time >= timeout:
+                current_status = status or "UNKNOWN"
+                raise KnowledgeBaseDatabaseTimeoutError(
+                    f"Knowledge base database did not reach ONLINE within {timeout} seconds. Current status: {current_status}",
+                    knowledge_base_id=uuid,
+                )
+
+            await asyncio.sleep(poll_interval)
+
 
 class KnowledgeBasesResourceWithRawResponse:
     def __init__(self, knowledge_bases: KnowledgeBasesResource) -> None:
@@ -638,6 +734,9 @@ def __init__(self, knowledge_bases: KnowledgeBasesResource) -> None:
         self.delete = to_raw_response_wrapper(
             knowledge_bases.delete,
         )
+        self.wait_until_database_online = to_raw_response_wrapper(
+            knowledge_bases.wait_until_database_online,
+        )
 
     @cached_property
     def data_sources(self) -> DataSourcesResourceWithRawResponse:
@@ -667,6 +766,9 @@ def __init__(self, knowledge_bases: AsyncKnowledgeBasesResource) -> None:
         self.delete = async_to_raw_response_wrapper(
             knowledge_bases.delete,
         )
+        self.wait_until_database_online = async_to_raw_response_wrapper(
+            knowledge_bases.wait_until_database_online,
+        )
 
     @cached_property
     def data_sources(self) -> AsyncDataSourcesResourceWithRawResponse:
@@ -696,6 +798,9 @@ def __init__(self, knowledge_bases: KnowledgeBasesResource) -> None:
         self.delete = to_streamed_response_wrapper(
             knowledge_bases.delete,
         )
+        self.wait_until_database_online = to_streamed_response_wrapper(
+            knowledge_bases.wait_until_database_online,
+        )
 
     @cached_property
     def data_sources(self) -> DataSourcesResourceWithStreamingResponse:
@@ -725,6 +830,9 @@ def __init__(self, knowledge_bases: AsyncKnowledgeBasesResource) -> None:
         self.delete = async_to_streamed_response_wrapper(
             knowledge_bases.delete,
         )
+        self.wait_until_database_online = async_to_streamed_response_wrapper(
+            knowledge_bases.wait_until_database_online,
+        )
 
     @cached_property
     def data_sources(self) -> AsyncDataSourcesResourceWithStreamingResponse: