Skip to content

feat: allow users to configure max wait time for prompt management queries for Python and JS #658

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
copybara-service wants to merge 1 commit into
base: main
Choose a base branch
from
Open

feat: allow users to configure max wait time for prompt management queries for Python and JS #658

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
15 changes: 12 additions & 3 deletions src/prompts.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,7 @@ export class Prompts extends BaseModule {
private async _waitForOperation(
operation: any,
timeout = 90,
maxWaitTime = 60000,
): Promise<string> {
let done = false;
let promptDatasetOperation: any;
Expand All @@ -44,7 +45,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) {
Expand Down Expand Up @@ -77,11 +77,11 @@ export class Prompts extends BaseModule {
private async _waitForProjectOperation(
operation: any,
timeout = 90,
maxWaitTime = 60000,
): Promise<void> {
let done = false;
const startTime = Date.now();
let sleepDuration = 5000; // ms
const maxWaitTime = 60000; // ms

if (!operation?.name) {
throw new Error('Invalid operation name.');
Expand Down Expand Up @@ -258,6 +258,7 @@ export class Prompts extends BaseModule {
const datasetResourceName = await this._waitForOperation(
createPromptDatasetOperation,
config?.timeout || 90,
(config?.maxWaitTime || 60) * 1000,
);

const datasetId = datasetResourceName.split('/').pop();
Expand All @@ -280,6 +281,7 @@ export class Prompts extends BaseModule {
const datasetVersionResourceName = await this._waitForOperation(
createDatasetVersionOperation,
config?.timeout || 90,
(config?.maxWaitTime || 60) * 1000,
);

const versionId = datasetVersionResourceName.split('/').pop();
Expand Down Expand Up @@ -456,6 +458,7 @@ export class Prompts extends BaseModule {
await this._waitForProjectOperation(
deletePromptOperation,
config?.timeout || 90,
(config?.maxWaitTime || 60) * 1000,
);
}

Expand All @@ -482,6 +485,7 @@ export class Prompts extends BaseModule {
await this._waitForProjectOperation(
deleteVersionOperation,
config?.timeout || 90,
(config?.maxWaitTime || 60) * 1000,
);
}

Expand All @@ -505,7 +509,11 @@ export class Prompts extends BaseModule {
versionId,
config,
});
await this._waitForProjectOperation(restorePromptOperation, 90);
await this._waitForProjectOperation(
restorePromptOperation,
config?.timeout || 90,
(config?.maxWaitTime || 60) * 1000,
);
const datasetVersionResource = await this.getDatasetVersionResourceInternal(
{
datasetId: promptId,
Expand Down Expand Up @@ -585,6 +593,7 @@ export class Prompts extends BaseModule {
const datasetVersionResourceName = await this._waitForOperation(
createDatasetVersionOperation,
config?.timeout || 90,
(config?.maxWaitTime || 60) * 1000,
);

const versionId = datasetVersionResourceName.split('/').pop();
Expand Down
12 changes: 12 additions & 0 deletions src/types/common.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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. */
Expand Down Expand Up @@ -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. */
Expand Down Expand Up @@ -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. */
Expand Down Expand Up @@ -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. */
Expand All @@ -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. */
Expand Down
Loading