feat: add wait_until_database_online helper for knowledge base databa…

[cipheraxat]

Overview

This PR implements wait_until_database_online() method for polling knowledge base database creation status, following the established pattern of agents.wait_until_ready().

Changes Made

1. New Exception Classes (src/gradient/_exceptions.py)

Added two new exception types for database polling:

• KnowledgeBaseDatabaseError: Raised when database creation fails with terminal status
• KnowledgeBaseDatabaseTimeoutError: Raised when polling exceeds timeout limit

Both exceptions store relevant identifiers for debugging:

• status attribute on KnowledgeBaseDatabaseError
• knowledge_base_id attribute on KnowledgeBaseDatabaseTimeoutError

2. Resource Methods (src/gradient/resources/knowledge_bases/knowledge_bases.py)

Synchronous Method

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

Asynchronous Method

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

Behavior:

• Polls retrieve(uuid) until database_status == "ONLINE"
• Success states: Returns KnowledgeBaseRetrieveResponse
• Failure states: Raises KnowledgeBaseDatabaseError on terminal status ("DECOMMISSIONED", "UNHEALTHY")
• Timeout: Raises KnowledgeBaseDatabaseTimeoutError if exceeded
• Validation: Raises ValueError on empty UUID

3. Wrapper Classes

Updated all four wrapper classes to expose the new method:

• KnowledgeBasesResourceWithRawResponse
• AsyncKnowledgeBasesResourceWithRawResponse
• KnowledgeBasesResourceWithStreamingResponse
• AsyncKnowledgeBasesResourceWithStreamingResponse

Enables usage patterns:

# Raw response wrapper
response = client.knowledge_bases.with_raw_response.wait_until_database_online(uuid)

# Streaming response wrapper
async with client.knowledge_bases.with_streaming_response.wait_until_database_online(uuid) as r:
    kb = await r.parse()

4. Comprehensive Tests (tests/api_resources/test_knowledge_bases.py)

Synchronous Tests (4):

1. test_method_wait_until_database_online: Success path with status progression
2. test_wait_until_database_online_timeout: Timeout handling
3. test_wait_until_database_online_failed: Failure status handling (UNHEALTHY)
4. test_wait_until_database_online_empty_uuid: Input validation

Asynchronous Tests (4):

• Identical coverage as sync tests with async/await patterns

Test Characteristics:

• Uses respx_mock to simulate server responses
• Follows existing pattern from agents.wait_until_ready tests
• Covers success, timeout, failure, and validation scenarios
• Validates exception attributes and error messages

Implementation Details

Polling Logic

start_time = time.time()

while True:
    kb_response = retrieve(uuid)
    status = kb_response.database_status
    
    # Success
    if status == "ONLINE":
        return kb_response
    
    # Failure
    if status in ("DECOMMISSIONED", "UNHEALTHY"):
        raise KnowledgeBaseDatabaseError(...)
    
    # Timeout check
    if time.time() - start_time >= timeout:
        raise KnowledgeBaseDatabaseTimeoutError(...)
    
    # Wait and retry
    sleep(poll_interval)

Key Design Decisions

1. Pattern Consistency: Mirrors agents.wait_until_ready() for familiarity
2. Configurable Polling: Timeout and interval parameters for flexibility
3. Terminal Status Recognition: "DECOMMISSIONED" and "UNHEALTHY" treated as failures
4. Comprehensive Error Info: Exceptions include context about failure state
5. Extra Parameters: Supports extra_headers, extra_query, extra_body for consistency

Testing & Validation

✓ Code Quality Checks

• Linting: Ruff + Black formatting - All checks passed
• Type Checking: Pyright - No errors or warnings
• Import Resolution: All imports successful
• Compilation: No syntax errors

✓ Functional Tests

• Method signatures validated
• Exception instantiation verified
• Wrapper classes confirmed to expose methods
• Test discovery confirmed (8 tests present)

✓ Integration

• Added import time for polling functionality
• Added import httpx for test mock responses
• All existing tests remain unmodified

Files Modified

1. src/gradient/_exceptions.py - Added 2 exception classes
2. src/gradient/resources/knowledge_bases/knowledge_bases.py - Added sync/async methods + wrappers
3. tests/api_resources/test_knowledge_bases.py - Added 8 unit tests
4. Minor formatting: tests/api_resources/test_agents.py, examples/agent_wait_until_ready.py

Commit Information

• Hash: 59ab161
• Message: feat: add wait_until_database_online helper for knowledge base database polling
• Changes: 5 files changed, 283 insertions(+), 36 deletions(-) (mostly test additions)

Usage Example

from gradient import Gradient
from gradient._exceptions import KnowledgeBaseDatabaseError, KnowledgeBaseDatabaseTimeoutError

client = Gradient()

try:
    # Wait for knowledge base database to be ready
    kb_response = client.knowledge_bases.wait_until_database_online(
        uuid="kb-12345678-1234-1234-1234-123456789012",
        timeout=300.0,  # 5 minutes max
        poll_interval=5.0  # Check every 5 seconds
    )
    
    print(f"Database is ONLINE: {kb_response.database_status}")
    
except KnowledgeBaseDatabaseTimeoutError as e:
    print(f"Timeout waiting for KB {e.knowledge_base_id}: {e}")
    
except KnowledgeBaseDatabaseError as e:
    print(f"Database creation failed with status {e.status}: {e}")

Async Example

import asyncio
from gradient import AsyncGradient

async def main():
    client = AsyncGradient()
    
    kb_response = await client.knowledge_bases.wait_until_database_online(
        uuid="kb-uuid-here",
        timeout=300.0,
        poll_interval=5.0
    )
    
    print(f"Database status: {kb_response.database_status}")

asyncio.run(main())

Related Issue

Closes #42

[bbatha]

Closing as a dupe of #49