Add wait_for_database helper function to poll for Knowledge Base database creation

[Copilot]

Overview

This PR adds a wait_for_database() helper function to simplify polling for Knowledge Base database status during creation. Previously, users had to manually write polling loops to wait for the database to become ready, which was error-prone and required handling multiple edge cases.

Problem

When creating a Knowledge Base, the database deployment can take several minutes. Users previously had to write manual polling loops like:

from gradient import Gradient
import time

client = Gradient()

while True:
    knowledge_base = client.knowledge_bases.retrieve("uuid")
    if knowledge_base.database_status == "ONLINE":
        break
    # Manual error handling, timeout tracking, etc.
    time.sleep(5)

This approach has several issues:

• No timeout handling
• No proper error handling for failed states
• Verbose and repetitive code
• Easy to miss edge cases

Solution

Added a wait_for_database() helper method that:

• Automatically polls the database status until it becomes ONLINE
• Handles failed states (DECOMMISSIONED, UNHEALTHY) with a custom KnowledgeBaseDatabaseError
• Supports configurable timeout (default: 600 seconds) and poll interval (default: 5 seconds)
• Raises APITimeoutError if the database doesn't become ready within the timeout
• Works for both synchronous and asynchronous clients

Usage

from gradient import Gradient
from gradient.resources.knowledge_bases import KnowledgeBaseDatabaseError
from gradient._exceptions import APITimeoutError

client = Gradient()

# Create a knowledge base
kb = client.knowledge_bases.create(
    name="My Knowledge Base",
    region="nyc1",
    embedding_model_uuid="your-embedding-model-uuid",
)

kb_uuid = kb.knowledge_base.uuid

try:
    # Wait for the database to be ready (default: 10 minute timeout, 5 second polling)
    result = client.knowledge_bases.wait_for_database(kb_uuid)
    print(f"Database is ready: {result.database_status}")
    
    # Or with custom parameters
    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")

Async Usage

from gradient import AsyncGradient

async_client = AsyncGradient()

result = await async_client.knowledge_bases.wait_for_database(kb_uuid)

Changes

• New method: KnowledgeBasesResource.wait_for_database() (synchronous)
• New method: AsyncKnowledgeBasesResource.wait_for_database() (asynchronous)
• New exception: KnowledgeBaseDatabaseError for failed database states
• Tests: 20 comprehensive unit tests covering success scenarios, failed states, timeouts, and parameter validation
• Documentation: Updated README with usage examples and error handling patterns

API Reference

def wait_for_database(
    uuid: str,
    *,
    timeout: float = 600.0,
    poll_interval: float = 5.0,
    extra_headers: Headers | None = None,
    extra_query: Query | None = None,
    extra_body: Body | None = None,
) -> KnowledgeBaseRetrieveResponse

Parameters:

• uuid - Knowledge base UUID to poll (required)
• timeout - Maximum wait time in seconds (default: 600)
• poll_interval - Time between status checks in seconds (default: 5)

Returns: KnowledgeBaseRetrieveResponse when database status is ONLINE

Raises:

• KnowledgeBaseDatabaseError - Database entered a failed state (DECOMMISSIONED or UNHEALTHY)
• APITimeoutError - Timeout exceeded before database became ONLINE
• ValueError - Invalid UUID parameter

Testing

All tests pass successfully:

pytest tests/api_resources/test_knowledge_bases.py -k "wait_for_database" -v
# 20 passed (10 sync + 10 async tests)

Tests cover:

• Successful database creation and polling
• Failed database states (UNHEALTHY, DECOMMISSIONED)
• Timeout scenarios
• Parameter validation
• Both loose and strict client modes

Backward Compatibility

✅ No breaking changes - this is a purely additive feature. Existing code continues to work unchanged, and the new helper is opt-in.

Original prompt

analys this full project. and Create a helper function to poll for knowledge base database creation.

When a Gradient AI Platform Knowledge Base is created the deployment of the database can run for several minutes. Currently users must write loops to poll for the deployment state, roughly.

from gradient import Gradient

client = Gradient()
while True:
knowledge_base = client.knowledge_bases.retrieve(
"uuid",
)
if knowledge_base.database_status != "ONLINE"
continue
This should of course handle failed states for errors and such

. it was an issue raised on the repo. do fix this full and implement what ever needed and make it perfect. test everything and double check everything.
analys this full project. and Create a helper function to poll for knowledge base database creation.

When a Gradient AI Platform Knowledge Base is created the deployment of the database can run for several minutes. Currently users must write loops to poll for the deployment state, roughly.

from gradient import Gradient

client = Gradient()
while True:
knowledge_base = client.knowledge_bases.retrieve(
"uuid",
)
if knowledge_base.database_status != "ONLINE"
continue
This should of course handle failed states for errors and such

. it was an issue raised on the repo. do fix this full and implement what ever needed and make it perfect. test everything and double check everything.

Created from VS Code via the [GitHub Pull Request](https://marketplace.visualstudio.com/items?itemName=GitHub.vscode-pull-request-github) extension.

---

✨ Let Copilot coding agent set things up for you — coding agent works faster and does higher quality work when set up for your repo.