diff --git a/src/index.ts b/src/index.ts index a068645..51dbd70 100644 --- a/src/index.ts +++ b/src/index.ts @@ -20,6 +20,14 @@ export { PermissionDeniedError, UnprocessableEntityError, } from './core/error'; + +// Export knowledge base helpers +export { + waitForDatabase, + type WaitForDatabaseOptions, + WaitForDatabaseTimeoutError, + WaitForDatabaseFailedError, +} from './resources/knowledge-bases/wait-for-database'; export { IndexingJobAbortedError, IndexingJobNotFoundError, diff --git a/src/internal/detect-platform.ts b/src/internal/detect-platform.ts index 1d49720..0120a78 100644 --- a/src/internal/detect-platform.ts +++ b/src/internal/detect-platform.ts @@ -101,8 +101,8 @@ const getPlatformProperties = (): PlatformProperties => { 'X-Stainless-OS': normalizePlatform((globalThis as any).process?.platform ?? 'unknown'), 'X-Stainless-Arch': normalizeArch((globalThis as any).process?.arch ?? 'unknown'), 'X-Stainless-Runtime': 'bun', - 'X-Stainless-Runtime-Version': (globalThis as any).Bun?.version ?? - (globalThis as any).process?.versions?.bun ?? 'unknown', + 'X-Stainless-Runtime-Version': + (globalThis as any).Bun?.version ?? (globalThis as any).process?.versions?.bun ?? 'unknown', }; } if (detectedPlatform === 'workerd') { diff --git a/src/resources/knowledge-bases/knowledge-bases.ts b/src/resources/knowledge-bases/knowledge-bases.ts index 8758749..aa4fa4b 100644 --- a/src/resources/knowledge-bases/knowledge-bases.ts +++ b/src/resources/knowledge-bases/knowledge-bases.ts @@ -36,6 +36,7 @@ import { import { APIPromise } from '../../core/api-promise'; import { RequestOptions } from '../../internal/request-options'; import { path } from '../../internal/utils/path'; +import { waitForDatabase } from './wait-for-database'; export class KnowledgeBases extends APIResource { dataSources: DataSourcesAPI.DataSources = new DataSourcesAPI.DataSources(this._client); @@ -137,6 +138,31 @@ export class KnowledgeBases extends APIResource { ...options, }); } + + /** + * Polls for knowledge base database creation to complete. + * + * This helper method polls the knowledge base status until the database is ONLINE, + * handling various error states and providing configurable timeout and polling intervals. + * + * @example + * ```ts + * const kb = await client.knowledgeBases.waitForDatabase('123e4567-e89b-12d3-a456-426614174000'); + * console.log('Database is ready:', kb.database_status); // 'ONLINE' + * ``` + * + * @param uuid - The knowledge base UUID to poll for + * @param options - Configuration options for polling behavior + * @returns Promise - The knowledge base with ONLINE database status + * @throws WaitForDatabaseTimeoutError - If polling times out + * @throws WaitForDatabaseFailedError - If the database enters a failed state + */ + async waitForDatabase( + uuid: string, + options?: import('./wait-for-database').WaitForDatabaseOptions, + ): Promise { + return waitForDatabase(this._client, uuid, options || {}); + } } /** @@ -424,6 +450,13 @@ export interface KnowledgeBaseListParams { KnowledgeBases.DataSources = DataSources; KnowledgeBases.IndexingJobs = IndexingJobs; +export { + waitForDatabase, + WaitForDatabaseOptions, + WaitForDatabaseTimeoutError, + WaitForDatabaseFailedError, +} from './wait-for-database'; + export declare namespace KnowledgeBases { export { type APIKnowledgeBase as APIKnowledgeBase, diff --git a/src/resources/knowledge-bases/wait-for-database.ts b/src/resources/knowledge-bases/wait-for-database.ts new file mode 100644 index 0000000..650c1a9 --- /dev/null +++ b/src/resources/knowledge-bases/wait-for-database.ts @@ -0,0 +1,197 @@ +// File generated from our OpenAPI spec by Stainless. See CONTRIBUTING.md for details. + +import { Gradient } from '../../client'; +import { GradientError } from '../../core/error'; +import { sleep } from '../../internal/utils/sleep'; +import { KnowledgeBaseRetrieveResponse } from './knowledge-bases'; +import { RequestOptions } from '../../internal/request-options'; + +export interface WaitForDatabaseOptions { + /** + * The polling interval in milliseconds. Defaults to 5000 (5 seconds). + */ + interval?: number; + + /** + * The maximum time to wait in milliseconds. Defaults to 600000 (10 minutes). + */ + timeout?: number; + + /** + * AbortSignal to cancel the polling operation. + */ + signal?: AbortSignal; + + /** + * Additional request options to pass through to the knowledge base retrieval request. + */ + requestOptions?: RequestOptions; +} + +/** + * Database status values that indicate a successful deployment. + */ +const ONLINE_STATUSES = ['ONLINE'] as const; + +/** + * Database status values that indicate a failed deployment. + */ +const FAILED_STATUSES = ['DECOMMISSIONED', 'UNHEALTHY'] as const; + +/** + * Database status values that indicate the deployment is still in progress. + */ +const PENDING_STATUSES = [ + 'CREATING', + 'POWEROFF', + 'REBUILDING', + 'REBALANCING', + 'FORKING', + 'MIGRATING', + 'RESIZING', + 'RESTORING', + 'POWERING_ON', +] as const; + +export class WaitForDatabaseTimeoutError extends GradientError { + constructor(message: string, kbId?: string, timeout?: number) { + super(message); + this.name = 'WaitForDatabaseTimeoutError'; + if (kbId) { + (this as any).knowledgeBaseId = kbId; + (this as any).timeout = timeout; + } + } +} + +export class WaitForDatabaseFailedError extends GradientError { + constructor(message: string, kbId?: string, status?: string) { + super(message); + this.name = 'WaitForDatabaseFailedError'; + if (kbId) { + (this as any).knowledgeBaseId = kbId; + (this as any).databaseStatus = status; + } + } +} + +/** + * Polls for knowledge base database creation to complete. + * + * This helper function polls the knowledge base status until the database is ONLINE, + * handling various error states and providing configurable timeout and polling intervals. + * + * @example + * ```ts + * import Gradient from '@digitalocean/gradient'; + * + * const client = new Gradient(); + * + * // Basic usage + * try { + * const kb = await client.knowledgeBases.waitForDatabase('123e4567-e89b-12d3-a456-426614174000'); + * console.log('Database is ready:', kb.database_status); // 'ONLINE' + * } catch (error) { + * if (error instanceof WaitForDatabaseTimeoutError) { + * console.log('Polling timed out'); + * } else if (error instanceof WaitForDatabaseFailedError) { + * console.log('Database deployment failed'); + * } + * } + * + * // With AbortSignal + * const controller = new AbortController(); + * setTimeout(() => controller.abort(), 30000); // Cancel after 30 seconds + * + * try { + * const kb = await client.knowledgeBases.waitForDatabase('123e4567-e89b-12d3-a456-426614174000', { + * signal: controller.signal + * }); + * } catch (error) { + * if (error.message === 'Operation was aborted') { + * console.log('Operation was cancelled'); + * } + * } + * ``` + * + * @param client - The Gradient client instance + * @param uuid - The knowledge base UUID to poll for + * @param options - Configuration options for polling behavior + * @returns Promise - The knowledge base with ONLINE database status + * @throws WaitForDatabaseTimeoutError - If polling times out + * @throws WaitForDatabaseFailedError - If the database enters a failed state + * @throws Error - If the operation is aborted via AbortSignal + */ +export async function waitForDatabase( + client: Gradient, + uuid: string, + options: WaitForDatabaseOptions, +): Promise { + const { interval = 5000, timeout = 600000, signal, requestOptions } = options; + + const startTime = Date.now(); + + while (true) { + // Check if operation was aborted + if (signal?.aborted) { + throw new Error('Operation was aborted'); + } + + const elapsed = Date.now() - startTime; + + if (elapsed > timeout) { + throw new WaitForDatabaseTimeoutError( + `Knowledge base database ${uuid} did not become ONLINE within ${timeout}ms`, + uuid, + timeout, + ); + } + + try { + const response = await client.knowledgeBases.retrieve(uuid, requestOptions); + const status = response.database_status; + + if (!status) { + // If database_status is not present, continue polling + await sleep(interval); + continue; + } + + // Check for successful completion + if (ONLINE_STATUSES.includes(status as any)) { + return response; + } + + // Check for failed states + if (FAILED_STATUSES.includes(status as any)) { + throw new WaitForDatabaseFailedError( + `Knowledge base database ${uuid} entered failed state: ${status}`, + uuid, + status, + ); + } + + // Check if still in progress + if (PENDING_STATUSES.includes(status as any)) { + await sleep(interval); + continue; + } + + // Unknown status - treat as error for safety + throw new WaitForDatabaseFailedError( + `Knowledge base database ${uuid} entered unknown state: ${status}`, + uuid, + status, + ); + } catch (error) { + // If it's our custom error, re-throw it + if (error instanceof WaitForDatabaseFailedError || error instanceof WaitForDatabaseTimeoutError) { + throw error; + } + + // For other errors (network issues, etc.), try waiting a bit longer before retrying + await sleep(Math.min(interval * 2, 30000)); // Max 30 seconds between retries + continue; + } + } +} diff --git a/tests/resources/knowledge-bases/wait-for-database.test.ts b/tests/resources/knowledge-bases/wait-for-database.test.ts new file mode 100644 index 0000000..523dff6 --- /dev/null +++ b/tests/resources/knowledge-bases/wait-for-database.test.ts @@ -0,0 +1,81 @@ +// File generated from our OpenAPI spec by Stainless. See CONTRIBUTING.md for details. + +import Gradient from '@digitalocean/gradient'; +import { + WaitForDatabaseFailedError, + WaitForDatabaseTimeoutError, + waitForDatabase, +} from '../../../src/resources/knowledge-bases/wait-for-database'; + +const client = new Gradient({ + accessToken: 'My Access Token', + modelAccessKey: 'My Model Access Key', + agentAccessKey: 'My Agent Access Key', + baseURL: process.env['TEST_API_BASE_URL'] ?? 'http://127.0.0.1:4010', +}); + +describe('waitForDatabase', () => { + const kbUuid = '123e4567-e89b-12d3-a456-426614174000'; + + beforeEach(() => { + jest.resetAllMocks(); + }); + + describe('error classes', () => { + it('should create WaitForDatabaseTimeoutError with correct properties', () => { + const error = new WaitForDatabaseTimeoutError('Test timeout', kbUuid, 1000); + + expect(error.name).toBe('WaitForDatabaseTimeoutError'); + expect(error.message).toBe('Test timeout'); + expect(error).toBeInstanceOf(Error); + }); + + it('should create WaitForDatabaseFailedError with correct properties', () => { + const error = new WaitForDatabaseFailedError('Test failure', kbUuid, 'DECOMMISSIONED'); + + expect(error.name).toBe('WaitForDatabaseFailedError'); + expect(error.message).toBe('Test failure'); + expect(error).toBeInstanceOf(Error); + }); + }); + + describe('function parameters', () => { + it('should accept correct parameters', () => { + expect(typeof waitForDatabase).toBe('function'); + expect(waitForDatabase.length).toBe(3); + }); + + it('should use default options when none provided', () => { + const options = {}; + expect(waitForDatabase).toBeDefined(); + // Function should exist and be callable (will fail at runtime due to mocking, but should compile) + }); + }); + + describe('status constants', () => { + it('should handle different database status values', () => { + // Test that status strings are handled correctly + const statuses = [ + 'ONLINE', + 'CREATING', + 'REBUILDING', + 'RESIZING', + 'POWERING_ON', + 'DECOMMISSIONED', + 'UNHEALTHY', + ]; + + statuses.forEach((status) => { + expect(typeof status).toBe('string'); + }); + }); + }); + + describe('integration test placeholder', () => { + it('should integrate with knowledge base retrieval', async () => { + // This is a placeholder test - actual integration tests would require a running mock server + expect(client.knowledgeBases).toBeDefined(); + expect(client.knowledgeBases.retrieve).toBeDefined(); + }); + }); +});