feat(vertexai): Add support for AbortSignal

Author: dlarocque

.changeset/long-keys-watch.md

@@ -0,0 +1,6 @@
+---
+'firebase': minor
+'@firebase/vertexai': minor
+---
+
+Add support for `AbortSignal`, allowing requests to be aborted.

common/api-review/vertexai.api.md

@@ -50,8 +50,8 @@ export class ChatSession {
     params?: StartChatParams | undefined;
     // (undocumented)
     requestOptions?: RequestOptions | undefined;
-    sendMessage(request: string | Array<string | Part>): Promise<GenerateContentResult>;
-    sendMessageStream(request: string | Array<string | Part>): Promise<GenerateContentStreamResult>;
+    sendMessage(request: string | Array<string | Part>, singleRequestOptions?: SingleRequestOptions): Promise<GenerateContentResult>;
+    sendMessageStream(request: string | Array<string | Part>, singleRequestOptions?: SingleRequestOptions): Promise<GenerateContentStreamResult>;
     }
 
 // @public
@@ -328,9 +328,9 @@ export interface GenerativeContentBlob {
 // @public
 export class GenerativeModel extends VertexAIModel {
     constructor(vertexAI: VertexAI, modelParams: ModelParams, requestOptions?: RequestOptions);
-    countTokens(request: CountTokensRequest | string | Array<string | Part>): Promise<CountTokensResponse>;
-    generateContent(request: GenerateContentRequest | string | Array<string | Part>): Promise<GenerateContentResult>;
-    generateContentStream(request: GenerateContentRequest | string | Array<string | Part>): Promise<GenerateContentStreamResult>;
+    countTokens(request: CountTokensRequest | string | Array<string | Part>, singleRequestOptions?: SingleRequestOptions): Promise<CountTokensResponse>;
+    generateContent(request: GenerateContentRequest | string | Array<string | Part>, singleRequestOptions?: SingleRequestOptions): Promise<GenerateContentResult>;
+    generateContentStream(request: GenerateContentRequest | string | Array<string | Part>, singleRequestOptions?: SingleRequestOptions): Promise<GenerateContentStreamResult>;
     // (undocumented)
     generationConfig: GenerationConfig;
     // (undocumented)
@@ -466,9 +466,9 @@ export interface ImagenInlineImage {
 // @beta
 export class ImagenModel extends VertexAIModel {
     constructor(vertexAI: VertexAI, modelParams: ImagenModelParams, requestOptions?: RequestOptions | undefined);
-    generateImages(prompt: string): Promise<ImagenGenerationResponse<ImagenInlineImage>>;
+    generateImages(prompt: string, singleRequestOptions?: SingleRequestOptions): Promise<ImagenGenerationResponse<ImagenInlineImage>>;
     // @internal
-    generateImagesGCS(prompt: string, gcsURI: string): Promise<ImagenGenerationResponse<ImagenGCSImage>>;
+    generateImagesGCS(prompt: string, gcsURI: string, singleRequestOptions?: SingleRequestOptions): Promise<ImagenGenerationResponse<ImagenGCSImage>>;
     generationConfig?: ImagenGenerationConfig;
     // (undocumented)
     requestOptions?: RequestOptions | undefined;
@@ -731,6 +731,11 @@ export interface Segment {
     startIndex: number;
 }
 
+// @public
+export interface SingleRequestOptions extends RequestOptions {
+    signal?: AbortSignal;
+}
+
 // @public
 export interface StartChatParams extends BaseParams {
     // (undocumented)

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

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,14 +103,15 @@ Sends a chat message and receives a non-streaming [GenerateContentResult](./vert
 <b>Signature:</b>
 
 ```typescript
-sendMessage(request: string | Array<string | Part>): Promise<GenerateContentResult>;
+sendMessage(request: string | Array<string | Part>, singleRequestOptions?: SingleRequestOptions): Promise<GenerateContentResult>;
 ```
 
 #### Parameters
 
 |  Parameter | Type | Description |
 |  --- | --- | --- |
 |  request | string \| Array&lt;string \| [Part](./vertexai.md#part)<!-- -->&gt; |  |
+|  singleRequestOptions | [SingleRequestOptions](./vertexai.singlerequestoptions.md#singlerequestoptions_interface) |  |
 
 <b>Returns:</b>
 
@@ -123,14 +124,15 @@ Sends a chat message and receives the response as a [GenerateContentStreamResult
 <b>Signature:</b>
 
 ```typescript
-sendMessageStream(request: string | Array<string | Part>): Promise<GenerateContentStreamResult>;
+sendMessageStream(request: string | Array<string | Part>, singleRequestOptions?: SingleRequestOptions): Promise<GenerateContentStreamResult>;
 ```
 
 #### Parameters
 
 |  Parameter | Type | Description |
 |  --- | --- | --- |
 |  request | string \| Array&lt;string \| [Part](./vertexai.md#part)<!-- -->&gt; |  |
+|  singleRequestOptions | [SingleRequestOptions](./vertexai.singlerequestoptions.md#singlerequestoptions_interface) |  |
 
 <b>Returns:</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,14 +118,15 @@ Counts the tokens in the provided request.
 <b>Signature:</b>
 
 ```typescript
-countTokens(request: CountTokensRequest | string | Array<string | Part>): Promise<CountTokensResponse>;
+countTokens(request: CountTokensRequest | string | Array<string | Part>, singleRequestOptions?: SingleRequestOptions): Promise<CountTokensResponse>;
 ```
 
 #### Parameters
 
 |  Parameter | Type | Description |
 |  --- | --- | --- |
 |  request | [CountTokensRequest](./vertexai.counttokensrequest.md#counttokensrequest_interface) \| string \| Array&lt;string \| [Part](./vertexai.md#part)<!-- -->&gt; |  |
+|  singleRequestOptions | [SingleRequestOptions](./vertexai.singlerequestoptions.md#singlerequestoptions_interface) |  |
 
 <b>Returns:</b>
 
@@ -138,14 +139,15 @@ Makes a single non-streaming call to the model and returns an object containing
 <b>Signature:</b>
 
 ```typescript
-generateContent(request: GenerateContentRequest | string | Array<string | Part>): Promise<GenerateContentResult>;
+generateContent(request: GenerateContentRequest | string | Array<string | Part>, singleRequestOptions?: SingleRequestOptions): Promise<GenerateContentResult>;
 ```
 
 #### Parameters
 
 |  Parameter | Type | Description |
 |  --- | --- | --- |
 |  request | [GenerateContentRequest](./vertexai.generatecontentrequest.md#generatecontentrequest_interface) \| string \| Array&lt;string \| [Part](./vertexai.md#part)<!-- -->&gt; |  |
+|  singleRequestOptions | [SingleRequestOptions](./vertexai.singlerequestoptions.md#singlerequestoptions_interface) |  |
 
 <b>Returns:</b>
 
@@ -158,14 +160,15 @@ Makes a single streaming call to the model and returns an object containing an i
 <b>Signature:</b>
 
 ```typescript
-generateContentStream(request: GenerateContentRequest | string | Array<string | Part>): Promise<GenerateContentStreamResult>;
+generateContentStream(request: GenerateContentRequest | string | Array<string | Part>, singleRequestOptions?: SingleRequestOptions): Promise<GenerateContentStreamResult>;
 ```
 
 #### Parameters
 
 |  Parameter | Type | Description |
 |  --- | --- | --- |
 |  request | [GenerateContentRequest](./vertexai.generatecontentrequest.md#generatecontentrequest_interface) \| string \| Array&lt;string \| [Part](./vertexai.md#part)<!-- -->&gt; |  |
+|  singleRequestOptions | [SingleRequestOptions](./vertexai.singlerequestoptions.md#singlerequestoptions_interface) |  |
 
 <b>Returns:</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) |  | <b><i>(Public Preview)</i></b> Generates images using the Imagen model and returns them as base64-encoded strings. |
+|  [generateImages(prompt, singleRequestOptions)](./vertexai.imagenmodel.md#imagenmodelgenerateimages) |  | <b><i>(Public Preview)</i></b> Generates images using the Imagen model and returns them as base64-encoded strings. |
 
 ## ImagenModel.(constructor)
 
@@ -118,14 +118,15 @@ If the prompt was not blocked, but one or more of the generated images were filt
 <b>Signature:</b>
 
 ```typescript
-generateImages(prompt: string): Promise<ImagenGenerationResponse<ImagenInlineImage>>;
+generateImages(prompt: string, singleRequestOptions?: SingleRequestOptions): Promise<ImagenGenerationResponse<ImagenInlineImage>>;
 ```
 
 #### Parameters
 
 |  Parameter | Type | Description |
 |  --- | --- | --- |
 |  prompt | string | A text prompt describing the image(s) to generate. |
+|  singleRequestOptions | [SingleRequestOptions](./vertexai.singlerequestoptions.md#singlerequestoptions_interface) |  |
 
 <b>Returns:</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 <code>timeout</code> and <code>baseUrl</code>) with request-specific controls like cancellation via <code>AbortSignal</code>.<!-- -->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. |

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)<!-- -->).
+
+<b>Signature:</b>
+
+```typescript
+export interface SingleRequestOptions extends RequestOptions 
+```
+<b>Extends:</b> [RequestOptions](./vertexai.requestoptions.md#requestoptions_interface)
+
+## Properties
+
+|  Property | Type | Description |
+|  --- | --- | --- |
+|  [signal](./vertexai.singlerequestoptions.md#singlerequestoptionssignal) | AbortSignal | An <code>AbortSignal</code> instance that allows cancelling ongoing requests (like <code>generateContent</code> or <code>generateImages</code>).<!-- -->If provided, calling <code>abort()</code> on the corresponding <code>AbortController</code> will attempt to cancel the underlying HTTP request. An <code>AbortError</code> 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.
+
+<b>Signature:</b>
+
+```typescript
+signal?: AbortSignal;
+```
+
+### Example
+
+
+```javascript
+const controller = new AbortController();
+const model = getGenerativeModel({
+  // ...
+});
+
+// To cancel request:
+controller.abort();
+
+```
+

packages/vertexai/src/methods/chat-session.test.ts

@@ -52,6 +52,64 @@ describe('ChatSession', () => {
         match.any
       );
     });
+    it('singleRequestOptions overrides requestOptions', async () => {
+      const generateContentStub = stub(
+        generateContentMethods,
+        'generateContent'
+      ).rejects('generateContent failed'); // not important
+      const requestOptions = {
+        timeout: 1000
+      };
+      const singleRequestOptions = {
+        timeout: 2000
+      };
+      const chatSession = new ChatSession(
+        fakeApiSettings,
+        'a-model',
+        undefined,
+        requestOptions
+      );
+      await expect(chatSession.sendMessage('hello', singleRequestOptions)).to.be
+        .rejected;
+      expect(generateContentStub).to.be.calledWith(
+        fakeApiSettings,
+        'a-model',
+        match.any,
+        match({
+          timeout: singleRequestOptions.timeout
+        })
+      );
+    });
+    it('singleRequestOptions is merged with requestOptions', async () => {
+      const generateContentStub = stub(
+        generateContentMethods,
+        'generateContent'
+      ).rejects('generateContent failed'); // not important
+      const abortController = new AbortController();
+      const requestOptions = {
+        timeout: 1000
+      };
+      const singleRequestOptions = {
+        signal: abortController.signal
+      };
+      const chatSession = new ChatSession(
+        fakeApiSettings,
+        'a-model',
+        undefined,
+        requestOptions
+      );
+      await expect(chatSession.sendMessage('hello', singleRequestOptions)).to.be
+        .rejected;
+      expect(generateContentStub).to.be.calledWith(
+        fakeApiSettings,
+        'a-model',
+        match.any,
+        match({
+          timeout: requestOptions.timeout,
+          signal: singleRequestOptions.signal
+        })
+      );
+    });
   });
   describe('sendMessageStream()', () => {
     it('generateContentStream errors should be catchable', async () => {
@@ -94,5 +152,63 @@ describe('ChatSession', () => {
       );
       clock.restore();
     });
+    it('singleRequestOptions overrides requestOptions', async () => {
+      const generateContentStub = stub(
+        generateContentMethods,
+        'generateContent'
+      ).rejects('generateContent failed'); // not important
+      const requestOptions = {
+        timeout: 1000
+      };
+      const singleRequestOptions = {
+        timeout: 2000
+      };
+      const chatSession = new ChatSession(
+        fakeApiSettings,
+        'a-model',
+        undefined,
+        requestOptions
+      );
+      await expect(chatSession.sendMessage('hello', singleRequestOptions)).to.be
+        .rejected;
+      expect(generateContentStub).to.be.calledWith(
+        fakeApiSettings,
+        'a-model',
+        match.any,
+        match({
+          timeout: singleRequestOptions.timeout
+        })
+      );
+    });
+    it('singleRequestOptions is merged with requestOptions', async () => {
+      const generateContentStub = stub(
+        generateContentMethods,
+        'generateContent'
+      ).rejects('generateContent failed'); // not important
+      const abortController = new AbortController();
+      const requestOptions = {
+        timeout: 1000
+      };
+      const singleRequestOptions = {
+        signal: abortController.signal
+      };
+      const chatSession = new ChatSession(
+        fakeApiSettings,
+        'a-model',
+        undefined,
+        requestOptions
+      );
+      await expect(chatSession.sendMessage('hello', singleRequestOptions)).to.be
+        .rejected;
+      expect(generateContentStub).to.be.calledWith(
+        fakeApiSettings,
+        'a-model',
+        match.any,
+        match({
+          timeout: requestOptions.timeout,
+          signal: singleRequestOptions.signal
+        })
+      );
+    });
   });
 });