diff --git a/.changeset/khaki-ladybugs-kiss.md b/.changeset/khaki-ladybugs-kiss.md new file mode 100644 index 00000000000..a860a7d299f --- /dev/null +++ b/.changeset/khaki-ladybugs-kiss.md @@ -0,0 +1,5 @@ +--- +'@firebase/vertexai': minor +--- + +Add an `AbortSignal` to request options for generation methods, allowing them to be aborted. diff --git a/common/api-review/vertexai.api.md b/common/api-review/vertexai.api.md index e7f00c2f4e0..d23e7e6ccb0 100644 --- a/common/api-review/vertexai.api.md +++ b/common/api-review/vertexai.api.md @@ -50,8 +50,8 @@ export class ChatSession { params?: StartChatParams | undefined; // (undocumented) requestOptions?: RequestOptions | undefined; - sendMessage(request: string | Array): Promise; - sendMessageStream(request: string | Array): Promise; + sendMessage(request: string | Array, singleRequestOptions?: SingleRequestOptions): Promise; + sendMessageStream(request: string | Array, singleRequestOptions?: SingleRequestOptions): Promise; } // @public @@ -325,9 +325,9 @@ export interface GenerativeContentBlob { // @public export class GenerativeModel extends VertexAIModel { constructor(vertexAI: VertexAI, modelParams: ModelParams, requestOptions?: RequestOptions); - countTokens(request: CountTokensRequest | string | Array): Promise; - generateContent(request: GenerateContentRequest | string | Array): Promise; - generateContentStream(request: GenerateContentRequest | string | Array): Promise; + countTokens(request: CountTokensRequest | string | Array, singleRequestOptions?: SingleRequestOptions): Promise; + generateContent(request: GenerateContentRequest | string | Array, singleRequestOptions?: SingleRequestOptions): Promise; + generateContentStream(request: GenerateContentRequest | string | Array, singleRequestOptions?: SingleRequestOptions): Promise; // (undocumented) generationConfig: GenerationConfig; // (undocumented) @@ -463,9 +463,9 @@ export interface ImagenInlineImage { // @beta export class ImagenModel extends VertexAIModel { constructor(vertexAI: VertexAI, modelParams: ImagenModelParams, requestOptions?: RequestOptions | undefined); - generateImages(prompt: string): Promise>; + generateImages(prompt: string, singleRequestOptions?: SingleRequestOptions): Promise>; // @internal - generateImagesGCS(prompt: string, gcsURI: string): Promise>; + generateImagesGCS(prompt: string, gcsURI: string, singleRequestOptions?: SingleRequestOptions): Promise>; generationConfig?: ImagenGenerationConfig; // (undocumented) requestOptions?: RequestOptions | undefined; @@ -719,6 +719,11 @@ export interface Segment { startIndex: number; } +// @public +export interface SingleRequestOptions extends RequestOptions { + signal?: AbortSignal; +} + // @public export interface StartChatParams extends BaseParams { // (undocumented) diff --git a/docs-devsite/_toc.yaml b/docs-devsite/_toc.yaml index 665222edb9d..0f5d51acec6 100644 --- a/docs-devsite/_toc.yaml +++ b/docs-devsite/_toc.yaml @@ -588,6 +588,8 @@ toc: path: /docs/reference/js/vertexai.schemashared.md - title: Segment path: /docs/reference/js/vertexai.segment.md + - title: SingleRequestOptions + path: /docs/reference/js/vertexai.singlerequestoptions.md - title: StartChatParams path: /docs/reference/js/vertexai.startchatparams.md - title: StringSchema diff --git a/docs-devsite/vertexai.chatsession.md b/docs-devsite/vertexai.chatsession.md index cc5a75ace16..5fa4a6f48a0 100644 --- a/docs-devsite/vertexai.chatsession.md +++ b/docs-devsite/vertexai.chatsession.md @@ -37,8 +37,8 @@ export declare class ChatSession | Method | Modifiers | Description | | --- | --- | --- | | [getHistory()](./vertexai.chatsession.md#chatsessiongethistory) | | Gets the chat history so far. Blocked prompts are not added to history. Neither blocked candidates nor the prompts that generated them are added to history. | -| [sendMessage(request)](./vertexai.chatsession.md#chatsessionsendmessage) | | Sends a chat message and receives a non-streaming [GenerateContentResult](./vertexai.generatecontentresult.md#generatecontentresult_interface) | -| [sendMessageStream(request)](./vertexai.chatsession.md#chatsessionsendmessagestream) | | Sends a chat message and receives the response as a [GenerateContentStreamResult](./vertexai.generatecontentstreamresult.md#generatecontentstreamresult_interface) containing an iterable stream and a response promise. | +| [sendMessage(request, singleRequestOptions)](./vertexai.chatsession.md#chatsessionsendmessage) | | Sends a chat message and receives a non-streaming [GenerateContentResult](./vertexai.generatecontentresult.md#generatecontentresult_interface) | +| [sendMessageStream(request, singleRequestOptions)](./vertexai.chatsession.md#chatsessionsendmessagestream) | | Sends a chat message and receives the response as a [GenerateContentStreamResult](./vertexai.generatecontentstreamresult.md#generatecontentstreamresult_interface) containing an iterable stream and a response promise. | ## ChatSession.(constructor) @@ -103,7 +103,7 @@ Sends a chat message and receives a non-streaming [GenerateContentResult]( Signature: ```typescript -sendMessage(request: string | Array): Promise; +sendMessage(request: string | Array, singleRequestOptions?: SingleRequestOptions): Promise; ``` #### Parameters @@ -111,6 +111,7 @@ sendMessage(request: string | Array): Promise> | | +| singleRequestOptions | [SingleRequestOptions](./vertexai.singlerequestoptions.md#singlerequestoptions_interface) | | Returns: @@ -123,7 +124,7 @@ Sends a chat message and receives the response as a [GenerateContentStream Signature: ```typescript -sendMessageStream(request: string | Array): Promise; +sendMessageStream(request: string | Array, singleRequestOptions?: SingleRequestOptions): Promise; ``` #### Parameters @@ -131,6 +132,7 @@ sendMessageStream(request: string | Array): Promise> | | +| singleRequestOptions | [SingleRequestOptions](./vertexai.singlerequestoptions.md#singlerequestoptions_interface) | | Returns: diff --git a/docs-devsite/vertexai.generativemodel.md b/docs-devsite/vertexai.generativemodel.md index b734e241e78..6d691ed9e5a 100644 --- a/docs-devsite/vertexai.generativemodel.md +++ b/docs-devsite/vertexai.generativemodel.md @@ -40,9 +40,9 @@ export declare class GenerativeModel extends VertexAIModel | Method | Modifiers | Description | | --- | --- | --- | -| [countTokens(request)](./vertexai.generativemodel.md#generativemodelcounttokens) | | Counts the tokens in the provided request. | -| [generateContent(request)](./vertexai.generativemodel.md#generativemodelgeneratecontent) | | Makes a single non-streaming call to the model and returns an object containing a single [GenerateContentResponse](./vertexai.generatecontentresponse.md#generatecontentresponse_interface). | -| [generateContentStream(request)](./vertexai.generativemodel.md#generativemodelgeneratecontentstream) | | Makes a single streaming call to the model and returns an object containing an iterable stream that iterates over all chunks in the streaming response as well as a promise that returns the final aggregated response. | +| [countTokens(request, singleRequestOptions)](./vertexai.generativemodel.md#generativemodelcounttokens) | | Counts the tokens in the provided request. | +| [generateContent(request, singleRequestOptions)](./vertexai.generativemodel.md#generativemodelgeneratecontent) | | Makes a single non-streaming call to the model and returns an object containing a single [GenerateContentResponse](./vertexai.generatecontentresponse.md#generatecontentresponse_interface). | +| [generateContentStream(request, singleRequestOptions)](./vertexai.generativemodel.md#generativemodelgeneratecontentstream) | | Makes a single streaming call to the model and returns an object containing an iterable stream that iterates over all chunks in the streaming response as well as a promise that returns the final aggregated response. | | [startChat(startChatParams)](./vertexai.generativemodel.md#generativemodelstartchat) | | Gets a new [ChatSession](./vertexai.chatsession.md#chatsession_class) instance which can be used for multi-turn chats. | ## GenerativeModel.(constructor) @@ -118,7 +118,7 @@ Counts the tokens in the provided request. Signature: ```typescript -countTokens(request: CountTokensRequest | string | Array): Promise; +countTokens(request: CountTokensRequest | string | Array, singleRequestOptions?: SingleRequestOptions): Promise; ``` #### Parameters @@ -126,6 +126,7 @@ countTokens(request: CountTokensRequest | string | Array): Promis | Parameter | Type | Description | | --- | --- | --- | | request | [CountTokensRequest](./vertexai.counttokensrequest.md#counttokensrequest_interface) \| string \| Array<string \| [Part](./vertexai.md#part)> | | +| singleRequestOptions | [SingleRequestOptions](./vertexai.singlerequestoptions.md#singlerequestoptions_interface) | | Returns: @@ -138,7 +139,7 @@ Makes a single non-streaming call to the model and returns an object containing Signature: ```typescript -generateContent(request: GenerateContentRequest | string | Array): Promise; +generateContent(request: GenerateContentRequest | string | Array, singleRequestOptions?: SingleRequestOptions): Promise; ``` #### Parameters @@ -146,6 +147,7 @@ generateContent(request: GenerateContentRequest | string | Array) | Parameter | Type | Description | | --- | --- | --- | | request | [GenerateContentRequest](./vertexai.generatecontentrequest.md#generatecontentrequest_interface) \| string \| Array<string \| [Part](./vertexai.md#part)> | | +| singleRequestOptions | [SingleRequestOptions](./vertexai.singlerequestoptions.md#singlerequestoptions_interface) | | Returns: @@ -158,7 +160,7 @@ Makes a single streaming call to the model and returns an object containing an i Signature: ```typescript -generateContentStream(request: GenerateContentRequest | string | Array): Promise; +generateContentStream(request: GenerateContentRequest | string | Array, singleRequestOptions?: SingleRequestOptions): Promise; ``` #### Parameters @@ -166,6 +168,7 @@ generateContentStream(request: GenerateContentRequest | string | Array> | | +| singleRequestOptions | [SingleRequestOptions](./vertexai.singlerequestoptions.md#singlerequestoptions_interface) | | Returns: diff --git a/docs-devsite/vertexai.imagenmodel.md b/docs-devsite/vertexai.imagenmodel.md index 63e15ff133a..d1f8885032c 100644 --- a/docs-devsite/vertexai.imagenmodel.md +++ b/docs-devsite/vertexai.imagenmodel.md @@ -42,7 +42,7 @@ export declare class ImagenModel extends VertexAIModel | Method | Modifiers | Description | | --- | --- | --- | -| [generateImages(prompt)](./vertexai.imagenmodel.md#imagenmodelgenerateimages) | | (Public Preview) Generates images using the Imagen model and returns them as base64-encoded strings. | +| [generateImages(prompt, singleRequestOptions)](./vertexai.imagenmodel.md#imagenmodelgenerateimages) | | (Public Preview) Generates images using the Imagen model and returns them as base64-encoded strings. | ## ImagenModel.(constructor) @@ -118,7 +118,7 @@ If the prompt was not blocked, but one or more of the generated images were filt Signature: ```typescript -generateImages(prompt: string): Promise>; +generateImages(prompt: string, singleRequestOptions?: SingleRequestOptions): Promise>; ``` #### Parameters @@ -126,6 +126,7 @@ generateImages(prompt: string): PromiseReturns: diff --git a/docs-devsite/vertexai.md b/docs-devsite/vertexai.md index fca51b42f4f..cc411582bb4 100644 --- a/docs-devsite/vertexai.md +++ b/docs-devsite/vertexai.md @@ -111,6 +111,7 @@ The Vertex AI in Firebase Web SDK. | [SchemaRequest](./vertexai.schemarequest.md#schemarequest_interface) | Final format for [Schema](./vertexai.schema.md#schema_class) params passed to backend requests. | | [SchemaShared](./vertexai.schemashared.md#schemashared_interface) | Basic [Schema](./vertexai.schema.md#schema_class) properties shared across several Schema-related types. | | [Segment](./vertexai.segment.md#segment_interface) | | +| [SingleRequestOptions](./vertexai.singlerequestoptions.md#singlerequestoptions_interface) | Options that can be provided per-request. Extends the base [RequestOptions](./vertexai.requestoptions.md#requestoptions_interface) (like timeout and baseUrl) with request-specific controls like cancellation via AbortSignal.Options specified here will override any default [RequestOptions](./vertexai.requestoptions.md#requestoptions_interface) configured on a model (e.g. [GenerativeModel](./vertexai.generativemodel.md#generativemodel_class)). | | [StartChatParams](./vertexai.startchatparams.md#startchatparams_interface) | Params for [GenerativeModel.startChat()](./vertexai.generativemodel.md#generativemodelstartchat). | | [TextPart](./vertexai.textpart.md#textpart_interface) | Content part interface if the part represents a text string. | | [ToolConfig](./vertexai.toolconfig.md#toolconfig_interface) | Tool config. This config is shared for all tools provided in the request. | diff --git a/docs-devsite/vertexai.singlerequestoptions.md b/docs-devsite/vertexai.singlerequestoptions.md new file mode 100644 index 00000000000..af5d46059fd --- /dev/null +++ b/docs-devsite/vertexai.singlerequestoptions.md @@ -0,0 +1,57 @@ +Project: /docs/reference/js/_project.yaml +Book: /docs/reference/_book.yaml +page_type: reference + +{% comment %} +DO NOT EDIT THIS FILE! +This is generated by the JS SDK team, and any local changes will be +overwritten. Changes should be made in the source code at +https://github.com/firebase/firebase-js-sdk +{% endcomment %} + +# SingleRequestOptions interface +Options that can be provided per-request. Extends the base [RequestOptions](./vertexai.requestoptions.md#requestoptions_interface) (like `timeout` and `baseUrl`) with request-specific controls like cancellation via `AbortSignal`. + +Options specified here will override any default [RequestOptions](./vertexai.requestoptions.md#requestoptions_interface) configured on a model (e.g. [GenerativeModel](./vertexai.generativemodel.md#generativemodel_class)). + +Signature: + +```typescript +export interface SingleRequestOptions extends RequestOptions +``` +Extends: [RequestOptions](./vertexai.requestoptions.md#requestoptions_interface) + +## Properties + +| Property | Type | Description | +| --- | --- | --- | +| [signal](./vertexai.singlerequestoptions.md#singlerequestoptionssignal) | AbortSignal | An AbortSignal instance that allows cancelling ongoing requests (like generateContent or generateImages).If provided, calling abort() on the corresponding AbortController will attempt to cancel the underlying HTTP request. An AbortError will be thrown if cancellation is successful.Note that this will not cancel the request in the backend, so billing will still be applied despite cancellation. | + +## SingleRequestOptions.signal + +An `AbortSignal` instance that allows cancelling ongoing requests (like `generateContent` or `generateImages`). + +If provided, calling `abort()` on the corresponding `AbortController` will attempt to cancel the underlying HTTP request. An `AbortError` will be thrown if cancellation is successful. + +Note that this will not cancel the request in the backend, so billing will still be applied despite cancellation. + +Signature: + +```typescript +signal?: AbortSignal; +``` + +### Example + + +```javascript +const controller = new AbortController(); +const model = getGenerativeModel({ + // ... +}); + +// To cancel request: +controller.abort(); + +``` + diff --git a/packages/vertexai/src/methods/chat-session.ts b/packages/vertexai/src/methods/chat-session.ts index dd22b29a7c8..42a741ef861 100644 --- a/packages/vertexai/src/methods/chat-session.ts +++ b/packages/vertexai/src/methods/chat-session.ts @@ -22,6 +22,7 @@ import { GenerateContentStreamResult, Part, RequestOptions, + SingleRequestOptions, StartChatParams } from '../types'; import { formatNewContent } from '../requests/request-helpers'; @@ -75,7 +76,8 @@ export class ChatSession { * {@link GenerateContentResult} */ async sendMessage( - request: string | Array + request: string | Array, + singleRequestOptions?: SingleRequestOptions ): Promise { await this._sendPromise; const newContent = formatNewContent(request); @@ -95,7 +97,11 @@ export class ChatSession { this._apiSettings, this.model, generateContentRequest, - this.requestOptions + // Merge requestOptions + { + ...singleRequestOptions, + ...this.requestOptions + } ) ) .then(result => { @@ -130,7 +136,8 @@ export class ChatSession { * and a response promise. */ async sendMessageStream( - request: string | Array + request: string | Array, + singleRequestOptions?: SingleRequestOptions ): Promise { await this._sendPromise; const newContent = formatNewContent(request); @@ -146,7 +153,11 @@ export class ChatSession { this._apiSettings, this.model, generateContentRequest, - this.requestOptions + // Merge requestOptions + { + ...singleRequestOptions, + ...this.requestOptions + } ); // Add onto the chain. diff --git a/packages/vertexai/src/methods/count-tokens.ts b/packages/vertexai/src/methods/count-tokens.ts index c9d43a5b6fd..b89dab44e64 100644 --- a/packages/vertexai/src/methods/count-tokens.ts +++ b/packages/vertexai/src/methods/count-tokens.ts @@ -18,7 +18,7 @@ import { CountTokensRequest, CountTokensResponse, - RequestOptions + SingleRequestOptions } from '../types'; import { Task, makeRequest } from '../requests/request'; import { ApiSettings } from '../types/internal'; @@ -27,7 +27,7 @@ export async function countTokens( apiSettings: ApiSettings, model: string, params: CountTokensRequest, - requestOptions?: RequestOptions + singleRequestOptions?: SingleRequestOptions ): Promise { const response = await makeRequest( model, @@ -35,7 +35,7 @@ export async function countTokens( apiSettings, false, JSON.stringify(params), - requestOptions + singleRequestOptions ); return response.json(); } diff --git a/packages/vertexai/src/methods/generate-content.ts b/packages/vertexai/src/methods/generate-content.ts index 0944b38016a..c67c36ac1d2 100644 --- a/packages/vertexai/src/methods/generate-content.ts +++ b/packages/vertexai/src/methods/generate-content.ts @@ -20,7 +20,7 @@ import { GenerateContentResponse, GenerateContentResult, GenerateContentStreamResult, - RequestOptions + SingleRequestOptions } from '../types'; import { Task, makeRequest } from '../requests/request'; import { createEnhancedContentResponse } from '../requests/response-helpers'; @@ -31,7 +31,7 @@ export async function generateContentStream( apiSettings: ApiSettings, model: string, params: GenerateContentRequest, - requestOptions?: RequestOptions + singleRequestOptions?: SingleRequestOptions ): Promise { const response = await makeRequest( model, @@ -39,7 +39,7 @@ export async function generateContentStream( apiSettings, /* stream */ true, JSON.stringify(params), - requestOptions + singleRequestOptions ); return processStream(response); } @@ -48,7 +48,7 @@ export async function generateContent( apiSettings: ApiSettings, model: string, params: GenerateContentRequest, - requestOptions?: RequestOptions + singleRequestOptions?: SingleRequestOptions ): Promise { const response = await makeRequest( model, @@ -56,7 +56,7 @@ export async function generateContent( apiSettings, /* stream */ false, JSON.stringify(params), - requestOptions + singleRequestOptions ); const responseJson: GenerateContentResponse = await response.json(); const enhancedResponse = createEnhancedContentResponse(responseJson); diff --git a/packages/vertexai/src/models/generative-model.ts b/packages/vertexai/src/models/generative-model.ts index b4cf464f025..5991e9bd62c 100644 --- a/packages/vertexai/src/models/generative-model.ts +++ b/packages/vertexai/src/models/generative-model.ts @@ -29,11 +29,12 @@ import { GenerationConfig, ModelParams, Part, - RequestOptions, SafetySetting, + RequestOptions, StartChatParams, Tool, - ToolConfig + ToolConfig, + SingleRequestOptions } from '../types'; import { ChatSession } from '../methods/chat-session'; import { countTokens } from '../methods/count-tokens'; @@ -77,7 +78,8 @@ export class GenerativeModel extends VertexAIModel { * and returns an object containing a single {@link GenerateContentResponse}. */ async generateContent( - request: GenerateContentRequest | string | Array + request: GenerateContentRequest | string | Array, + singleRequestOptions?: SingleRequestOptions ): Promise { const formattedParams = formatGenerateContentInput(request); return generateContent( @@ -91,7 +93,11 @@ export class GenerativeModel extends VertexAIModel { systemInstruction: this.systemInstruction, ...formattedParams }, - this.requestOptions + // Merge request options + { + ...singleRequestOptions, + ...this.requestOptions + } ); } @@ -102,7 +108,8 @@ export class GenerativeModel extends VertexAIModel { * a promise that returns the final aggregated response. */ async generateContentStream( - request: GenerateContentRequest | string | Array + request: GenerateContentRequest | string | Array, + singleRequestOptions?: SingleRequestOptions ): Promise { const formattedParams = formatGenerateContentInput(request); return generateContentStream( @@ -116,7 +123,11 @@ export class GenerativeModel extends VertexAIModel { systemInstruction: this.systemInstruction, ...formattedParams }, - this.requestOptions + // Merge request options + { + ...singleRequestOptions, + ...this.requestOptions + } ); } @@ -142,9 +153,19 @@ export class GenerativeModel extends VertexAIModel { * Counts the tokens in the provided request. */ async countTokens( - request: CountTokensRequest | string | Array + request: CountTokensRequest | string | Array, + singleRequestOptions?: SingleRequestOptions ): Promise { const formattedParams = formatGenerateContentInput(request); - return countTokens(this._apiSettings, this.model, formattedParams); + return countTokens( + this._apiSettings, + this.model, + formattedParams, + // Merge request options + { + ...singleRequestOptions, + ...this.requestOptions + } + ); } } diff --git a/packages/vertexai/src/models/imagen-model.ts b/packages/vertexai/src/models/imagen-model.ts index 89c740852a3..7e7de956aa9 100644 --- a/packages/vertexai/src/models/imagen-model.ts +++ b/packages/vertexai/src/models/imagen-model.ts @@ -26,7 +26,8 @@ import { RequestOptions, ImagenModelParams, ImagenGenerationResponse, - ImagenSafetySettings + ImagenSafetySettings, + SingleRequestOptions } from '../types'; import { VertexAIModel } from './vertexai-model'; @@ -102,7 +103,8 @@ export class ImagenModel extends VertexAIModel { * @beta */ async generateImages( - prompt: string + prompt: string, + singleRequestOptions?: SingleRequestOptions ): Promise> { const body = createPredictRequestBody(prompt, { ...this.generationConfig, @@ -114,7 +116,11 @@ export class ImagenModel extends VertexAIModel { this._apiSettings, /* stream */ false, JSON.stringify(body), - this.requestOptions + // Merge request options + { + ...singleRequestOptions, + ...this.requestOptions + } ); return handlePredictResponse(response); } @@ -140,7 +146,8 @@ export class ImagenModel extends VertexAIModel { */ async generateImagesGCS( prompt: string, - gcsURI: string + gcsURI: string, + singleRequestOptions?: SingleRequestOptions ): Promise> { const body = createPredictRequestBody(prompt, { gcsURI, @@ -153,7 +160,10 @@ export class ImagenModel extends VertexAIModel { this._apiSettings, /* stream */ false, JSON.stringify(body), - this.requestOptions + { + ...singleRequestOptions, + ...this.requestOptions + } ); return handlePredictResponse(response); } diff --git a/packages/vertexai/src/requests/request.ts b/packages/vertexai/src/requests/request.ts index 47e4c6ab446..002930cfbc8 100644 --- a/packages/vertexai/src/requests/request.ts +++ b/packages/vertexai/src/requests/request.ts @@ -15,7 +15,11 @@ * limitations under the License. */ -import { ErrorDetails, RequestOptions, VertexAIErrorCode } from '../types'; +import { + ErrorDetails, + SingleRequestOptions, + VertexAIErrorCode +} from '../types'; import { VertexAIError } from '../errors'; import { ApiSettings } from '../types/internal'; import { @@ -27,6 +31,9 @@ import { } from '../constants'; import { logger } from '../logger'; +const TIMEOUT_EXPIRED_MESSAGE = 'Timeout has expired.'; +const ABORT_ERROR_NAME = 'AbortError'; + export enum Task { GENERATE_CONTENT = 'generateContent', STREAM_GENERATE_CONTENT = 'streamGenerateContent', @@ -40,7 +47,7 @@ export class RequestUrl { public task: Task, public apiSettings: ApiSettings, public stream: boolean, - public requestOptions?: RequestOptions + public requestOptions?: SingleRequestOptions ) {} toString(): string { // TODO: allow user-set option if that feature becomes available @@ -115,9 +122,15 @@ export async function constructRequest( apiSettings: ApiSettings, stream: boolean, body: string, - requestOptions?: RequestOptions + singleRequestOptions?: SingleRequestOptions ): Promise<{ url: string; fetchOptions: RequestInit }> { - const url = new RequestUrl(model, task, apiSettings, stream, requestOptions); + const url = new RequestUrl( + model, + task, + apiSettings, + stream, + singleRequestOptions + ); return { url: url.toString(), fetchOptions: { @@ -134,11 +147,49 @@ export async function makeRequest( apiSettings: ApiSettings, stream: boolean, body: string, - requestOptions?: RequestOptions + singleRequestOptions?: SingleRequestOptions ): Promise { - const url = new RequestUrl(model, task, apiSettings, stream, requestOptions); + const url = new RequestUrl( + model, + task, + apiSettings, + stream, + singleRequestOptions + ); let response; - let fetchTimeoutId: string | number | NodeJS.Timeout | undefined; + + const externalSignal = singleRequestOptions?.signal; + const timeoutMillis = + singleRequestOptions?.timeout != null && singleRequestOptions.timeout >= 0 + ? singleRequestOptions.timeout + : DEFAULT_FETCH_TIMEOUT_MS; + const internalAbortController = new AbortController(); + const fetchTimeoutId = setTimeout(() => { + internalAbortController.abort(TIMEOUT_EXPIRED_MESSAGE); + logger.debug( + `Aborting request to ${url} due to timeout (${timeoutMillis}ms)` + ); + }, timeoutMillis); + + if (externalSignal) { + if (externalSignal.aborted) { + clearTimeout(fetchTimeoutId); + throw new DOMException( + externalSignal.reason ?? 'Aborted externally before fetch', + ABORT_ERROR_NAME + ); + } + + const externalAbortListener = (): void => { + logger.debug(`Aborting request to ${url} due to external abort signal.`); + internalAbortController.abort(externalSignal.reason); + }; + + externalSignal.addEventListener('abort', externalAbortListener, { + once: true + }); + } + try { const request = await constructRequest( model, @@ -146,16 +197,9 @@ export async function makeRequest( apiSettings, stream, body, - requestOptions + singleRequestOptions ); - // Timeout is 180s by default - const timeoutMillis = - requestOptions?.timeout != null && requestOptions.timeout >= 0 - ? requestOptions.timeout - : DEFAULT_FETCH_TIMEOUT_MS; - const abortController = new AbortController(); - fetchTimeoutId = setTimeout(() => abortController.abort(), timeoutMillis); - request.fetchOptions.signal = abortController.signal; + request.fetchOptions.signal = internalAbortController.signal; response = await fetch(request.url, request.fetchOptions); if (!response.ok) { diff --git a/packages/vertexai/src/types/requests.ts b/packages/vertexai/src/types/requests.ts index 5058b457365..f1fb7536d12 100644 --- a/packages/vertexai/src/types/requests.ts +++ b/packages/vertexai/src/types/requests.ts @@ -143,6 +143,44 @@ export interface RequestOptions { baseUrl?: string; } +/** + * Options that can be provided per-request. + * Extends the base {@link RequestOptions} (like `timeout` and `baseUrl`) + * with request-specific controls like cancellation via `AbortSignal`. + * + * Options specified here will override any default {@link RequestOptions} + * configured on a model (e.g. {@link GenerativeModel}). + * + * @public + */ +export interface SingleRequestOptions extends RequestOptions { + /** + * An `AbortSignal` instance that allows cancelling ongoing requests (like `generateContent` or + * `generateImages`). + * + * If provided, calling `abort()` on the corresponding `AbortController` + * will attempt to cancel the underlying HTTP request. An `AbortError` will be thrown + * if cancellation is successful. + * + * Note that this will not cancel the request in the backend, so billing will + * still be applied despite cancellation. + * + * @example + * ```javascript + * const controller = new AbortController(); + * const model = getGenerativeModel({ + * // ... + * }); + * + * // To cancel request: + * controller.abort(); + * ``` + * @see https://developer.mozilla.org/en-US/docs/Web/API/AbortSignal + */ + + signal?: AbortSignal; +} + /** * Defines a tool that model can call to access external knowledge. * @public