From 7a9ec5fc9bf75789fbda8b02700f18df473ba3f8 Mon Sep 17 00:00:00 2001 From: Ayush Agrawal Date: Wed, 29 Jul 2026 13:13:00 -0700 Subject: [PATCH] feat: allow users to configure max wait time for prompt management queries for Python and JS PiperOrigin-RevId: 956048771 --- src/prompts.ts | 32 ++++++++++++++++++++++---------- src/types/common.ts | 12 ++++++++++++ 2 files changed, 34 insertions(+), 10 deletions(-) diff --git a/src/prompts.ts b/src/prompts.ts index a73053a7..10772240 100644 --- a/src/prompts.ts +++ b/src/prompts.ts @@ -16,10 +16,14 @@ export class Prompts extends BaseModule { super(); } + private static readonly _DEFAULT_TIMEOUT = 90; + private static readonly _DEFAULT_MAX_WAIT_TIME = 60; + /* eslint-disable @typescript-eslint/no-explicit-any */ private async _waitForOperation( operation: any, - timeout = 90, + timeout = Prompts._DEFAULT_TIMEOUT, + maxWaitTime = Prompts._DEFAULT_MAX_WAIT_TIME * 1000, ): Promise { let done = false; let promptDatasetOperation: any; @@ -44,7 +48,6 @@ export class Prompts extends BaseModule { const startTime = Date.now(); let sleepDuration = 5000; // ms - const maxWaitTime = 60000; // ms while (!done) { if (Date.now() - startTime > timeout * 1000) { @@ -76,12 +79,12 @@ export class Prompts extends BaseModule { private async _waitForProjectOperation( operation: any, - timeout = 90, + timeout = Prompts._DEFAULT_TIMEOUT, + maxWaitTime = Prompts._DEFAULT_MAX_WAIT_TIME * 1000, ): Promise { let done = false; const startTime = Date.now(); let sleepDuration = 5000; // ms - const maxWaitTime = 60000; // ms if (!operation?.name) { throw new Error('Invalid operation name.'); @@ -257,7 +260,8 @@ export class Prompts extends BaseModule { const datasetResourceName = await this._waitForOperation( createPromptDatasetOperation, - config?.timeout || 90, + config?.timeout ?? Prompts._DEFAULT_TIMEOUT, + (config?.maxWaitTime ?? Prompts._DEFAULT_MAX_WAIT_TIME) * 1000, ); const datasetId = datasetResourceName.split('/').pop(); @@ -279,7 +283,8 @@ export class Prompts extends BaseModule { const datasetVersionResourceName = await this._waitForOperation( createDatasetVersionOperation, - config?.timeout || 90, + config?.timeout ?? Prompts._DEFAULT_TIMEOUT, + (config?.maxWaitTime ?? Prompts._DEFAULT_MAX_WAIT_TIME) * 1000, ); const versionId = datasetVersionResourceName.split('/').pop(); @@ -455,7 +460,8 @@ export class Prompts extends BaseModule { }); await this._waitForProjectOperation( deletePromptOperation, - config?.timeout || 90, + config?.timeout ?? Prompts._DEFAULT_TIMEOUT, + (config?.maxWaitTime ?? Prompts._DEFAULT_MAX_WAIT_TIME) * 1000, ); } @@ -481,7 +487,8 @@ export class Prompts extends BaseModule { }); await this._waitForProjectOperation( deleteVersionOperation, - config?.timeout || 90, + config?.timeout ?? Prompts._DEFAULT_TIMEOUT, + (config?.maxWaitTime ?? Prompts._DEFAULT_MAX_WAIT_TIME) * 1000, ); } @@ -505,7 +512,11 @@ export class Prompts extends BaseModule { versionId, config, }); - await this._waitForProjectOperation(restorePromptOperation, 90); + await this._waitForProjectOperation( + restorePromptOperation, + config?.timeout ?? Prompts._DEFAULT_TIMEOUT, + (config?.maxWaitTime ?? Prompts._DEFAULT_MAX_WAIT_TIME) * 1000, + ); const datasetVersionResource = await this.getDatasetVersionResourceInternal( { datasetId: promptId, @@ -584,7 +595,8 @@ export class Prompts extends BaseModule { const datasetVersionResourceName = await this._waitForOperation( createDatasetVersionOperation, - config?.timeout || 90, + config?.timeout ?? Prompts._DEFAULT_TIMEOUT, + (config?.maxWaitTime ?? Prompts._DEFAULT_MAX_WAIT_TIME) * 1000, ); const versionId = datasetVersionResourceName.split('/').pop(); diff --git a/src/types/common.ts b/src/types/common.ts index 1ae2e68b..4c1f5580 100644 --- a/src/types/common.ts +++ b/src/types/common.ts @@ -3339,6 +3339,8 @@ export declare interface DeletePromptConfig { abortSignal?: AbortSignal; /** Timeout for the delete prompt operation in seconds. Defaults to 90. */ timeout?: number; + /** Maximum interval between polling requests in seconds. Defaults to 60. */ + maxWaitTime?: number; } /** Parameters for deleting a prompt dataset. */ @@ -3392,6 +3394,10 @@ export declare interface RestoreVersionConfig { be charged usage for any applicable operations. */ abortSignal?: AbortSignal; + /** Timeout for the restore prompt version operation in seconds. Defaults to 90. */ + timeout?: number; + /** Maximum interval between polling requests in seconds. Defaults to 60. */ + maxWaitTime?: number; } /** Parameters for restoring a prompt version. */ @@ -3434,6 +3440,8 @@ export declare interface UpdatePromptConfig { timeout?: number; /** Customer-managed encryption key spec for a prompt dataset. If set, this prompt dataset and all sub-resources of this prompt dataset will be secured by this key. */ encryptionSpec?: genaiTypes.EncryptionSpec; + /** The maximum interval between polling requests in seconds. If not set, the default interval is 60 seconds. */ + maxWaitTime?: number; } /** Parameters for creating a dataset resource to store prompts. */ @@ -3766,6 +3774,8 @@ export declare interface CreatePromptConfig { encryptionSpec?: genaiTypes.EncryptionSpec; /** The display name for the prompt version. If not set, a default name with a timestamp will be used. */ versionDisplayName?: string; + /** The maximum interval between requests in seconds. If not set, the default interval is 60 seconds. */ + maxWaitTime?: number; } /** Config for creating a prompt version. */ @@ -3787,6 +3797,8 @@ export declare interface CreatePromptVersionConfig { promptDisplayName?: string; /** Customer-managed encryption key spec for a prompt dataset. If set, this prompt dataset and all sub-resources of this prompt dataset will be secured by this key. */ encryptionSpec?: genaiTypes.EncryptionSpec; + /** The maximum interval between requests in seconds. If not set, the default interval is 60 seconds. */ + maxWaitTime?: number; } /** Config for getting a prompt. */